Removed all.
diff --git a/.gitattributes b/.gitattributes
deleted file mode 100644
index fac7bf8..0000000
--- a/.gitattributes
+++ /dev/null
@@ -1,6 +0,0 @@
-r/R/RcppExports.R linguist-generated=true
-r/R/arrowExports.R linguist-generated=true
-r/src/RcppExports.cpp linguist-generated=true
-r/src/arrowExports.cpp linguist-generated=true
-r/man/*.Rd linguist-generated=true
-
diff --git a/.github/ISSUE_TEMPLATE/bug_report.md b/.github/ISSUE_TEMPLATE/bug_report.md
deleted file mode 100644
index 5600dab..0000000
--- a/.github/ISSUE_TEMPLATE/bug_report.md
+++ /dev/null
@@ -1,20 +0,0 @@
----
-name: Bug report
-about: Create a report to help us improve
-title: ''
-labels: bug
-assignees: ''
-
----
-
-**Describe the bug**
-A clear and concise description of what the bug is.
-
-**To Reproduce**
-Steps to reproduce the behavior:
-
-**Expected behavior**
-A clear and concise description of what you expected to happen.
-
-**Additional context**
-Add any other context about the problem here.
diff --git a/.github/ISSUE_TEMPLATE/feature_request.md b/.github/ISSUE_TEMPLATE/feature_request.md
deleted file mode 100644
index d9883dd..0000000
--- a/.github/ISSUE_TEMPLATE/feature_request.md
+++ /dev/null
@@ -1,21 +0,0 @@
----
-name: Feature request
-about: Suggest an idea for this project
-title: ''
-labels: enhancement
-assignees: ''
-
----
-
-**Is your feature request related to a problem or challenge? Please describe what you are trying to do.**
-A clear and concise description of what the problem is. Ex. I'm always frustrated when [...]
-(This section helps Arrow developers understand the context and *why* for this feature, in addition to the *what*)
-
-**Describe the solution you'd like**
-A clear and concise description of what you want to happen.
-
-**Describe alternatives you've considered**
-A clear and concise description of any alternative solutions or features you've considered.
-
-**Additional context**
-Add any other context or screenshots about the feature request here.
diff --git a/.github/pull_request_template.md b/.github/pull_request_template.md
deleted file mode 100644
index db170e3..0000000
--- a/.github/pull_request_template.md
+++ /dev/null
@@ -1,31 +0,0 @@
-# Which issue does this PR close?
-
-<!---
-We generally require a GitHub issue to be filed for all bug fixes and enhancements and this helps us generate change logs for our releases. You can link an issue to this PR using the GitHub syntax. For example `Closes #123` indicates that this PR will close issue #123.
--->
-
-Closes #.
-
-# Rationale for this change
-
- <!---
- Why are you proposing this change? If this is already explained clearly in the issue then this section is not needed.
- Explaining clearly why changes are proposed helps reviewers understand your changes and offer better suggestions for fixes.
--->
-
-# What changes are included in this PR?
-
-<!---
-There is no need to duplicate the description in the issue here but it is sometimes worth providing a summary of the individual changes in this PR.
--->
-
-# Are there any user-facing changes?
-
-
-<!---
-If there are user-facing changes then we may require documentation to be updated before approving the PR.
--->
-
-<!---
-If there are any breaking changes to public APIs, please add the `breaking change` label.
--->
diff --git a/.github/workflows/cancel.yml b/.github/workflows/cancel.yml
deleted file mode 100644
index b4fb904..0000000
--- a/.github/workflows/cancel.yml
+++ /dev/null
@@ -1,53 +0,0 @@
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-
-name: Cancel stale runs
-
-on:
- workflow_run:
- # The name of another workflow (whichever one) that always runs on PRs
- workflows: ['Dev']
- types: ['requested']
-
-jobs:
- cancel-stale-workflow-runs:
- name: "Cancel stale workflow runs"
- runs-on: ubuntu-latest
- steps:
- # Unfortunately, we need to define a separate cancellation step for
- # each workflow where we want to cancel stale runs.
- - uses: potiuk/cancel-workflow-runs@master
- name: "Cancel stale Dev runs"
- with:
- cancelMode: allDuplicates
- token: ${{ secrets.GITHUB_TOKEN }}
- workflowFileName: dev.yml
- skipEventTypes: '["push", "schedule"]'
- - uses: potiuk/cancel-workflow-runs@master
- name: "Cancel stale Integration runs"
- with:
- cancelMode: allDuplicates
- token: ${{ secrets.GITHUB_TOKEN }}
- workflowFileName: integration.yml
- skipEventTypes: '["push", "schedule"]'
- - uses: potiuk/cancel-workflow-runs@master
- name: "Cancel stale Rust runs"
- with:
- cancelMode: allDuplicates
- token: ${{ secrets.GITHUB_TOKEN }}
- workflowFileName: rust.yml
- skipEventTypes: '["push", "schedule"]'
diff --git a/.github/workflows/comment_bot.yml b/.github/workflows/comment_bot.yml
deleted file mode 100644
index 6ca0953..0000000
--- a/.github/workflows/comment_bot.yml
+++ /dev/null
@@ -1,72 +0,0 @@
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-
-name: Comment Bot
-
-on:
- # TODO(kszucs): support pull_request_review_comment
- issue_comment:
- types:
- - created
- - edited
-
-jobs:
- crossbow:
- name: Listen!
- if: startsWith(github.event.comment.body, '@github-actions crossbow')
- runs-on: ubuntu-latest
- steps:
- - name: Checkout Arrow
- uses: actions/checkout@v2
- with:
- repository: apache/arrow
- - name: Set up Python
- uses: actions/setup-python@v2
- with:
- python-version: 3.8
- - name: Install Archery and Crossbow dependencies
- run: pip install -e dev/archery[bot]
- - name: Handle Github comment event
- env:
- ARROW_GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
- CROSSBOW_GITHUB_TOKEN: ${{ secrets.CROSSBOW_GITHUB_TOKEN }}
- run: |
- archery trigger-bot \
- --event-name ${{ github.event_name }} \
- --event-payload ${{ github.event_path }}
-
- rebase:
- name: "Rebase"
- if: startsWith(github.event.comment.body, '@github-actions rebase')
- runs-on: ubuntu-latest
- steps:
- - uses: actions/checkout@v2
- - uses: r-lib/actions/pr-fetch@master
- with:
- repo-token: ${{ secrets.GITHUB_TOKEN }}
- - name: Rebase on ${{ github.repository }} master
- run: |
- set -ex
- git config user.name "$(git log -1 --pretty=format:%an)"
- git config user.email "$(git log -1 --pretty=format:%ae)"
- git remote add upstream https://github.com/${{ github.repository }}
- git fetch --unshallow upstream master
- git rebase upstream/master
- - uses: r-lib/actions/pr-push@master
- with:
- repo-token: ${{ secrets.GITHUB_TOKEN }}
- args: "--force"
diff --git a/.github/workflows/dev.yml b/.github/workflows/dev.yml
deleted file mode 100644
index 548f0dd..0000000
--- a/.github/workflows/dev.yml
+++ /dev/null
@@ -1,67 +0,0 @@
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-
-name: Dev
-
-on:
- # always trigger
- push:
- pull_request:
-
-env:
- ARCHERY_DOCKER_USER: ${{ secrets.DOCKERHUB_USER }}
- ARCHERY_DOCKER_PASSWORD: ${{ secrets.DOCKERHUB_TOKEN }}
-
-jobs:
-
- rat:
- name: Release Audit Tool (RAT)
- runs-on: ubuntu-latest
- steps:
- - name: Checkout Arrow
- uses: actions/checkout@v2
- with:
- repository: apache/arrow
- submodules: true
- fetch-depth: 0
- - name: Checkout Arrow Rust
- uses: actions/checkout@v2
- with:
- path: rust
- fetch-depth: 0
- - name: Setup Python
- uses: actions/setup-python@v1
- with:
- python-version: 3.8
- - name: Setup Archery
- run: pip install -e dev/archery[lint]
- - name: Lint
- run: archery lint --rat
-
- prettier:
- name: Use prettier to check formatting of documents
- runs-on: ubuntu-latest
- steps:
- - uses: actions/checkout@v2
- - uses: actions/setup-node@v2
- with:
- node-version: "14"
- - name: Prettier check
- run: |
- # if you encounter error, try rerun the command below with --write instead of --check
- # and commit the changes
- npx prettier@2.3.0 --check {arrow,arrow-flight,dev,integration-testing,parquet}/**/*.md README.md CODE_OF_CONDUCT.md CONTRIBUTING.md
diff --git a/.github/workflows/dev_pr.yml b/.github/workflows/dev_pr.yml
deleted file mode 100644
index 8202853..0000000
--- a/.github/workflows/dev_pr.yml
+++ /dev/null
@@ -1,45 +0,0 @@
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-
-name: Dev PR
-
-on:
- pull_request:
-
-jobs:
- process:
- name: Process
- runs-on: ubuntu-latest
- steps:
- - uses: actions/checkout@v2
-
- - name: Assign GitHub labels
- if: |
- github.event_name == 'pull_request_target' &&
- (github.event.action == 'opened' ||
- github.event.action == 'synchronize')
- uses: actions/labeler@2.2.0
- with:
- repo-token: ${{ secrets.GITHUB_TOKEN }}
- configuration-path: .github/workflows/dev_pr/labeler.yml
- sync-labels: true
-
- #- name: Checks if PR needs rebase
- # uses: eps1lon/actions-label-merge-conflict@releases/2.x
- # with:
- # dirtyLabel: "needs-rebase"
- # repoToken: "${{ secrets.GITHUB_TOKEN }}"
diff --git a/.github/workflows/dev_pr/labeler.yml b/.github/workflows/dev_pr/labeler.yml
deleted file mode 100644
index 63b4154..0000000
--- a/.github/workflows/dev_pr/labeler.yml
+++ /dev/null
@@ -1,25 +0,0 @@
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-
-arrow:
- - arrow/**/*
-
-arrow-flight:
- - arrow-flight/**/*
-
-parquet:
- - parquet/**/*
diff --git a/.github/workflows/integration.yml b/.github/workflows/integration.yml
deleted file mode 100644
index cab6dd3..0000000
--- a/.github/workflows/integration.yml
+++ /dev/null
@@ -1,48 +0,0 @@
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-
-name: Integration
-
-on:
- push:
- pull_request:
-
-jobs:
-
- docker:
- name: Integration Test
- runs-on: ubuntu-latest
- steps:
- - name: Checkout Arrow
- uses: actions/checkout@v2
- with:
- repository: apache/arrow
- submodules: true
- fetch-depth: 0
- - name: Checkout Arrow Rust
- uses: actions/checkout@v2
- with:
- path: rust
- fetch-depth: 0
- - name: Setup Python
- uses: actions/setup-python@v2
- with:
- python-version: 3.8
- - name: Setup Archery
- run: pip install -e dev/archery[docker]
- - name: Execute Docker Build
- run: archery docker run -e ARCHERY_INTEGRATION_WITH_RUST=1 conda-integration
diff --git a/.github/workflows/miri.yaml b/.github/workflows/miri.yaml
deleted file mode 100644
index b74170f..0000000
--- a/.github/workflows/miri.yaml
+++ /dev/null
@@ -1,60 +0,0 @@
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-
-name: Rust
-
-on:
- # always trigger
- push:
- pull_request:
-
-jobs:
-
- miri-checks:
- name: MIRI
- runs-on: ubuntu-latest
- strategy:
- matrix:
- arch: [amd64]
- rust: [nightly-2021-03-24]
- steps:
- - uses: actions/checkout@v2
- with:
- submodules: true
- - uses: actions/cache@v2
- with:
- path: |
- ~/.cargo/registry
- ~/.cargo/git
- target
- key: ${{ runner.os }}-cargo-miri-${{ hashFiles('**/Cargo.lock') }}
- - name: Setup Rust toolchain
- run: |
- rustup toolchain install ${{ matrix.rust }}
- rustup default ${{ matrix.rust }}
- rustup component add rustfmt clippy miri
- - name: Run Miri Checks
- env:
- RUST_BACKTRACE: full
- RUST_LOG: 'trace'
- run: |
- export MIRIFLAGS="-Zmiri-disable-isolation"
- cargo miri setup
- cargo clean
- # Currently only the arrow crate is tested with miri
- # IO related tests and some unsupported tests are skipped
- cargo miri test -p arrow -- --skip csv --skip ipc --skip json || true
diff --git a/.github/workflows/rust.yml b/.github/workflows/rust.yml
deleted file mode 100644
index 34e0449..0000000
--- a/.github/workflows/rust.yml
+++ /dev/null
@@ -1,417 +0,0 @@
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-
-name: Rust
-
-on:
- # always trigger
- push:
- pull_request:
-
-jobs:
-
- # build the library, a compilation step used by multiple steps below
- linux-build-lib:
- name: Build Libraries on AMD64 Rust ${{ matrix.rust }}
- runs-on: ubuntu-latest
- strategy:
- matrix:
- arch: [amd64]
- rust: [stable]
- container:
- image: ${{ matrix.arch }}/rust
- env:
- # Disable full debug symbol generation to speed up CI build and keep memory down
- # "1" means line tables only, which is useful for panic tracebacks.
- RUSTFLAGS: "-C debuginfo=1"
- steps:
- - uses: actions/checkout@v2
- - name: Cache Cargo
- uses: actions/cache@v2
- with:
- # these represent dependencies downloaded by cargo
- # and thus do not depend on the OS, arch nor rust version.
- path: /github/home/.cargo
- key: cargo-cache-
- - name: Cache Rust dependencies
- uses: actions/cache@v2
- with:
- # these represent compiled steps of both dependencies and arrow
- # and thus are specific for a particular OS, arch and rust version.
- path: /github/home/target
- key: ${{ runner.os }}-${{ matrix.arch }}-target-cache-${{ matrix.rust }}-
- - name: Setup Rust toolchain
- run: |
- rustup toolchain install ${{ matrix.rust }}
- rustup default ${{ matrix.rust }}
- rustup component add rustfmt
- - name: Build Workspace
- run: |
- export CARGO_HOME="/github/home/.cargo"
- export CARGO_TARGET_DIR="/github/home/target"
- cargo build
-
- # test the crate
- linux-test:
- name: Test Workspace on AMD64 Rust ${{ matrix.rust }}
- needs: [linux-build-lib]
- runs-on: ubuntu-latest
- strategy:
- matrix:
- arch: [amd64]
- rust: [stable]
- container:
- image: ${{ matrix.arch }}/rust
- env:
- # Disable full debug symbol generation to speed up CI build and keep memory down
- # "1" means line tables only, which is useful for panic tracebacks.
- RUSTFLAGS: "-C debuginfo=1"
- ARROW_TEST_DATA: /__w/arrow-rs/arrow-rs/testing/data
- PARQUET_TEST_DATA: /__w/arrow-rs/arrow-rs/parquet-testing/data
- steps:
- - uses: actions/checkout@v2
- with:
- submodules: true
- - name: Cache Cargo
- uses: actions/cache@v2
- with:
- path: /github/home/.cargo
- # this key equals the ones on `linux-build-lib` for re-use
- key: cargo-cache-
- - name: Cache Rust dependencies
- uses: actions/cache@v2
- with:
- path: /github/home/target
- # this key equals the ones on `linux-build-lib` for re-use
- key: ${{ runner.os }}-${{ matrix.arch }}-target-cache-${{ matrix.rust }}
- - name: Setup Rust toolchain
- run: |
- rustup toolchain install ${{ matrix.rust }}
- rustup default ${{ matrix.rust }}
- rustup component add rustfmt
- - name: Run tests
- run: |
- export CARGO_HOME="/github/home/.cargo"
- export CARGO_TARGET_DIR="/github/home/target"
- # run tests on all workspace members with default feature list
- cargo test
- cd arrow
- # re-run tests on arrow workspace with additional features
- cargo test --features=prettyprint
- # run test on arrow with minimal set of features
- cargo test --no-default-features
- cargo run --example builders
- cargo run --example dynamic_types
- cargo run --example read_csv
- cargo run --example read_csv_infer_schema
-
- # test the --features "simd" of the arrow crate. This requires nightly.
- linux-test-simd:
- name: Test SIMD on AMD64 Rust ${{ matrix.rust }}
- runs-on: ubuntu-latest
- strategy:
- matrix:
- arch: [amd64]
- rust: [nightly-2021-03-24]
- container:
- image: ${{ matrix.arch }}/rust
- env:
- # Disable full debug symbol generation to speed up CI build and keep memory down
- # "1" means line tables only, which is useful for panic tracebacks.
- RUSTFLAGS: "-C debuginfo=1"
- ARROW_TEST_DATA: /__w/arrow-rs/arrow-rs/testing/data
- steps:
- - uses: actions/checkout@v2
- with:
- submodules: true
- - name: Cache Cargo
- uses: actions/cache@v2
- with:
- path: /github/home/.cargo
- # this key equals the ones on `linux-build-lib` for re-use
- key: cargo-cache-
- - name: Cache Rust dependencies
- uses: actions/cache@v2
- with:
- path: /github/home/target
- # this key equals the ones on `linux-build-lib` for re-use
- key: ${{ runner.os }}-${{ matrix.arch }}-target-cache-${{ matrix.rust }}
- - name: Setup Rust toolchain
- run: |
- rustup toolchain install ${{ matrix.rust }}
- rustup default ${{ matrix.rust }}
- rustup component add rustfmt
- - name: Run tests
- run: |
- export CARGO_HOME="/github/home/.cargo"
- export CARGO_TARGET_DIR="/github/home/target"
- cd arrow
- cargo test --features "simd"
-
- windows-and-macos:
- name: Test on ${{ matrix.os }} Rust ${{ matrix.rust }}
- runs-on: ${{ matrix.os }}
- strategy:
- matrix:
- os: [windows-latest, macos-latest]
- rust: [stable]
- steps:
- - uses: actions/checkout@v2
- with:
- submodules: true
- # TODO: this won't cache anything, which is expensive. Setup this action
- # with a OS-dependent path.
- - name: Setup Rust toolchain
- run: |
- rustup toolchain install ${{ matrix.rust }}
- rustup default ${{ matrix.rust }}
- rustup component add rustfmt
- - name: Run tests
- shell: bash
- run: |
- export ARROW_TEST_DATA=$(pwd)/testing/data
- export PARQUET_TEST_DATA=$(pwd)/parquet-testing/data
- # do not produce debug symbols to keep memory usage down
- export RUSTFLAGS="-C debuginfo=0"
- cargo test
-
- clippy:
- name: Clippy
- needs: [linux-build-lib]
- runs-on: ubuntu-latest
- strategy:
- matrix:
- arch: [amd64]
- rust: [stable]
- container:
- image: ${{ matrix.arch }}/rust
- env:
- # Disable full debug symbol generation to speed up CI build and keep memory down
- # "1" means line tables only, which is useful for panic tracebacks.
- RUSTFLAGS: "-C debuginfo=1"
- steps:
- - uses: actions/checkout@v2
- with:
- submodules: true
- - name: Cache Cargo
- uses: actions/cache@v2
- with:
- path: /github/home/.cargo
- # this key equals the ones on `linux-build-lib` for re-use
- key: cargo-cache-
- - name: Cache Rust dependencies
- uses: actions/cache@v2
- with:
- path: /github/home/target
- # this key equals the ones on `linux-build-lib` for re-use
- key: ${{ runner.os }}-${{ matrix.arch }}-target-cache-${{ matrix.rust }}
- - name: Setup Rust toolchain
- run: |
- rustup toolchain install ${{ matrix.rust }}
- rustup default ${{ matrix.rust }}
- rustup component add rustfmt clippy
- - name: Run clippy
- run: |
- export CARGO_HOME="/github/home/.cargo"
- export CARGO_TARGET_DIR="/github/home/target"
- cargo clippy --all-targets --workspace -- -D warnings -A clippy::redundant_field_names
-
- lint:
- name: Lint
- runs-on: ubuntu-latest
- container:
- image: amd64/rust
- steps:
- - uses: actions/checkout@v2
- - name: Setup toolchain
- run: |
- rustup toolchain install stable
- rustup default stable
- rustup component add rustfmt
- - name: Run
- run: cargo fmt --all -- --check
- coverage:
- name: Coverage
- runs-on: ubuntu-latest
- strategy:
- matrix:
- arch: [amd64]
- rust: [stable]
- steps:
- - uses: actions/checkout@v2
- with:
- submodules: true
- - name: Cache Cargo
- uses: actions/cache@v2
- with:
- path: /home/runner/.cargo
- # this key is not equal because the user is different than on a container (runner vs github)
- key: cargo-coverage-cache-
- - name: Cache Rust dependencies
- uses: actions/cache@v2
- with:
- path: /home/runner/target
- # this key is not equal because coverage uses different compilation flags.
- key: ${{ runner.os }}-${{ matrix.arch }}-target-coverage-cache-${{ matrix.rust }}-
- - name: Run coverage
- run: |
- export CARGO_HOME="/home/runner/.cargo"
- export CARGO_TARGET_DIR="/home/runner/target"
-
- export ARROW_TEST_DATA=$(pwd)/testing/data
- export PARQUET_TEST_DATA=$(pwd)/parquet-testing/data
-
- # 2020-11-15: There is a cargo-tarpaulin regression in 0.17.0
- # see https://github.com/xd009642/tarpaulin/issues/618
- cargo install --version 0.16.0 cargo-tarpaulin
- cargo tarpaulin --out Xml
- - name: Report coverage
- continue-on-error: true
- run: bash <(curl -s https://codecov.io/bash)
-
- # test FFI against the C-Data interface exposed by pyarrow
- pyarrow-integration-test:
- name: Test Pyarrow C Data Interface
- runs-on: ubuntu-latest
- strategy:
- matrix:
- rust: [stable]
- steps:
- - uses: actions/checkout@v2
- with:
- submodules: true
- - name: Setup Rust toolchain
- run: |
- rustup toolchain install ${{ matrix.rust }}
- rustup default ${{ matrix.rust }}
- rustup component add rustfmt clippy
- - name: Cache Cargo
- uses: actions/cache@v2
- with:
- path: /home/runner/.cargo
- key: cargo-maturin-cache-
- - name: Cache Rust dependencies
- uses: actions/cache@v2
- with:
- path: /home/runner/target
- # this key is not equal because maturin uses different compilation flags.
- key: ${{ runner.os }}-${{ matrix.arch }}-target-maturin-cache-${{ matrix.rust }}-
- - uses: actions/setup-python@v2
- with:
- python-version: '3.7'
- - name: Install Python dependencies
- run: python -m pip install --upgrade pip setuptools wheel
- - name: Run tests
- run: |
- export CARGO_HOME="/home/runner/.cargo"
- export CARGO_TARGET_DIR="/home/runner/target"
-
- cd arrow-pyarrow-integration-testing
-
- python -m venv venv
- source venv/bin/activate
-
- pip install maturin==0.8.2 toml==0.10.1 pyarrow==1.0.0
- maturin develop
- python -m unittest discover tests
-
- # test the arrow crate builds against wasm32 in stable rust
- wasm32-build:
- name: Build wasm32 on AMD64 Rust ${{ matrix.rust }}
- runs-on: ubuntu-latest
- strategy:
- matrix:
- arch: [amd64]
- rust: [nightly-2021-03-24]
- container:
- image: ${{ matrix.arch }}/rust
- env:
- # Disable full debug symbol generation to speed up CI build and keep memory down
- # "1" means line tables only, which is useful for panic tracebacks.
- RUSTFLAGS: "-C debuginfo=1"
- ARROW_TEST_DATA: /__w/arrow-rs/arrow-rs/testing/data
- PARQUET_TEST_DATA: /__w/arrow/arrow/parquet-testing/data
- steps:
- - uses: actions/checkout@v2
- with:
- submodules: true
- - name: Cache Cargo
- uses: actions/cache@v2
- with:
- path: /github/home/.cargo
- # this key equals the ones on `linux-build-lib` for re-use
- key: cargo-cache-
- - name: Cache Rust dependencies
- uses: actions/cache@v2
- with:
- path: /github/home/target
- key: ${{ runner.os }}-${{ matrix.arch }}-target-wasm32-cache-${{ matrix.rust }}
- - name: Setup Rust toolchain
- run: |
- rustup toolchain install ${{ matrix.rust }}
- rustup override set ${{ matrix.rust }}
- rustup component add rustfmt
- rustup target add wasm32-unknown-unknown
- - name: Build arrow crate
- run: |
- export CARGO_HOME="/github/home/.cargo"
- export CARGO_TARGET_DIR="/github/home/target"
- cd arrow
- cargo build --target wasm32-unknown-unknown
-
- # test the projects can build without default features
- default-build:
- name: Check No Defaults on AMD64 Rust ${{ matrix.rust }}
- runs-on: ubuntu-latest
- strategy:
- matrix:
- arch: [amd64]
- rust: [stable]
- container:
- image: ${{ matrix.arch }}/rust
- env:
- # Disable full debug symbol generation to speed up CI build and keep memory down
- # "1" means line tables only, which is useful for panic tracebacks.
- RUSTFLAGS: "-C debuginfo=1"
- ARROW_TEST_DATA: /__w/arrow-rs/arrow-rs/testing/data
- PARQUET_TEST_DATA: /__w/arrow/arrow/parquet-testing/data
- steps:
- - uses: actions/checkout@v2
- with:
- submodules: true
- - name: Cache Cargo
- uses: actions/cache@v2
- with:
- path: /github/home/.cargo
- # this key equals the ones on `linux-build-lib` for re-use
- key: cargo-cache-
- - name: Cache Rust dependencies
- uses: actions/cache@v2
- with:
- path: /github/home/target
- key: ${{ runner.os }}-${{ matrix.arch }}-target-wasm32-cache-${{ matrix.rust }}
- - name: Setup Rust toolchain
- run: |
- rustup toolchain install ${{ matrix.rust }}
- rustup override set ${{ matrix.rust }}
- rustup component add rustfmt
- - name: Build arrow crate
- run: |
- export CARGO_HOME="/github/home/.cargo"
- export CARGO_TARGET_DIR="/github/home/target"
- cd arrow
- cargo check --all-targets --no-default-features
diff --git a/.github_changelog_generator b/.github_changelog_generator
deleted file mode 100644
index 52452ed..0000000
--- a/.github_changelog_generator
+++ /dev/null
@@ -1,29 +0,0 @@
-#!/bin/bash
-#
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-#
-
-# point to the old changelog in apache/arrow
-front-matter=For older versions, see [apache/arrow/CHANGELOG.md](https://github.com/apache/arrow/blob/master/CHANGELOG.md)\n
-# some issues are just documentation
-add-sections={"documentation":{"prefix":"**Documentation updates:**","labels":["documentation"]}}
-# uncomment to not show PRs. TBD if we shown them or not.
-#pull-requests=false
-# so that the component is shown associated with the issue
-issue-line-labels=arrow,parquet,arrow-flight
-exclude-labels=development-process,invalid
diff --git a/.gitignore b/.gitignore
deleted file mode 100644
index e8d9955..0000000
--- a/.gitignore
+++ /dev/null
@@ -1,8 +0,0 @@
-Cargo.lock
-target
-rusty-tags.vi
-.history
-.flatbuffers/
-.idea/
-.vscode
-venv/*
diff --git a/.gitmodules b/.gitmodules
deleted file mode 100644
index 97da927..0000000
--- a/.gitmodules
+++ /dev/null
@@ -1,6 +0,0 @@
-[submodule "testing"]
- path = testing
- url = https://github.com/apache/arrow-testing
-[submodule "parquet-testing"]
- path = parquet-testing
- url = https://github.com/apache/parquet-testing.git
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
deleted file mode 100644
index 5331a53..0000000
--- a/.pre-commit-config.yaml
+++ /dev/null
@@ -1,69 +0,0 @@
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-
-# To use this, install the python package `pre-commit` and
-# run once `pre-commit install`. This will setup a git pre-commit-hook
-# that is executed on each commit and will report the linting problems.
-# To run all hooks on all files use `pre-commit run -a`
-
-repos:
- - repo: local
- hooks:
- - id: rat
- name: Release Audit Tool
- language: system
- entry: bash -c "git archive HEAD --prefix=apache-arrow/ --output=arrow-src.tar && ./dev/release/run-rat.sh arrow-src.tar"
- always_run: true
- pass_filenames: false
- - id: rustfmt
- name: Rust Format
- language: system
- entry: bash -c "cargo +stable fmt --all -- --check"
- files: ^.*\.rs$
- types:
- - file
- - rust
- - id: cmake-format
- name: CMake Format
- language: python
- entry: python run-cmake-format.py
- types: [cmake]
- additional_dependencies:
- - cmake_format==0.5.2
- - id: hadolint
- name: Docker Format
- language: docker_image
- types:
- - dockerfile
- entry: --entrypoint /bin/hadolint hadolint/hadolint:latest -
- exclude: ^dev/.*$
- - repo: git://github.com/pre-commit/pre-commit-hooks
- sha: v1.2.3
- hooks:
- - id: flake8
- name: Python Format
- files: ^(python|dev|integration)/
- types:
- - file
- - python
- - id: flake8
- name: Cython Format
- files: ^python/
- types:
- - file
- - cython
- args: [--config=python/.flake8.cython]
diff --git a/CHANGELOG.md b/CHANGELOG.md
deleted file mode 100644
index c901a4e..0000000
--- a/CHANGELOG.md
+++ /dev/null
@@ -1,78 +0,0 @@
-For older versions, see [apache/arrow/CHANGELOG.md](https://github.com/apache/arrow/blob/master/CHANGELOG.md)
-
-# Changelog
-
-## [4.0.1](https://github.com/apache/arrow-rs/tree/4.0.1) (2021-05-16)
-
-[Full Changelog](https://github.com/apache/arrow-rs/compare/8707fd2b2d17b17bd3e79be0255a18ffaea6914a...4.0.1)
-
-**Implemented enhancements:**
-
-- Improve performance of bound checking option [\#280](https://github.com/apache/arrow-rs/issues/280) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)]
-- extend compute kernel arity to include nullary functions [\#276](https://github.com/apache/arrow-rs/issues/276)
-- Add support for pretty-printing Decimal numbers [\#230](https://github.com/apache/arrow-rs/issues/230) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)]
-- CSV Reader String Dictionary Support [\#228](https://github.com/apache/arrow-rs/issues/228) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)]
-- Add Builder interface for adding Arrays to record batches [\#210](https://github.com/apache/arrow-rs/issues/210) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)]
-- Support auto-vectorization for min/max [\#209](https://github.com/apache/arrow-rs/issues/209) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)]
-- Support LargeUtf8 in sort kernel [\#25](https://github.com/apache/arrow-rs/issues/25) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)]
-
-**Fixed bugs:**
-
-- Rust 1.52 Clippy error [\#266](https://github.com/apache/arrow-rs/issues/266)
-- NaNs can break parquet statistics [\#255](https://github.com/apache/arrow-rs/issues/255) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)]
-- u64::MAX does not roundtrip through parquet [\#254](https://github.com/apache/arrow-rs/issues/254) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)]
-- Integration tests failing to compile \(flatbuffer\) [\#249](https://github.com/apache/arrow-rs/issues/249) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)]
-- Fix compatibility quirks between arrow and parquet structs [\#245](https://github.com/apache/arrow-rs/issues/245) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)]
-- Unable to write non-null Arrow structs to Parquet [\#244](https://github.com/apache/arrow-rs/issues/244) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)]
-- Dev PR / Process \(pull\_request\) Failing on PRs [\#242](https://github.com/apache/arrow-rs/issues/242)
-- schema: missing field `metadata` when deserialize [\#241](https://github.com/apache/arrow-rs/issues/241) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)]
-- Arrow does not compile due to flatbuffers upgrade [\#238](https://github.com/apache/arrow-rs/issues/238) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)]
-- Sort with limit panics for the limit includes some but not all nulls, for large arrays [\#235](https://github.com/apache/arrow-rs/issues/235) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)]
-- Read list field correctly in \<struct\<list\>\> [\#167](https://github.com/apache/arrow-rs/issues/167) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)]
-
-**Documentation updates:**
-
-- Comment out the instructions in the PR template [\#277](https://github.com/apache/arrow-rs/issues/277)
-- Update links to datafusion and ballista in README.md [\#19](https://github.com/apache/arrow-rs/issues/19)
-- Update "repository" in Cargo.toml [\#12](https://github.com/apache/arrow-rs/issues/12)
-- Improve docs for NullArray, new\_null\_array and new\_empty\_array [\#240](https://github.com/apache/arrow-rs/pull/240) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([alamb](https://github.com/alamb))
-
-**Merged pull requests:**
-
-- Remove old release scripts [\#293](https://github.com/apache/arrow-rs/pull/293) ([alamb](https://github.com/alamb))
-- manually bump development version [\#288](https://github.com/apache/arrow-rs/pull/288) ([nevi-me](https://github.com/nevi-me))
-- Fix subtraction underflow when sorting string arrays with many nulls [\#285](https://github.com/apache/arrow-rs/pull/285) ([medwards](https://github.com/medwards))
-- Speed up bound checking in `take` [\#281](https://github.com/apache/arrow-rs/pull/281) ([Dandandan](https://github.com/Dandandan))
-- Update PR template by commenting out instructions [\#278](https://github.com/apache/arrow-rs/pull/278) ([nevi-me](https://github.com/nevi-me))
-- Added Decimal support to pretty-print display utility \(\#230\) [\#273](https://github.com/apache/arrow-rs/pull/273) ([mgill25](https://github.com/mgill25))
-- Fix null struct and list roundtrip [\#270](https://github.com/apache/arrow-rs/pull/270) ([nevi-me](https://github.com/nevi-me))
-- 1.52 clippy fixes [\#267](https://github.com/apache/arrow-rs/pull/267) ([nevi-me](https://github.com/nevi-me))
-- Fix typo in csv/reader.rs [\#265](https://github.com/apache/arrow-rs/pull/265) ([domoritz](https://github.com/domoritz))
-- Fix empty Schema::metadata deserialization error [\#260](https://github.com/apache/arrow-rs/pull/260) ([hulunbier](https://github.com/hulunbier))
-- update datafusion and ballista doc links [\#259](https://github.com/apache/arrow-rs/pull/259) ([Jimexist](https://github.com/Jimexist))
-- support full u32 and u64 roundtrip through parquet [\#258](https://github.com/apache/arrow-rs/pull/258) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] ([crepererum](https://github.com/crepererum))
-- fix NaN handling in parquet statistics [\#256](https://github.com/apache/arrow-rs/pull/256) ([crepererum](https://github.com/crepererum))
-- \[MINOR\] Added env to run rust in integration. [\#253](https://github.com/apache/arrow-rs/pull/253) ([jorgecarleitao](https://github.com/jorgecarleitao))
-- \[Minor\] Made integration tests always run. [\#248](https://github.com/apache/arrow-rs/pull/248) ([jorgecarleitao](https://github.com/jorgecarleitao))
-- fix parquet max\_definition for non-null structs [\#246](https://github.com/apache/arrow-rs/pull/246) ([nevi-me](https://github.com/nevi-me))
-- Disabled rebase needed until demonstrate working. [\#243](https://github.com/apache/arrow-rs/pull/243) ([jorgecarleitao](https://github.com/jorgecarleitao))
-- pin flatbuffers to 0.8.4 [\#239](https://github.com/apache/arrow-rs/pull/239) ([ritchie46](https://github.com/ritchie46))
-- sort\_primitive result is capped to the min of limit or values.len [\#236](https://github.com/apache/arrow-rs/pull/236) ([medwards](https://github.com/medwards))
-- Read list field correctly [\#234](https://github.com/apache/arrow-rs/pull/234) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] ([nevi-me](https://github.com/nevi-me))
-- Fix code examples for RecordBatch::try\_from\_iter [\#231](https://github.com/apache/arrow-rs/pull/231) ([alamb](https://github.com/alamb))
-- Support string dictionaries in csv reader \(\#228\) [\#229](https://github.com/apache/arrow-rs/pull/229) ([tustvold](https://github.com/tustvold))
-- support LargeUtf8 in sort kernel [\#26](https://github.com/apache/arrow-rs/pull/26) ([ritchie46](https://github.com/ritchie46))
-- Removed unused files [\#22](https://github.com/apache/arrow-rs/pull/22) ([jorgecarleitao](https://github.com/jorgecarleitao))
-- ARROW-12504: Buffer::from\_slice\_ref set correct capacity [\#18](https://github.com/apache/arrow-rs/pull/18) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([tustvold](https://github.com/tustvold))
-- Add GitHub templates [\#17](https://github.com/apache/arrow-rs/pull/17) ([andygrove](https://github.com/andygrove))
-- ARROW-12493: Add support for writing dictionary arrays to CSV and JSON [\#16](https://github.com/apache/arrow-rs/pull/16) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([tustvold](https://github.com/tustvold))
-- ARROW-12426: \[Rust\] Fix concatentation of arrow dictionaries [\#15](https://github.com/apache/arrow-rs/pull/15) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([tustvold](https://github.com/tustvold))
-- Update repository and homepage urls [\#14](https://github.com/apache/arrow-rs/pull/14) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] [[arrow-flight](https://github.com/apache/arrow-rs/labels/arrow-flight)] [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] ([Dandandan](https://github.com/Dandandan))
-- Added rebase-needed bot [\#13](https://github.com/apache/arrow-rs/pull/13) ([jorgecarleitao](https://github.com/jorgecarleitao))
-- Added Integration tests against arrow [\#10](https://github.com/apache/arrow-rs/pull/10) ([jorgecarleitao](https://github.com/jorgecarleitao))
-- ARROW-12343: \[Rust\] Support auto-vectorization for min/max [\#9](https://github.com/apache/arrow-rs/pull/9) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([Dandandan](https://github.com/Dandandan))
-- ARROW-12411: \[Rust\] Create RecordBatches from Iterators [\#7](https://github.com/apache/arrow-rs/pull/7) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([alamb](https://github.com/alamb))
-
-
-
-\* *This Changelog was automatically generated by [github_changelog_generator](https://github.com/github-changelog-generator/github-changelog-generator)*
diff --git a/CODE_OF_CONDUCT.md b/CODE_OF_CONDUCT.md
deleted file mode 100644
index 9a24b9b..0000000
--- a/CODE_OF_CONDUCT.md
+++ /dev/null
@@ -1,24 +0,0 @@
-<!---
- Licensed to the Apache Software Foundation (ASF) under one
- or more contributor license agreements. See the NOTICE file
- distributed with this work for additional information
- regarding copyright ownership. The ASF licenses this file
- to you under the Apache License, Version 2.0 (the
- "License"); you may not use this file except in compliance
- with the License. You may obtain a copy of the License at
-
- http://www.apache.org/licenses/LICENSE-2.0
-
- Unless required by applicable law or agreed to in writing,
- software distributed under the License is distributed on an
- "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
- KIND, either express or implied. See the License for the
- specific language governing permissions and limitations
- under the License.
--->
-
-# Code of Conduct
-
-- [Code of Conduct for The Apache Software Foundation][1]
-
-[1]: https://www.apache.org/foundation/policies/conduct.html
diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md
deleted file mode 100644
index 18d6a7b..0000000
--- a/CONTRIBUTING.md
+++ /dev/null
@@ -1,77 +0,0 @@
-<!---
- Licensed to the Apache Software Foundation (ASF) under one
- or more contributor license agreements. See the NOTICE file
- distributed with this work for additional information
- regarding copyright ownership. The ASF licenses this file
- to you under the Apache License, Version 2.0 (the
- "License"); you may not use this file except in compliance
- with the License. You may obtain a copy of the License at
-
- http://www.apache.org/licenses/LICENSE-2.0
-
- Unless required by applicable law or agreed to in writing,
- software distributed under the License is distributed on an
- "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
- KIND, either express or implied. See the License for the
- specific language governing permissions and limitations
- under the License.
--->
-
-# How to contribute to Apache Arrow
-
-## Did you find a bug?
-
-The Arrow project uses JIRA as a bug tracker. To report a bug, you'll have
-to first create an account on the
-[Apache Foundation JIRA](https://issues.apache.org/jira/). The JIRA server
-hosts bugs and issues for multiple Apache projects. The JIRA project name
-for Arrow is "ARROW".
-
-To be assigned to an issue, ask an Arrow JIRA admin to go to
-[Arrow Roles](https://issues.apache.org/jira/plugins/servlet/project-config/ARROW/roles),
-click "Add users to a role," and add you to the "Contributor" role. Most
-committers are authorized to do this; if you're a committer and aren't
-able to load that project admin page, have someone else add you to the
-necessary role.
-
-Before you create a new bug entry, we recommend you first
-[search](https://issues.apache.org/jira/projects/ARROW/issues/ARROW-5140?filter=allopenissues)
-among existing Arrow issues.
-
-When you create a new JIRA entry, please don't forget to fill the "Component"
-field. Arrow has many subcomponents and this helps triaging and filtering
-tremendously. Also, we conventionally prefix the issue title with the component
-name in brackets, such as "[C++] Crash in Array::Frobnicate()", so as to make
-lists more easy to navigate, and we'd be grateful if you did the same.
-
-## Did you write a patch that fixes a bug or brings an improvement?
-
-First create a JIRA entry as described above. Then, submit your changes
-as a GitHub Pull Request. We'll ask you to prefix the pull request title
-with the JIRA issue number and the component name in brackets.
-(for example: "ARROW-2345: [C++] Fix crash in Array::Frobnicate()").
-Respecting this convention makes it easier for us to process the backlog
-of submitted Pull Requests.
-
-### Minor Fixes
-
-Any functionality change should have a JIRA opened. For minor changes that
-affect documentation, you do not need to open up a JIRA. Instead you can
-prefix the title of your PR with "MINOR: " if meets the following guidelines:
-
-- 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/contributing.html).
diff --git a/Cargo.toml b/Cargo.toml
deleted file mode 100644
index de7d36f..0000000
--- a/Cargo.toml
+++ /dev/null
@@ -1,31 +0,0 @@
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-
-[workspace]
-members = [
- "arrow",
- "parquet",
- "parquet_derive",
- "parquet_derive_test",
- "arrow-flight",
- "integration-testing",
-]
-
-# this package is excluded because it requires different compilation flags, thereby significantly changing
-# how it is compiled within the workspace, causing the whole workspace to be compiled from scratch
-# this way, this is a stand-alone package that compiles independently of the others.
-exclude = ["arrow-pyarrow-integration-testing"]
diff --git a/LICENSE.txt b/LICENSE.txt
deleted file mode 100644
index 4cec07f..0000000
--- a/LICENSE.txt
+++ /dev/null
@@ -1,2220 +0,0 @@
-
- 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.
-
---------------------------------------------------------------------------------
-
-src/plasma/fling.cc and src/plasma/fling.h: Apache 2.0
-
-Copyright 2013 Sharvil Nanavati
-
-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.
-
---------------------------------------------------------------------------------
-
-src/plasma/thirdparty/ae: Modified / 3-Clause BSD
-
-Copyright (c) 2006-2010, Salvatore Sanfilippo <antirez at gmail dot com>
-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 Redis 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.
-
---------------------------------------------------------------------------------
-
-src/plasma/thirdparty/dlmalloc.c: CC0
-
-This is a version (aka dlmalloc) of malloc/free/realloc written by
-Doug Lea and released to the public domain, as explained at
-http://creativecommons.org/publicdomain/zero/1.0/ Send questions,
-comments, complaints, performance data, etc to dl@cs.oswego.edu
-
---------------------------------------------------------------------------------
-
-src/plasma/common.cc (some portions)
-
-Copyright (c) Austin Appleby (aappleby (AT) gmail)
-
-Some portions of this file are derived from code in the MurmurHash project
-
-All code is released to the public domain. For business purposes, Murmurhash is
-under the MIT license.
-
-https://sites.google.com/site/murmurhash/
-
---------------------------------------------------------------------------------
-
-src/arrow/util (some portions): Apache 2.0, and 3-clause BSD
-
-Some portions of this module are derived from code in the Chromium project,
-copyright (c) Google inc and (c) The Chromium Authors and licensed under the
-Apache 2.0 License or the under the 3-clause BSD license:
-
- Copyright (c) 2013 The Chromium Authors. 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.
-
---------------------------------------------------------------------------------
-
-This project includes code from Daniel Lemire's FrameOfReference project.
-
-https://github.com/lemire/FrameOfReference/blob/6ccaf9e97160f9a3b299e23a8ef739e711ef0c71/src/bpacking.cpp
-
-Copyright: 2013 Daniel Lemire
-Home page: http://lemire.me/en/
-Project page: https://github.com/lemire/FrameOfReference
-License: Apache License Version 2.0 http://www.apache.org/licenses/LICENSE-2.0
-
---------------------------------------------------------------------------------
-
-This project includes code from the TensorFlow project
-
-Copyright 2015 The TensorFlow Authors. All Rights Reserved.
-
-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 project includes code from the NumPy project.
-
-https://github.com/numpy/numpy/blob/e1f191c46f2eebd6cb892a4bfe14d9dd43a06c4e/numpy/core/src/multiarray/multiarraymodule.c#L2910
-
-https://github.com/numpy/numpy/blob/68fd82271b9ea5a9e50d4e761061dfcca851382a/numpy/core/src/multiarray/datetime.c
-
-Copyright (c) 2005-2017, NumPy Developers.
-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 the NumPy Developers nor the names of any
- 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.
-
---------------------------------------------------------------------------------
-
-This project includes code from the Boost project
-
-Boost Software License - Version 1.0 - August 17th, 2003
-
-Permission is hereby granted, free of charge, to any person or organization
-obtaining a copy of the software and accompanying documentation covered by
-this license (the "Software") to use, reproduce, display, distribute,
-execute, and transmit the Software, and to prepare derivative works of the
-Software, and to permit third-parties to whom the Software is furnished to
-do so, all subject to the following:
-
-The copyright notices in the Software and this entire statement, including
-the above license grant, this restriction and the following disclaimer,
-must be included in all copies of the Software, in whole or in part, and
-all derivative works of the Software, unless such copies or derivative
-works are solely in the form of machine-executable object code generated by
-a source language processor.
-
-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, TITLE AND NON-INFRINGEMENT. IN NO EVENT
-SHALL THE COPYRIGHT HOLDERS OR ANYONE DISTRIBUTING THE SOFTWARE BE LIABLE
-FOR ANY DAMAGES OR OTHER LIABILITY, WHETHER IN CONTRACT, TORT OR OTHERWISE,
-ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
-DEALINGS IN THE SOFTWARE.
-
---------------------------------------------------------------------------------
-
-This project includes code from the FlatBuffers project
-
-Copyright 2014 Google Inc.
-
-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 project includes code from the tslib project
-
-Copyright 2015 Microsoft Corporation. All rights reserved.
-
-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 project includes code from the jemalloc project
-
-https://github.com/jemalloc/jemalloc
-
-Copyright (C) 2002-2017 Jason Evans <jasone@canonware.com>.
-All rights reserved.
-Copyright (C) 2007-2012 Mozilla Foundation. All rights reserved.
-Copyright (C) 2009-2017 Facebook, 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:
-1. Redistributions of source code must retain the above copyright notice(s),
- this list of conditions and the following disclaimer.
-2. Redistributions in binary form must reproduce the above copyright notice(s),
- 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 HOLDER(S) ``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(S) 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 project includes code from the Go project, BSD 3-clause license + PATENTS
-weak patent termination clause
-(https://github.com/golang/go/blob/master/PATENTS).
-
-Copyright (c) 2009 The Go Authors. 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.
-
---------------------------------------------------------------------------------
-
-This project includes code from the hs2client
-
-https://github.com/cloudera/hs2client
-
-Copyright 2016 Cloudera Inc.
-
-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.
-
---------------------------------------------------------------------------------
-
-The script ci/scripts/util_wait_for_it.sh has the following license
-
-Copyright (c) 2016 Giles Hall
-
-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.
-
---------------------------------------------------------------------------------
-
-The script r/configure has the following license (MIT)
-
-Copyright (c) 2017, Jeroen Ooms and Jim Hester
-
-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.
-
---------------------------------------------------------------------------------
-
-cpp/src/arrow/util/logging.cc, cpp/src/arrow/util/logging.h and
-cpp/src/arrow/util/logging-test.cc are adapted from
-Ray Project (https://github.com/ray-project/ray) (Apache 2.0).
-
-Copyright (c) 2016 Ray Project (https://github.com/ray-project/ray)
-
-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.
-
---------------------------------------------------------------------------------
-The files cpp/src/arrow/vendored/datetime/date.h, cpp/src/arrow/vendored/datetime/tz.h,
-cpp/src/arrow/vendored/datetime/tz_private.h, cpp/src/arrow/vendored/datetime/ios.h,
-cpp/src/arrow/vendored/datetime/ios.mm,
-cpp/src/arrow/vendored/datetime/tz.cpp are adapted from
-Howard Hinnant's date library (https://github.com/HowardHinnant/date)
-It is licensed under MIT license.
-
-The MIT License (MIT)
-Copyright (c) 2015, 2016, 2017 Howard Hinnant
-Copyright (c) 2016 Adrian Colomitchi
-Copyright (c) 2017 Florian Dang
-Copyright (c) 2017 Paul Thompson
-Copyright (c) 2018 Tomasz Kamiński
-
-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.
-
---------------------------------------------------------------------------------
-
-The file cpp/src/arrow/util/utf8.h includes code adapted from the page
- https://bjoern.hoehrmann.de/utf-8/decoder/dfa/
-with the following license (MIT)
-
-Copyright (c) 2008-2009 Bjoern Hoehrmann <bjoern@hoehrmann.de>
-
-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.
-
---------------------------------------------------------------------------------
-
-The file cpp/src/arrow/vendored/string_view.hpp has the following license
-
-Boost Software License - Version 1.0 - August 17th, 2003
-
-Permission is hereby granted, free of charge, to any person or organization
-obtaining a copy of the software and accompanying documentation covered by
-this license (the "Software") to use, reproduce, display, distribute,
-execute, and transmit the Software, and to prepare derivative works of the
-Software, and to permit third-parties to whom the Software is furnished to
-do so, all subject to the following:
-
-The copyright notices in the Software and this entire statement, including
-the above license grant, this restriction and the following disclaimer,
-must be included in all copies of the Software, in whole or in part, and
-all derivative works of the Software, unless such copies or derivative
-works are solely in the form of machine-executable object code generated by
-a source language processor.
-
-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, TITLE AND NON-INFRINGEMENT. IN NO EVENT
-SHALL THE COPYRIGHT HOLDERS OR ANYONE DISTRIBUTING THE SOFTWARE BE LIABLE
-FOR ANY DAMAGES OR OTHER LIABILITY, WHETHER IN CONTRACT, TORT OR OTHERWISE,
-ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
-DEALINGS IN THE SOFTWARE.
-
---------------------------------------------------------------------------------
-
-The files in cpp/src/arrow/vendored/xxhash/ have the following license
-(BSD 2-Clause License)
-
-xxHash Library
-Copyright (c) 2012-2014, Yann Collet
-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.
-
-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.
-
-You can contact the author at :
-- xxHash homepage: http://www.xxhash.com
-- xxHash source repository : https://github.com/Cyan4973/xxHash
-
---------------------------------------------------------------------------------
-
-The files in cpp/src/arrow/vendored/double-conversion/ have the following license
-(BSD 3-Clause License)
-
-Copyright 2006-2011, the V8 project authors. 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.
-
---------------------------------------------------------------------------------
-
-The files in cpp/src/arrow/vendored/uriparser/ have the following license
-(BSD 3-Clause License)
-
-uriparser - RFC 3986 URI parsing library
-
-Copyright (C) 2007, Weijia Song <songweijia@gmail.com>
-Copyright (C) 2007, Sebastian Pipping <sebastian@pipping.org>
-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 the <ORGANIZATION> 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.
-
---------------------------------------------------------------------------------
-
-The files under dev/tasks/conda-recipes have the following license
-
-BSD 3-clause license
-Copyright (c) 2015-2018, conda-forge
-All rights reserved.
-
-Redistribution and use in source and binary forms, with or without modification,
-are permitted provided that the following conditions are met:
-
-1. Redistributions of source code must retain the above copyright notice, this
- list of conditions and the following disclaimer.
-
-2. 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.
-
-3. Neither the name of the copyright holder 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 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.
-
---------------------------------------------------------------------------------
-
-The files in cpp/src/arrow/vendored/utfcpp/ have the following license
-
-Copyright 2006-2018 Nemanja Trifunovic
-
-Permission is hereby granted, free of charge, to any person or organization
-obtaining a copy of the software and accompanying documentation covered by
-this license (the "Software") to use, reproduce, display, distribute,
-execute, and transmit the Software, and to prepare derivative works of the
-Software, and to permit third-parties to whom the Software is furnished to
-do so, all subject to the following:
-
-The copyright notices in the Software and this entire statement, including
-the above license grant, this restriction and the following disclaimer,
-must be included in all copies of the Software, in whole or in part, and
-all derivative works of the Software, unless such copies or derivative
-works are solely in the form of machine-executable object code generated by
-a source language processor.
-
-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, TITLE AND NON-INFRINGEMENT. IN NO EVENT
-SHALL THE COPYRIGHT HOLDERS OR ANYONE DISTRIBUTING THE SOFTWARE BE LIABLE
-FOR ANY DAMAGES OR OTHER LIABILITY, WHETHER IN CONTRACT, TORT OR OTHERWISE,
-ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
-DEALINGS IN THE SOFTWARE.
-
---------------------------------------------------------------------------------
-
-This project includes code from Apache Kudu.
-
- * cpp/cmake_modules/CompilerInfo.cmake is based on Kudu's cmake_modules/CompilerInfo.cmake
-
-Copyright: 2016 The Apache Software Foundation.
-Home page: https://kudu.apache.org/
-License: http://www.apache.org/licenses/LICENSE-2.0
-
---------------------------------------------------------------------------------
-
-This project includes code from Apache Impala (incubating), formerly
-Impala. The Impala code and rights were donated to the ASF as part of the
-Incubator process after the initial code imports into Apache Parquet.
-
-Copyright: 2012 Cloudera, Inc.
-Copyright: 2016 The Apache Software Foundation.
-Home page: http://impala.apache.org/
-License: http://www.apache.org/licenses/LICENSE-2.0
-
---------------------------------------------------------------------------------
-
-This project includes code from Apache Aurora.
-
-* dev/release/{release,changelog,release-candidate} are based on the scripts from
- Apache Aurora
-
-Copyright: 2016 The Apache Software Foundation.
-Home page: https://aurora.apache.org/
-License: http://www.apache.org/licenses/LICENSE-2.0
-
---------------------------------------------------------------------------------
-
-This project includes code from the Google styleguide.
-
-* cpp/build-support/cpplint.py is based on the scripts from the Google styleguide.
-
-Copyright: 2009 Google Inc. All rights reserved.
-Homepage: https://github.com/google/styleguide
-License: 3-clause BSD
-
---------------------------------------------------------------------------------
-
-This project includes code from Snappy.
-
-* cpp/cmake_modules/{SnappyCMakeLists.txt,SnappyConfig.h} are based on code
- from Google's Snappy project.
-
-Copyright: 2009 Google Inc. All rights reserved.
-Homepage: https://github.com/google/snappy
-License: 3-clause BSD
-
---------------------------------------------------------------------------------
-
-This project includes code from the manylinux project.
-
-* python/manylinux1/scripts/{build_python.sh,python-tag-abi-tag.py,
- requirements.txt} are based on code from the manylinux project.
-
-Copyright: 2016 manylinux
-Homepage: https://github.com/pypa/manylinux
-License: The MIT License (MIT)
-
---------------------------------------------------------------------------------
-
-This project includes code from the cymove project:
-
-* python/pyarrow/includes/common.pxd includes code from the cymove project
-
-The MIT License (MIT)
-Copyright (c) 2019 Omer Ozarslan
-
-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.
-
---------------------------------------------------------------------------------
-
-The projects includes code from the Ursabot project under the dev/archery
-directory.
-
-License: BSD 2-Clause
-
-Copyright 2019 RStudio, Inc.
-
-Redistribution and use in source and binary forms, with or without
-modification, are permitted provided that the following conditions are met:
-
-1. Redistributions of source code must retain the above copyright notice, this
- list of conditions and the following disclaimer.
-
-2. 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 project include code from mingw-w64.
-
-* cpp/src/arrow/util/cpu-info.cc has a polyfill for mingw-w64 < 5
-
-Copyright (c) 2009 - 2013 by the mingw-w64 project
-Homepage: https://mingw-w64.org
-License: Zope Public License (ZPL) Version 2.1.
-
----------------------------------------------------------------------------------
-
-This project include code from Google's Asylo project.
-
-* cpp/src/arrow/result.h is based on status_or.h
-
-Copyright (c) Copyright 2017 Asylo authors
-Homepage: https://asylo.dev/
-License: Apache 2.0
-
---------------------------------------------------------------------------------
-
-This project includes code from Google's protobuf project
-
-* cpp/src/arrow/result.h ARROW_ASSIGN_OR_RAISE is based off ASSIGN_OR_RETURN
-
-Copyright 2008 Google Inc. All rights reserved.
-Homepage: https://developers.google.com/protocol-buffers/
-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.
- * 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.
-
---------------------------------------------------------------------------------
-
-3rdparty dependency LLVM is statically linked in certain binary distributions.
-Additionally some sections of source code have been derived from sources in LLVM
-and have been clearly labeled as such. LLVM has the following license:
-
-==============================================================================
-The LLVM Project is under the Apache License v2.0 with LLVM Exceptions:
-==============================================================================
-
- 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.
-
-
----- LLVM Exceptions to the Apache 2.0 License ----
-
-As an exception, if, as a result of your compiling your source code, portions
-of this Software are embedded into an Object form of such source code, you
-may redistribute such embedded portions in such Object form without complying
-with the conditions of Sections 4(a), 4(b) and 4(d) of the License.
-
-In addition, if you combine or link compiled forms of this Software with
-software that is licensed under the GPLv2 ("Combined Software") and if a
-court of competent jurisdiction determines that the patent provision (Section
-3), the indemnity provision (Section 9) or other Section of the License
-conflicts with the conditions of the GPLv2, you may retroactively and
-prospectively choose to deem waived or otherwise exclude such Section(s) of
-the License, but only in their entirety and only with respect to the Combined
-Software.
-
-==============================================================================
-Software from third parties included in the LLVM Project:
-==============================================================================
-The LLVM Project contains third party software which is under different license
-terms. All such code will be identified clearly using at least one of two
-mechanisms:
-1) It will be in a separate directory tree with its own `LICENSE.txt` or
- `LICENSE` file at the top containing the specific license and restrictions
- which apply to that software, or
-2) It will contain specific license and restriction terms at the top of every
- file.
-
---------------------------------------------------------------------------------
-
-3rdparty dependency gRPC is statically linked in certain binary
-distributions, like the python wheels. gRPC has the following license:
-
-Copyright 2014 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.
-
---------------------------------------------------------------------------------
-
-3rdparty dependency Apache Thrift is statically linked in certain binary
-distributions, like the python wheels. Apache Thrift has the following license:
-
-Apache Thrift
-Copyright (C) 2006 - 2019, The Apache Software Foundation
-
-This product includes software developed at
-The Apache Software Foundation (http://www.apache.org/).
-
-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.
-
---------------------------------------------------------------------------------
-
-3rdparty dependency Apache ORC is statically linked in certain binary
-distributions, like the python wheels. Apache ORC has the following license:
-
-Apache ORC
-Copyright 2013-2019 The Apache Software Foundation
-
-This product includes software developed by The Apache Software
-Foundation (http://www.apache.org/).
-
-This product includes software developed by Hewlett-Packard:
-(c) Copyright [2014-2015] Hewlett-Packard Development Company, L.P
-
-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.
-
---------------------------------------------------------------------------------
-
-3rdparty dependency zstd is statically linked in certain binary
-distributions, like the python wheels. ZSTD has the following license:
-
-BSD License
-
-For Zstandard software
-
-Copyright (c) 2016-present, Facebook, 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 Facebook 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 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.
-
---------------------------------------------------------------------------------
-
-3rdparty dependency lz4 is statically linked in certain binary
-distributions, like the python wheels. lz4 has the following license:
-
-LZ4 Library
-Copyright (c) 2011-2016, Yann Collet
-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.
-
-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.
-
---------------------------------------------------------------------------------
-
-3rdparty dependency Brotli is statically linked in certain binary
-distributions, like the python wheels. Brotli has the following license:
-
-Copyright (c) 2009, 2010, 2013-2016 by the Brotli Authors.
-
-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.
-
---------------------------------------------------------------------------------
-
-3rdparty dependency rapidjson is statically linked in certain binary
-distributions, like the python wheels. rapidjson and its dependencies have the
-following licenses:
-
-Tencent is pleased to support the open source community by making RapidJSON
-available.
-
-Copyright (C) 2015 THL A29 Limited, a Tencent company, and Milo Yip.
-All rights reserved.
-
-If you have downloaded a copy of the RapidJSON binary from Tencent, please note
-that the RapidJSON binary is licensed under the MIT License.
-If you have downloaded a copy of the RapidJSON source code from Tencent, please
-note that RapidJSON source code is licensed under the MIT License, except for
-the third-party components listed below which are subject to different license
-terms. Your integration of RapidJSON into your own projects may require
-compliance with the MIT License, as well as the other licenses applicable to
-the third-party components included within RapidJSON. To avoid the problematic
-JSON license in your own projects, it's sufficient to exclude the
-bin/jsonchecker/ directory, as it's the only code under the JSON license.
-A copy of the MIT License is included in this file.
-
-Other dependencies and licenses:
-
- Open Source Software Licensed Under the BSD License:
- --------------------------------------------------------------------
-
- The msinttypes r29
- Copyright (c) 2006-2013 Alexander Chemeris
- 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 copyright holder 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 REGENTS 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 REGENTS AND 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.
-
- Open Source Software Licensed Under the JSON License:
- --------------------------------------------------------------------
-
- json.org
- Copyright (c) 2002 JSON.org
- All Rights Reserved.
-
- JSON_checker
- Copyright (c) 2002 JSON.org
- All Rights Reserved.
-
-
- Terms of the JSON License:
- ---------------------------------------------------
-
- 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 shall be used for Good, not Evil.
-
- 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.
-
-
- Terms of the MIT License:
- --------------------------------------------------------------------
-
- 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.
-
---------------------------------------------------------------------------------
-
-3rdparty dependency snappy is statically linked in certain binary
-distributions, like the python wheels. snappy has the following license:
-
-Copyright 2011, 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.
-
-===
-
-Some of the benchmark data in testdata/ is licensed differently:
-
- - fireworks.jpeg is Copyright 2013 Steinar H. Gunderson, and
- is licensed under the Creative Commons Attribution 3.0 license
- (CC-BY-3.0). See https://creativecommons.org/licenses/by/3.0/
- for more information.
-
- - kppkn.gtb is taken from the Gaviota chess tablebase set, and
- is licensed under the MIT License. See
- https://sites.google.com/site/gaviotachessengine/Home/endgame-tablebases-1
- for more information.
-
- - paper-100k.pdf is an excerpt (bytes 92160 to 194560) from the paper
- “Combinatorial Modeling of Chromatin Features Quantitatively Predicts DNA
- Replication Timing in _Drosophila_” by Federico Comoglio and Renato Paro,
- which is licensed under the CC-BY license. See
- http://www.ploscompbiol.org/static/license for more ifnormation.
-
- - alice29.txt, asyoulik.txt, plrabn12.txt and lcet10.txt are from Project
- Gutenberg. The first three have expired copyrights and are in the public
- domain; the latter does not have expired copyright, but is still in the
- public domain according to the license information
- (http://www.gutenberg.org/ebooks/53).
-
---------------------------------------------------------------------------------
-
-3rdparty dependency gflags is statically linked in certain binary
-distributions, like the python wheels. gflags has the following license:
-
-Copyright (c) 2006, 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.
-
---------------------------------------------------------------------------------
-
-3rdparty dependency glog is statically linked in certain binary
-distributions, like the python wheels. glog has the following license:
-
-Copyright (c) 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.
-
-
-A function gettimeofday in utilities.cc is based on
-
-http://www.google.com/codesearch/p?hl=en#dR3YEbitojA/COPYING&q=GetSystemTimeAsFileTime%20license:bsd
-
-The license of this code is:
-
-Copyright (c) 2003-2008, Jouni Malinen <j@w1.fi> and contributors
-All Rights Reserved.
-
-Redistribution and use in source and binary forms, with or without
-modification, are permitted provided that the following conditions are
-met:
-
-1. Redistributions of source code must retain the above copyright
- notice, this list of conditions and the following disclaimer.
-
-2. 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.
-
-3. Neither the name(s) of the above-listed copyright holder(s) 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.
-
---------------------------------------------------------------------------------
-
-3rdparty dependency re2 is statically linked in certain binary
-distributions, like the python wheels. re2 has the following license:
-
-Copyright (c) 2009 The RE2 Authors. 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.
-
---------------------------------------------------------------------------------
-
-3rdparty dependency c-ares is statically linked in certain binary
-distributions, like the python wheels. c-ares has the following license:
-
-# c-ares license
-
-Copyright (c) 2007 - 2018, Daniel Stenberg with many contributors, see AUTHORS
-file.
-
-Copyright 1998 by the Massachusetts Institute of Technology.
-
-Permission to use, copy, modify, and distribute this software and its
-documentation for any purpose and without fee is hereby granted, provided that
-the above copyright notice appear in all copies and that both that copyright
-notice and this permission notice appear in supporting documentation, and that
-the name of M.I.T. not be used in advertising or publicity pertaining to
-distribution of the software without specific, written prior permission.
-M.I.T. makes no representations about the suitability of this software for any
-purpose. It is provided "as is" without express or implied warranty.
-
---------------------------------------------------------------------------------
-
-3rdparty dependency zlib is redistributed as a dynamically linked shared
-library in certain binary distributions, like the python wheels. In the future
-this will likely change to static linkage. zlib has the following license:
-
-zlib.h -- interface of the 'zlib' general purpose compression library
- version 1.2.11, January 15th, 2017
-
- Copyright (C) 1995-2017 Jean-loup Gailly and Mark Adler
-
- This software is provided 'as-is', without any express or implied
- warranty. In no event will the authors be held liable for any damages
- arising from the use of this software.
-
- Permission is granted to anyone to use this software for any purpose,
- including commercial applications, and to alter it and redistribute it
- freely, subject to the following restrictions:
-
- 1. The origin of this software must not be misrepresented; you must not
- claim that you wrote the original software. If you use this software
- in a product, an acknowledgment in the product documentation would be
- appreciated but is not required.
- 2. Altered source versions must be plainly marked as such, and must not be
- misrepresented as being the original software.
- 3. This notice may not be removed or altered from any source distribution.
-
- Jean-loup Gailly Mark Adler
- jloup@gzip.org madler@alumni.caltech.edu
-
---------------------------------------------------------------------------------
-
-3rdparty dependency openssl is redistributed as a dynamically linked shared
-library in certain binary distributions, like the python wheels. openssl
-preceding version 3 has the following license:
-
- LICENSE ISSUES
- ==============
-
- The OpenSSL toolkit stays under a double license, i.e. both the conditions of
- the OpenSSL License and the original SSLeay license apply to the toolkit.
- See below for the actual license texts.
-
- OpenSSL License
- ---------------
-
-/* ====================================================================
- * Copyright (c) 1998-2019 The OpenSSL Project. All rights reserved.
- *
- * Redistribution and use in source and binary forms, with or without
- * modification, are permitted provided that the following conditions
- * are met:
- *
- * 1. Redistributions of source code must retain the above copyright
- * notice, this list of conditions and the following disclaimer.
- *
- * 2. 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.
- *
- * 3. All advertising materials mentioning features or use of this
- * software must display the following acknowledgment:
- * "This product includes software developed by the OpenSSL Project
- * for use in the OpenSSL Toolkit. (http://www.openssl.org/)"
- *
- * 4. The names "OpenSSL Toolkit" and "OpenSSL Project" must not be used to
- * endorse or promote products derived from this software without
- * prior written permission. For written permission, please contact
- * openssl-core@openssl.org.
- *
- * 5. Products derived from this software may not be called "OpenSSL"
- * nor may "OpenSSL" appear in their names without prior written
- * permission of the OpenSSL Project.
- *
- * 6. Redistributions of any form whatsoever must retain the following
- * acknowledgment:
- * "This product includes software developed by the OpenSSL Project
- * for use in the OpenSSL Toolkit (http://www.openssl.org/)"
- *
- * THIS SOFTWARE IS PROVIDED BY THE OpenSSL PROJECT ``AS IS'' AND ANY
- * EXPRESSED 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 OpenSSL PROJECT OR
- * ITS 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 product includes cryptographic software written by Eric Young
- * (eay@cryptsoft.com). This product includes software written by Tim
- * Hudson (tjh@cryptsoft.com).
- *
- */
-
- Original SSLeay License
- -----------------------
-
-/* Copyright (C) 1995-1998 Eric Young (eay@cryptsoft.com)
- * All rights reserved.
- *
- * This package is an SSL implementation written
- * by Eric Young (eay@cryptsoft.com).
- * The implementation was written so as to conform with Netscapes SSL.
- *
- * This library is free for commercial and non-commercial use as long as
- * the following conditions are aheared to. The following conditions
- * apply to all code found in this distribution, be it the RC4, RSA,
- * lhash, DES, etc., code; not just the SSL code. The SSL documentation
- * included with this distribution is covered by the same copyright terms
- * except that the holder is Tim Hudson (tjh@cryptsoft.com).
- *
- * Copyright remains Eric Young's, and as such any Copyright notices in
- * the code are not to be removed.
- * If this package is used in a product, Eric Young should be given attribution
- * as the author of the parts of the library used.
- * This can be in the form of a textual message at program startup or
- * in documentation (online or textual) provided with the package.
- *
- * Redistribution and use in source and binary forms, with or without
- * modification, are permitted provided that the following conditions
- * are met:
- * 1. Redistributions of source code must retain the copyright
- * notice, this list of conditions and the following disclaimer.
- * 2. 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.
- * 3. All advertising materials mentioning features or use of this software
- * must display the following acknowledgement:
- * "This product includes cryptographic software written by
- * Eric Young (eay@cryptsoft.com)"
- * The word 'cryptographic' can be left out if the rouines from the library
- * being used are not cryptographic related :-).
- * 4. If you include any Windows specific code (or a derivative thereof) from
- * the apps directory (application code) you must include an acknowledgement:
- * "This product includes software written by Tim Hudson (tjh@cryptsoft.com)"
- *
- * THIS SOFTWARE IS PROVIDED BY ERIC YOUNG ``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 AUTHOR 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.
- *
- * The licence and distribution terms for any publically available version or
- * derivative of this code cannot be changed. i.e. this code cannot simply be
- * copied and put under another distribution licence
- * [including the GNU Public Licence.]
- */
-
---------------------------------------------------------------------------------
-
-This project includes code from the rtools-backports project.
-
-* ci/scripts/PKGBUILD and ci/scripts/r_windows_build.sh are based on code
- from the rtools-backports project.
-
-Copyright: Copyright (c) 2013 - 2019, Алексей and Jeroen Ooms.
-All rights reserved.
-Homepage: https://github.com/r-windows/rtools-backports
-License: 3-clause BSD
-
---------------------------------------------------------------------------------
-
-Some code from pandas has been adapted for the pyarrow codebase. pandas is
-available under the 3-clause BSD license, which follows:
-
-pandas license
-==============
-
-Copyright (c) 2011-2012, Lambda Foundry, Inc. and PyData Development Team
-All rights reserved.
-
-Copyright (c) 2008-2011 AQR Capital Management, LLC
-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 the copyright holder nor the names of any
- 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 HOLDER 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.
-
---------------------------------------------------------------------------------
-
-Some bits from DyND, in particular aspects of the build system, have been
-adapted from libdynd and dynd-python under the terms of the BSD 2-clause
-license
-
-The BSD 2-Clause License
-
- Copyright (C) 2011-12, Dynamic NDArray Developers
- 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.
-
- 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.
-
-Dynamic NDArray Developers list:
-
- * Mark Wiebe
- * Continuum Analytics
-
---------------------------------------------------------------------------------
-
-Some source code from Ibis (https://github.com/cloudera/ibis) has been adapted
-for PyArrow. Ibis is released under the Apache License, Version 2.0.
-
---------------------------------------------------------------------------------
-
-This project includes code from the autobrew project.
-
-* r/tools/autobrew and dev/tasks/homebrew-formulae/autobrew/apache-arrow.rb
- are based on code from the autobrew project.
-
-Copyright (c) 2019, Jeroen Ooms
-License: MIT
-Homepage: https://github.com/jeroen/autobrew
-
---------------------------------------------------------------------------------
-
-dev/tasks/homebrew-formulae/apache-arrow.rb has the following license:
-
-BSD 2-Clause License
-
-Copyright (c) 2009-present, Homebrew contributors
-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.
-
-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.
-
-----------------------------------------------------------------------
-
-cpp/src/arrow/vendored/base64.cpp has the following license
-
-ZLIB License
-
-Copyright (C) 2004-2017 René Nyffenegger
-
-This source code is provided 'as-is', without any express or implied
-warranty. In no event will the author be held liable for any damages arising
-from the use of this software.
-
-Permission is granted to anyone to use this software for any purpose, including
-commercial applications, and to alter it and redistribute it freely, subject to
-the following restrictions:
-
-1. The origin of this source code must not be misrepresented; you must not
- claim that you wrote the original source code. If you use this source code
- in a product, an acknowledgment in the product documentation would be
- appreciated but is not required.
-
-2. Altered source versions must be plainly marked as such, and must not be
- misrepresented as being the original source code.
-
-3. This notice may not be removed or altered from any source distribution.
-
-René Nyffenegger rene.nyffenegger@adp-gmbh.ch
-
---------------------------------------------------------------------------------
-
-The file cpp/src/arrow/vendored/optional.hpp has the following license
-
-Boost Software License - Version 1.0 - August 17th, 2003
-
-Permission is hereby granted, free of charge, to any person or organization
-obtaining a copy of the software and accompanying documentation covered by
-this license (the "Software") to use, reproduce, display, distribute,
-execute, and transmit the Software, and to prepare derivative works of the
-Software, and to permit third-parties to whom the Software is furnished to
-do so, all subject to the following:
-
-The copyright notices in the Software and this entire statement, including
-the above license grant, this restriction and the following disclaimer,
-must be included in all copies of the Software, in whole or in part, and
-all derivative works of the Software, unless such copies or derivative
-works are solely in the form of machine-executable object code generated by
-a source language processor.
-
-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, TITLE AND NON-INFRINGEMENT. IN NO EVENT
-SHALL THE COPYRIGHT HOLDERS OR ANYONE DISTRIBUTING THE SOFTWARE BE LIABLE
-FOR ANY DAMAGES OR OTHER LIABILITY, WHETHER IN CONTRACT, TORT OR OTHERWISE,
-ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
-DEALINGS IN THE SOFTWARE.
-
---------------------------------------------------------------------------------
-
-This project includes code from Folly.
-
- * cpp/src/arrow/vendored/ProducerConsumerQueue.h
-
-is based on Folly's
-
- * folly/Portability.h
- * folly/lang/Align.h
- * folly/ProducerConsumerQueue.h
-
-Copyright: Copyright (c) Facebook, Inc. and its affiliates.
-Home page: https://github.com/facebook/folly
-License: http://www.apache.org/licenses/LICENSE-2.0
-
---------------------------------------------------------------------------------
-
-The file cpp/src/arrow/vendored/musl/strptime.c has the following license
-
-Copyright © 2005-2020 Rich Felker, et al.
-
-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.
-
---------------------------------------------------------------------------------
-
-The file cpp/cmake_modules/BuildUtils.cmake contains code from
-
-https://gist.github.com/cristianadam/ef920342939a89fae3e8a85ca9459b49
-
-which is made available under the MIT license
-
-Copyright (c) 2019 Cristian Adam
-
-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.
-
---------------------------------------------------------------------------------
-
-The files in cpp/src/arrow/vendored/portable-snippets/ contain code from
-
-https://github.com/nemequ/portable-snippets
-
-and have the following copyright notice:
-
-Each source file contains a preamble explaining the license situation
-for that file, which takes priority over this file. With the
-exception of some code pulled in from other repositories (such as
-µnit, an MIT-licensed project which is used for testing), the code is
-public domain, released using the CC0 1.0 Universal dedication (*).
-
-(*) https://creativecommons.org/publicdomain/zero/1.0/legalcode
-
---------------------------------------------------------------------------------
-
-The files in cpp/src/arrow/vendored/fast_float/ contain code from
-
-https://github.com/lemire/fast_float
-
-which is made available under the Apache License 2.0.
-
---------------------------------------------------------------------------------
-
-The file python/pyarrow/vendored/version.py contains code from
-
-https://github.com/pypa/packaging/
-
-which is made available under both the Apache license v2.0 and the
-BSD 2-clause license.
diff --git a/NOTICE.txt b/NOTICE.txt
deleted file mode 100644
index a609791..0000000
--- a/NOTICE.txt
+++ /dev/null
@@ -1,84 +0,0 @@
-Apache Arrow
-Copyright 2016-2019 The Apache Software Foundation
-
-This product includes software developed at
-The Apache Software Foundation (http://www.apache.org/).
-
-This product includes software from the SFrame project (BSD, 3-clause).
-* Copyright (C) 2015 Dato, Inc.
-* Copyright (c) 2009 Carnegie Mellon University.
-
-This product includes software from the Feather project (Apache 2.0)
-https://github.com/wesm/feather
-
-This product includes software from the DyND project (BSD 2-clause)
-https://github.com/libdynd
-
-This product includes software from the LLVM project
- * distributed under the University of Illinois Open Source
-
-This product includes software from the google-lint project
- * Copyright (c) 2009 Google Inc. All rights reserved.
-
-This product includes software from the mman-win32 project
- * Copyright https://code.google.com/p/mman-win32/
- * Licensed under the MIT License;
-
-This product includes software from the LevelDB project
- * Copyright (c) 2011 The LevelDB Authors. All rights reserved.
- * Use of this source code is governed by a BSD-style license that can be
- * Moved from Kudu http://github.com/cloudera/kudu
-
-This product includes software from the CMake project
- * Copyright 2001-2009 Kitware, Inc.
- * Copyright 2012-2014 Continuum Analytics, Inc.
- * All rights reserved.
-
-This product includes software from https://github.com/matthew-brett/multibuild (BSD 2-clause)
- * Copyright (c) 2013-2016, Matt Terry and Matthew Brett; all rights reserved.
-
-This product includes software from the Ibis project (Apache 2.0)
- * Copyright (c) 2015 Cloudera, Inc.
- * https://github.com/cloudera/ibis
-
-This product includes software from Dremio (Apache 2.0)
- * Copyright (C) 2017-2018 Dremio Corporation
- * https://github.com/dremio/dremio-oss
-
-This product includes software from Google Guava (Apache 2.0)
- * Copyright (C) 2007 The Guava Authors
- * https://github.com/google/guava
-
-This product include software from CMake (BSD 3-Clause)
- * CMake - Cross Platform Makefile Generator
- * Copyright 2000-2019 Kitware, Inc. and Contributors
-
-The web site includes files generated by Jekyll.
-
---------------------------------------------------------------------------------
-
-This product includes code from Apache Kudu, which includes the following in
-its NOTICE file:
-
- Apache Kudu
- Copyright 2016 The Apache Software Foundation
-
- This product includes software developed at
- The Apache Software Foundation (http://www.apache.org/).
-
- Portions of this software were developed at
- Cloudera, Inc (http://www.cloudera.com/).
-
---------------------------------------------------------------------------------
-
-This product includes code from Apache ORC, which includes the following in
-its NOTICE file:
-
- Apache ORC
- Copyright 2013-2019 The Apache Software Foundation
-
- This product includes software developed by The Apache Software
- Foundation (http://www.apache.org/).
-
- This product includes software developed by Hewlett-Packard:
- (c) Copyright [2014-2015] Hewlett-Packard Development Company, L.P
diff --git a/README.md b/README.md
deleted file mode 100644
index d10bf67..0000000
--- a/README.md
+++ /dev/null
@@ -1,186 +0,0 @@
-<!---
- Licensed to the Apache Software Foundation (ASF) under one
- or more contributor license agreements. See the NOTICE file
- distributed with this work for additional information
- regarding copyright ownership. The ASF licenses this file
- to you under the Apache License, Version 2.0 (the
- "License"); you may not use this file except in compliance
- with the License. You may obtain a copy of the License at
-
- http://www.apache.org/licenses/LICENSE-2.0
-
- Unless required by applicable law or agreed to in writing,
- software distributed under the License is distributed on an
- "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
- KIND, either express or implied. See the License for the
- specific language governing permissions and limitations
- under the License.
--->
-
-# Native Rust implementation of Apache Arrow
-
-[](https://codecov.io/gh/apache/arrow?branch=master)
-
-Welcome to the implementation of Arrow, the popular in-memory columnar format, in [Rust](https://www.rust-lang.org/).
-
-This part of the Arrow project is divided in 4 main components:
-
-| Crate | Description | Documentation |
-| ------------ | ------------------------------------------------------------------ | ------------------------------------------------------------------------------------- |
-| Arrow | Core functionality (memory layout, arrays, low level computations) | [(README)](arrow/README.md) |
-| Parquet | Parquet support | [(README)](parquet/README.md) |
-| Arrow-flight | Arrow data between processes | [(README)](arrow-flight/README.md) |
-| DataFusion | In-memory query engine with SQL support | [(README)](https://github.com/apache/arrow-datafusion/blob/master/README.md) |
-| Ballista | Distributed query execution | [(README)](https://github.com/apache/arrow-datafusion/blob/master/ballista/README.md) |
-
-Independently, they support a vast array of functionality for in-memory computations.
-
-Together, they allow users to write an SQL query or a `DataFrame` (using the `datafusion` crate), run it against a parquet file (using the `parquet` crate), evaluate it in-memory using Arrow's columnar format (using the `arrow` crate), and send to another process (using the `arrow-flight` crate).
-
-Generally speaking, the `arrow` crate offers functionality to develop code that uses Arrow arrays, and `datafusion` offers most operations typically found in SQL, with the notable exceptions of:
-
-- `join`
-- `window` functions
-
-There are too many features to enumerate here, but some notable mentions:
-
-- `Arrow` implements all formats in the specification except certain dictionaries
-- `Arrow` supports SIMD operations to some of its vertical operations
-- `DataFusion` supports `async` execution
-- `DataFusion` supports user-defined functions, aggregates, and whole execution nodes
-
-You can find more details about each crate in their respective READMEs.
-
-## Arrow Rust Community
-
-We use the official [ASF Slack](https://s.apache.org/slack-invite) for informal discussions and coordination. This is
-a great place to meet other contributors and get guidance on where to contribute. Join us in the `arrow-rust` channel.
-
-We use [ASF JIRA](https://issues.apache.org/jira/secure/Dashboard.jspa) as the system of record for new features
-and bug fixes and this plays a critical role in the release process.
-
-For design discussions we generally collaborate on Google documents and file a JIRA linking to the document.
-
-There is also a bi-weekly Rust-specific sync call for the Arrow Rust community. This is hosted on Google Meet
-at https://meet.google.com/ctp-yujs-aee on alternate Wednesday's at 09:00 US/Pacific, 12:00 US/Eastern. During
-US daylight savings time this corresponds to 16:00 UTC and at other times this is 17:00 UTC.
-
-## Developer's guide to Arrow Rust
-
-### How to compile
-
-This is a standard cargo project with workspaces. To build it, you need to have `rust` and `cargo`:
-
-```bash
-cd /rust && cargo build
-```
-
-You can also use rust's official docker image:
-
-```bash
-docker run --rm -v $(pwd)/rust:/rust -it rust /bin/bash -c "cd /rust && cargo build"
-```
-
-The command above assumes that are in the root directory of the project, not in the same
-directory as this README.md.
-
-You can also compile specific workspaces:
-
-```bash
-cd /rust/arrow && cargo build
-```
-
-### Git Submodules
-
-Before running tests and examples, it is necessary to set up the local development environment.
-
-The tests rely on test data that is contained in git submodules.
-
-To pull down this data run the following:
-
-```bash
-git submodule update --init
-```
-
-This populates data in two git submodules:
-
-- `../parquet_testing/data` (sourced from https://github.com/apache/parquet-testing.git)
-- `../testing` (sourced from https://github.com/apache/arrow-testing)
-
-By default, `cargo test` will look for these directories at their
-standard location. The following environment variables can be used to override the location:
-
-```bash
-# Optionaly specify a different location for test data
-export PARQUET_TEST_DATA=$(cd ../parquet-testing/data; pwd)
-export ARROW_TEST_DATA=$(cd ../testing/data; pwd)
-```
-
-From here on, this is a pure Rust project and `cargo` can be used to run tests, benchmarks, docs and examples as usual.
-
-### Running the tests
-
-Run tests using the Rust standard `cargo test` command:
-
-```bash
-# run all tests.
-cargo test
-
-
-# run only tests for the arrow crate
-cargo test -p arrow
-```
-
-## Code Formatting
-
-Our CI uses `rustfmt` to check code formatting. Before submitting a
-PR be sure to run the following and check for lint issues:
-
-```bash
-cargo +stable fmt --all -- --check
-```
-
-## Clippy Lints
-
-We recommend using `clippy` for checking lints during development. While we do not yet enforce `clippy` checks, we recommend not introducing new `clippy` errors or warnings.
-
-Run the following to check for clippy lints.
-
-```
-cargo clippy
-```
-
-If you use Visual Studio Code with the `rust-analyzer` plugin, you can enable `clippy` to run each time you save a file. See https://users.rust-lang.org/t/how-to-use-clippy-in-vs-code-with-rust-analyzer/41881.
-
-One of the concerns with `clippy` is that it often produces a lot of false positives, or that some recommendations may hurt readability. We do not have a policy of which lints are ignored, but if you disagree with a `clippy` lint, you may disable the lint and briefly justify it.
-
-Search for `allow(clippy::` in the codebase to identify lints that are ignored/allowed. We currently prefer ignoring lints on the lowest unit possible.
-
-- If you are introducing a line that returns a lint warning or error, you may disable the lint on that line.
-- If you have several lints on a function or module, you may disable the lint on the function or module.
-- If a lint is pervasive across multiple modules, you may disable it at the crate level.
-
-## Git Pre-Commit Hook
-
-We can use [git pre-commit hook](https://git-scm.com/book/en/v2/Customizing-Git-Git-Hooks) to automate various kinds of git pre-commit checking/formatting.
-
-Suppose you are in the root directory of the project.
-
-First check if the file already exists:
-
-```bash
-ls -l .git/hooks/pre-commit
-```
-
-If the file already exists, to avoid mistakenly **overriding**, you MAY have to check
-the link source or file content. Else if not exist, let's safely soft link [pre-commit.sh](pre-commit.sh) as file `.git/hooks/pre-commit`:
-
-```
-ln -s ../../rust/pre-commit.sh .git/hooks/pre-commit
-```
-
-If sometimes you want to commit without checking, just run `git commit` with `--no-verify`:
-
-```bash
-git commit --no-verify -m "... commit message ..."
-```
diff --git a/arrow-flight/Cargo.toml b/arrow-flight/Cargo.toml
deleted file mode 100644
index c6027f8..0000000
--- a/arrow-flight/Cargo.toml
+++ /dev/null
@@ -1,45 +0,0 @@
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-
-[package]
-name = "arrow-flight"
-description = "Apache Arrow Flight"
-version = "5.0.0-SNAPSHOT"
-edition = "2018"
-authors = ["Apache Arrow <dev@arrow.apache.org>"]
-homepage = "https://github.com/apache/arrow-rs"
-repository = "https://github.com/apache/arrow-rs"
-license = "Apache-2.0"
-
-[dependencies]
-arrow = { path = "../arrow", version = "5.0.0-SNAPSHOT" }
-tonic = "0.4"
-bytes = "1"
-prost = "0.7"
-prost-derive = "0.7"
-tokio = { version = "1.0", features = ["macros", "rt", "rt-multi-thread"] }
-futures = { version = "0.3", default-features = false, features = ["alloc"]}
-
-[build-dependencies]
-tonic-build = "0.4"
-# Pin specific version of the tonic-build dependencies to avoid auto-generated
-# (and checked in) arrow.flight.protocol.rs from changing
-proc-macro2 = "=1.0.24"
-
-#[lib]
-#name = "flight"
-#path = "src/lib.rs"
diff --git a/arrow-flight/README.md b/arrow-flight/README.md
deleted file mode 100644
index 4205ebb..0000000
--- a/arrow-flight/README.md
+++ /dev/null
@@ -1,28 +0,0 @@
-<!---
- Licensed to the Apache Software Foundation (ASF) under one
- or more contributor license agreements. See the NOTICE file
- distributed with this work for additional information
- regarding copyright ownership. The ASF licenses this file
- to you under the Apache License, Version 2.0 (the
- "License"); you may not use this file except in compliance
- with the License. You may obtain a copy of the License at
-
- http://www.apache.org/licenses/LICENSE-2.0
-
- Unless required by applicable law or agreed to in writing,
- software distributed under the License is distributed on an
- "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
- KIND, either express or implied. See the License for the
- specific language governing permissions and limitations
- under the License.
--->
-
-# Apache Arrow Flight
-
-[](https://crates.io/crates/arrow-flight)
-
-Apache Arrow Flight is a gRPC based protocol for exchanging Arrow data between processes. See the blog post [Introducing Apache Arrow Flight: A Framework for Fast Data Transport](https://arrow.apache.org/blog/2019/10/13/introducing-arrow-flight/) for more information.
-
-This crate simply provides the Rust implementation of the [Flight.proto](../../format/Flight.proto) gRPC protocol and provides an example that demonstrates how to build a Flight server implemented with Tonic.
-
-Note that building a Flight server also requires an implementation of Arrow IPC which is based on the Flatbuffers serialization framework. The Rust implementation of Arrow IPC is not yet complete although the generated Flatbuffers code is available as part of the core Arrow crate.
diff --git a/arrow-flight/build.rs b/arrow-flight/build.rs
deleted file mode 100644
index 1cbfceb..0000000
--- a/arrow-flight/build.rs
+++ /dev/null
@@ -1,54 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-use std::{
- env,
- fs::OpenOptions,
- io::{Read, Write},
- path::Path,
-};
-
-fn main() -> Result<(), Box<dyn std::error::Error>> {
- // override the build location, in order to check in the changes to proto files
- env::set_var("OUT_DIR", "src");
-
- // The current working directory can vary depending on how the project is being
- // built or released so we build an absolute path to the proto file
- let path = Path::new("../format/Flight.proto");
- if path.exists() {
- // avoid rerunning build if the file has not changed
- println!("cargo:rerun-if-changed=../format/Flight.proto");
-
- tonic_build::compile_protos("../format/Flight.proto")?;
- // read file contents to string
- let mut file = OpenOptions::new()
- .read(true)
- .open("src/arrow.flight.protocol.rs")?;
- let mut buffer = String::new();
- file.read_to_string(&mut buffer)?;
- // append warning that file was auto-generate
- let mut file = OpenOptions::new()
- .write(true)
- .truncate(true)
- .open("src/arrow.flight.protocol.rs")?;
- file.write_all("// This file was automatically generated through the build.rs script, and should not be edited.\n\n".as_bytes())?;
- file.write_all(buffer.as_bytes())?;
- }
-
- // As the proto file is checked in, the build should not fail if the file is not found
- Ok(())
-}
diff --git a/arrow-flight/examples/server.rs b/arrow-flight/examples/server.rs
deleted file mode 100644
index 75d0537..0000000
--- a/arrow-flight/examples/server.rs
+++ /dev/null
@@ -1,131 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-use std::pin::Pin;
-
-use futures::Stream;
-use tonic::transport::Server;
-use tonic::{Request, Response, Status, Streaming};
-
-use arrow_flight::{
- flight_service_server::FlightService, flight_service_server::FlightServiceServer,
- Action, ActionType, Criteria, Empty, FlightData, FlightDescriptor, FlightInfo,
- HandshakeRequest, HandshakeResponse, PutResult, SchemaResult, Ticket,
-};
-
-#[derive(Clone)]
-pub struct FlightServiceImpl {}
-
-#[tonic::async_trait]
-impl FlightService for FlightServiceImpl {
- type HandshakeStream = Pin<
- Box<dyn Stream<Item = Result<HandshakeResponse, Status>> + Send + Sync + 'static>,
- >;
- type ListFlightsStream =
- Pin<Box<dyn Stream<Item = Result<FlightInfo, Status>> + Send + Sync + 'static>>;
- type DoGetStream =
- Pin<Box<dyn Stream<Item = Result<FlightData, Status>> + Send + Sync + 'static>>;
- type DoPutStream =
- Pin<Box<dyn Stream<Item = Result<PutResult, Status>> + Send + Sync + 'static>>;
- type DoActionStream = Pin<
- Box<
- dyn Stream<Item = Result<arrow_flight::Result, Status>>
- + Send
- + Sync
- + 'static,
- >,
- >;
- type ListActionsStream =
- Pin<Box<dyn Stream<Item = Result<ActionType, Status>> + Send + Sync + 'static>>;
- type DoExchangeStream =
- Pin<Box<dyn Stream<Item = Result<FlightData, Status>> + Send + Sync + 'static>>;
-
- async fn handshake(
- &self,
- _request: Request<Streaming<HandshakeRequest>>,
- ) -> Result<Response<Self::HandshakeStream>, Status> {
- Err(Status::unimplemented("Not yet implemented"))
- }
-
- async fn list_flights(
- &self,
- _request: Request<Criteria>,
- ) -> Result<Response<Self::ListFlightsStream>, Status> {
- Err(Status::unimplemented("Not yet implemented"))
- }
-
- async fn get_flight_info(
- &self,
- _request: Request<FlightDescriptor>,
- ) -> Result<Response<FlightInfo>, Status> {
- Err(Status::unimplemented("Not yet implemented"))
- }
-
- async fn get_schema(
- &self,
- _request: Request<FlightDescriptor>,
- ) -> Result<Response<SchemaResult>, Status> {
- Err(Status::unimplemented("Not yet implemented"))
- }
-
- async fn do_get(
- &self,
- _request: Request<Ticket>,
- ) -> Result<Response<Self::DoGetStream>, Status> {
- Err(Status::unimplemented("Not yet implemented"))
- }
-
- async fn do_put(
- &self,
- _request: Request<Streaming<FlightData>>,
- ) -> Result<Response<Self::DoPutStream>, Status> {
- Err(Status::unimplemented("Not yet implemented"))
- }
-
- async fn do_action(
- &self,
- _request: Request<Action>,
- ) -> Result<Response<Self::DoActionStream>, Status> {
- Err(Status::unimplemented("Not yet implemented"))
- }
-
- async fn list_actions(
- &self,
- _request: Request<Empty>,
- ) -> Result<Response<Self::ListActionsStream>, Status> {
- Err(Status::unimplemented("Not yet implemented"))
- }
-
- async fn do_exchange(
- &self,
- _request: Request<Streaming<FlightData>>,
- ) -> Result<Response<Self::DoExchangeStream>, Status> {
- Err(Status::unimplemented("Not yet implemented"))
- }
-}
-
-#[tokio::main]
-async fn main() -> Result<(), Box<dyn std::error::Error>> {
- let addr = "[::1]:50051".parse()?;
- let service = FlightServiceImpl {};
-
- let svc = FlightServiceServer::new(service);
-
- Server::builder().add_service(svc).serve(addr).await?;
-
- Ok(())
-}
diff --git a/arrow-flight/src/arrow.flight.protocol.rs b/arrow-flight/src/arrow.flight.protocol.rs
deleted file mode 100644
index 5fce526..0000000
--- a/arrow-flight/src/arrow.flight.protocol.rs
+++ /dev/null
@@ -1,1039 +0,0 @@
-// This file was automatically generated through the build.rs script, and should not be edited.
-
-///
-/// The request that a client provides to a server on handshake.
-#[derive(Clone, PartialEq, ::prost::Message)]
-pub struct HandshakeRequest {
- ///
- /// A defined protocol version
- #[prost(uint64, tag = "1")]
- pub protocol_version: u64,
- ///
- /// Arbitrary auth/handshake info.
- #[prost(bytes = "vec", tag = "2")]
- pub payload: ::prost::alloc::vec::Vec<u8>,
-}
-#[derive(Clone, PartialEq, ::prost::Message)]
-pub struct HandshakeResponse {
- ///
- /// A defined protocol version
- #[prost(uint64, tag = "1")]
- pub protocol_version: u64,
- ///
- /// Arbitrary auth/handshake info.
- #[prost(bytes = "vec", tag = "2")]
- pub payload: ::prost::alloc::vec::Vec<u8>,
-}
-///
-/// A message for doing simple auth.
-#[derive(Clone, PartialEq, ::prost::Message)]
-pub struct BasicAuth {
- #[prost(string, tag = "2")]
- pub username: ::prost::alloc::string::String,
- #[prost(string, tag = "3")]
- pub password: ::prost::alloc::string::String,
-}
-#[derive(Clone, PartialEq, ::prost::Message)]
-pub struct Empty {}
-///
-/// Describes an available action, including both the name used for execution
-/// along with a short description of the purpose of the action.
-#[derive(Clone, PartialEq, ::prost::Message)]
-pub struct ActionType {
- #[prost(string, tag = "1")]
- pub r#type: ::prost::alloc::string::String,
- #[prost(string, tag = "2")]
- pub description: ::prost::alloc::string::String,
-}
-///
-/// A service specific expression that can be used to return a limited set
-/// of available Arrow Flight streams.
-#[derive(Clone, PartialEq, ::prost::Message)]
-pub struct Criteria {
- #[prost(bytes = "vec", tag = "1")]
- pub expression: ::prost::alloc::vec::Vec<u8>,
-}
-///
-/// An opaque action specific for the service.
-#[derive(Clone, PartialEq, ::prost::Message)]
-pub struct Action {
- #[prost(string, tag = "1")]
- pub r#type: ::prost::alloc::string::String,
- #[prost(bytes = "vec", tag = "2")]
- pub body: ::prost::alloc::vec::Vec<u8>,
-}
-///
-/// An opaque result returned after executing an action.
-#[derive(Clone, PartialEq, ::prost::Message)]
-pub struct Result {
- #[prost(bytes = "vec", tag = "1")]
- pub body: ::prost::alloc::vec::Vec<u8>,
-}
-///
-/// Wrap the result of a getSchema call
-#[derive(Clone, PartialEq, ::prost::Message)]
-pub struct SchemaResult {
- /// schema of the dataset as described in Schema.fbs::Schema.
- #[prost(bytes = "vec", tag = "1")]
- pub schema: ::prost::alloc::vec::Vec<u8>,
-}
-///
-/// 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.
-#[derive(Clone, PartialEq, ::prost::Message)]
-pub struct FlightDescriptor {
- #[prost(enumeration = "flight_descriptor::DescriptorType", tag = "1")]
- pub r#type: i32,
- ///
- /// Opaque value used to express a command. Should only be defined when
- /// type = CMD.
- #[prost(bytes = "vec", tag = "2")]
- pub cmd: ::prost::alloc::vec::Vec<u8>,
- ///
- /// List of strings identifying a particular dataset. Should only be defined
- /// when type = PATH.
- #[prost(string, repeated, tag = "3")]
- pub path: ::prost::alloc::vec::Vec<::prost::alloc::string::String>,
-}
-/// Nested message and enum types in `FlightDescriptor`.
-pub mod flight_descriptor {
- ///
- /// Describes what type of descriptor is defined.
- #[derive(
- Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord, ::prost::Enumeration,
- )]
- #[repr(i32)]
- pub 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,
- }
-}
-///
-/// The access coordinates for retrieval of a dataset. With a FlightInfo, a
-/// consumer is able to determine how to retrieve a dataset.
-#[derive(Clone, PartialEq, ::prost::Message)]
-pub struct FlightInfo {
- /// schema of the dataset as described in Schema.fbs::Schema.
- #[prost(bytes = "vec", tag = "1")]
- pub schema: ::prost::alloc::vec::Vec<u8>,
- ///
- /// The descriptor associated with this info.
- #[prost(message, optional, tag = "2")]
- pub flight_descriptor: ::core::option::Option<FlightDescriptor>,
- ///
- /// A list of endpoints associated with the flight. To consume the whole
- /// flight, all endpoints must be consumed.
- #[prost(message, repeated, tag = "3")]
- pub endpoint: ::prost::alloc::vec::Vec<FlightEndpoint>,
- /// Set these to -1 if unknown.
- #[prost(int64, tag = "4")]
- pub total_records: i64,
- #[prost(int64, tag = "5")]
- pub total_bytes: i64,
-}
-///
-/// A particular stream or split associated with a flight.
-#[derive(Clone, PartialEq, ::prost::Message)]
-pub struct FlightEndpoint {
- ///
- /// Token used to retrieve this stream.
- #[prost(message, optional, tag = "1")]
- pub ticket: ::core::option::Option<Ticket>,
- ///
- /// A list of URIs where this ticket can be redeemed. If the list is
- /// empty, the expectation is that the ticket can only be redeemed on the
- /// current service where the ticket was generated.
- #[prost(message, repeated, tag = "2")]
- pub location: ::prost::alloc::vec::Vec<Location>,
-}
-///
-/// A location where a Flight service will accept retrieval of a particular
-/// stream given a ticket.
-#[derive(Clone, PartialEq, ::prost::Message)]
-pub struct Location {
- #[prost(string, tag = "1")]
- pub uri: ::prost::alloc::string::String,
-}
-///
-/// An opaque identifier that the service can use to retrieve a particular
-/// portion of a stream.
-#[derive(Clone, PartialEq, ::prost::Message)]
-pub struct Ticket {
- #[prost(bytes = "vec", tag = "1")]
- pub ticket: ::prost::alloc::vec::Vec<u8>,
-}
-///
-/// A batch of Arrow data as part of a stream of batches.
-#[derive(Clone, PartialEq, ::prost::Message)]
-pub struct FlightData {
- ///
- /// The descriptor of the data. This is only relevant when a client is
- /// starting a new DoPut stream.
- #[prost(message, optional, tag = "1")]
- pub flight_descriptor: ::core::option::Option<FlightDescriptor>,
- ///
- /// Header for message data as described in Message.fbs::Message.
- #[prost(bytes = "vec", tag = "2")]
- pub data_header: ::prost::alloc::vec::Vec<u8>,
- ///
- /// Application-defined metadata.
- #[prost(bytes = "vec", tag = "3")]
- pub app_metadata: ::prost::alloc::vec::Vec<u8>,
- ///
- /// 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).
- #[prost(bytes = "vec", tag = "1000")]
- pub data_body: ::prost::alloc::vec::Vec<u8>,
-}
-///*
-/// The response message associated with the submission of a DoPut.
-#[derive(Clone, PartialEq, ::prost::Message)]
-pub struct PutResult {
- #[prost(bytes = "vec", tag = "1")]
- pub app_metadata: ::prost::alloc::vec::Vec<u8>,
-}
-#[doc = r" Generated client implementations."]
-pub mod flight_service_client {
- #![allow(unused_variables, dead_code, missing_docs)]
- use tonic::codegen::*;
- #[doc = ""]
- #[doc = " A flight service is an endpoint for retrieving or storing Arrow data. A"]
- #[doc = " flight service can expose one or more predefined endpoints that can be"]
- #[doc = " accessed using the Arrow Flight Protocol. Additionally, a flight service"]
- #[doc = " can expose a set of actions that are available."]
- pub struct FlightServiceClient<T> {
- inner: tonic::client::Grpc<T>,
- }
- impl FlightServiceClient<tonic::transport::Channel> {
- #[doc = r" Attempt to create a new client by connecting to a given endpoint."]
- pub async fn connect<D>(dst: D) -> Result<Self, tonic::transport::Error>
- where
- D: std::convert::TryInto<tonic::transport::Endpoint>,
- D::Error: Into<StdError>,
- {
- let conn = tonic::transport::Endpoint::new(dst)?.connect().await?;
- Ok(Self::new(conn))
- }
- }
- impl<T> FlightServiceClient<T>
- where
- T: tonic::client::GrpcService<tonic::body::BoxBody>,
- T::ResponseBody: Body + HttpBody + Send + 'static,
- T::Error: Into<StdError>,
- <T::ResponseBody as HttpBody>::Error: Into<StdError> + Send,
- {
- pub fn new(inner: T) -> Self {
- let inner = tonic::client::Grpc::new(inner);
- Self { inner }
- }
- pub fn with_interceptor(
- inner: T,
- interceptor: impl Into<tonic::Interceptor>,
- ) -> Self {
- let inner = tonic::client::Grpc::with_interceptor(inner, interceptor);
- Self { inner }
- }
- #[doc = ""]
- #[doc = " Handshake between client and server. Depending on the server, the"]
- #[doc = " handshake may be required to determine the token that should be used for"]
- #[doc = " future operations. Both request and response are streams to allow multiple"]
- #[doc = " round-trips depending on auth mechanism."]
- pub async fn handshake(
- &mut self,
- request: impl tonic::IntoStreamingRequest<Message = super::HandshakeRequest>,
- ) -> Result<
- tonic::Response<tonic::codec::Streaming<super::HandshakeResponse>>,
- tonic::Status,
- > {
- self.inner.ready().await.map_err(|e| {
- tonic::Status::new(
- tonic::Code::Unknown,
- format!("Service was not ready: {}", e.into()),
- )
- })?;
- let codec = tonic::codec::ProstCodec::default();
- let path = http::uri::PathAndQuery::from_static(
- "/arrow.flight.protocol.FlightService/Handshake",
- );
- self.inner
- .streaming(request.into_streaming_request(), path, codec)
- .await
- }
- #[doc = ""]
- #[doc = " Get a list of available streams given a particular criteria. Most flight"]
- #[doc = " services will expose one or more streams that are readily available for"]
- #[doc = " retrieval. This api allows listing the streams available for"]
- #[doc = " consumption. A user can also provide a criteria. The criteria can limit"]
- #[doc = " the subset of streams that can be listed via this interface. Each flight"]
- #[doc = " service allows its own definition of how to consume criteria."]
- pub async fn list_flights(
- &mut self,
- request: impl tonic::IntoRequest<super::Criteria>,
- ) -> Result<
- tonic::Response<tonic::codec::Streaming<super::FlightInfo>>,
- tonic::Status,
- > {
- self.inner.ready().await.map_err(|e| {
- tonic::Status::new(
- tonic::Code::Unknown,
- format!("Service was not ready: {}", e.into()),
- )
- })?;
- let codec = tonic::codec::ProstCodec::default();
- let path = http::uri::PathAndQuery::from_static(
- "/arrow.flight.protocol.FlightService/ListFlights",
- );
- self.inner
- .server_streaming(request.into_request(), path, codec)
- .await
- }
- #[doc = ""]
- #[doc = " For a given FlightDescriptor, get information about how the flight can be"]
- #[doc = " consumed. This is a useful interface if the consumer of the interface"]
- #[doc = " already can identify the specific flight to consume. This interface can"]
- #[doc = " also allow a consumer to generate a flight stream through a specified"]
- #[doc = " descriptor. For example, a flight descriptor might be something that"]
- #[doc = " includes a SQL statement or a Pickled Python operation that will be"]
- #[doc = " executed. In those cases, the descriptor will not be previously available"]
- #[doc = " within the list of available streams provided by ListFlights but will be"]
- #[doc = " available for consumption for the duration defined by the specific flight"]
- #[doc = " service."]
- pub async fn get_flight_info(
- &mut self,
- request: impl tonic::IntoRequest<super::FlightDescriptor>,
- ) -> Result<tonic::Response<super::FlightInfo>, tonic::Status> {
- self.inner.ready().await.map_err(|e| {
- tonic::Status::new(
- tonic::Code::Unknown,
- format!("Service was not ready: {}", e.into()),
- )
- })?;
- let codec = tonic::codec::ProstCodec::default();
- let path = http::uri::PathAndQuery::from_static(
- "/arrow.flight.protocol.FlightService/GetFlightInfo",
- );
- self.inner.unary(request.into_request(), path, codec).await
- }
- #[doc = ""]
- #[doc = " For a given FlightDescriptor, get the Schema as described in Schema.fbs::Schema"]
- #[doc = " This is used when a consumer needs the Schema of flight stream. Similar to"]
- #[doc = " GetFlightInfo this interface may generate a new flight that was not previously"]
- #[doc = " available in ListFlights."]
- pub async fn get_schema(
- &mut self,
- request: impl tonic::IntoRequest<super::FlightDescriptor>,
- ) -> Result<tonic::Response<super::SchemaResult>, tonic::Status> {
- self.inner.ready().await.map_err(|e| {
- tonic::Status::new(
- tonic::Code::Unknown,
- format!("Service was not ready: {}", e.into()),
- )
- })?;
- let codec = tonic::codec::ProstCodec::default();
- let path = http::uri::PathAndQuery::from_static(
- "/arrow.flight.protocol.FlightService/GetSchema",
- );
- self.inner.unary(request.into_request(), path, codec).await
- }
- #[doc = ""]
- #[doc = " Retrieve a single stream associated with a particular descriptor"]
- #[doc = " associated with the referenced ticket. A Flight can be composed of one or"]
- #[doc = " more streams where each stream can be retrieved using a separate opaque"]
- #[doc = " ticket that the flight service uses for managing a collection of streams."]
- pub async fn do_get(
- &mut self,
- request: impl tonic::IntoRequest<super::Ticket>,
- ) -> Result<
- tonic::Response<tonic::codec::Streaming<super::FlightData>>,
- tonic::Status,
- > {
- self.inner.ready().await.map_err(|e| {
- tonic::Status::new(
- tonic::Code::Unknown,
- format!("Service was not ready: {}", e.into()),
- )
- })?;
- let codec = tonic::codec::ProstCodec::default();
- let path = http::uri::PathAndQuery::from_static(
- "/arrow.flight.protocol.FlightService/DoGet",
- );
- self.inner
- .server_streaming(request.into_request(), path, codec)
- .await
- }
- #[doc = ""]
- #[doc = " Push a stream to the flight service associated with a particular"]
- #[doc = " flight stream. This allows a client of a flight service to upload a stream"]
- #[doc = " of data. Depending on the particular flight service, a client consumer"]
- #[doc = " could be allowed to upload a single stream per descriptor or an unlimited"]
- #[doc = " number. In the latter, the service might implement a 'seal' action that"]
- #[doc = " can be applied to a descriptor once all streams are uploaded."]
- pub async fn do_put(
- &mut self,
- request: impl tonic::IntoStreamingRequest<Message = super::FlightData>,
- ) -> Result<
- tonic::Response<tonic::codec::Streaming<super::PutResult>>,
- tonic::Status,
- > {
- self.inner.ready().await.map_err(|e| {
- tonic::Status::new(
- tonic::Code::Unknown,
- format!("Service was not ready: {}", e.into()),
- )
- })?;
- let codec = tonic::codec::ProstCodec::default();
- let path = http::uri::PathAndQuery::from_static(
- "/arrow.flight.protocol.FlightService/DoPut",
- );
- self.inner
- .streaming(request.into_streaming_request(), path, codec)
- .await
- }
- #[doc = ""]
- #[doc = " Open a bidirectional data channel for a given descriptor. This"]
- #[doc = " allows clients to send and receive arbitrary Arrow data and"]
- #[doc = " application-specific metadata in a single logical stream. In"]
- #[doc = " contrast to DoGet/DoPut, this is more suited for clients"]
- #[doc = " offloading computation (rather than storage) to a Flight service."]
- pub async fn do_exchange(
- &mut self,
- request: impl tonic::IntoStreamingRequest<Message = super::FlightData>,
- ) -> Result<
- tonic::Response<tonic::codec::Streaming<super::FlightData>>,
- tonic::Status,
- > {
- self.inner.ready().await.map_err(|e| {
- tonic::Status::new(
- tonic::Code::Unknown,
- format!("Service was not ready: {}", e.into()),
- )
- })?;
- let codec = tonic::codec::ProstCodec::default();
- let path = http::uri::PathAndQuery::from_static(
- "/arrow.flight.protocol.FlightService/DoExchange",
- );
- self.inner
- .streaming(request.into_streaming_request(), path, codec)
- .await
- }
- #[doc = ""]
- #[doc = " Flight services can support an arbitrary number of simple actions in"]
- #[doc = " addition to the possible ListFlights, GetFlightInfo, DoGet, DoPut"]
- #[doc = " operations that are potentially available. DoAction allows a flight client"]
- #[doc = " to do a specific action against a flight service. An action includes"]
- #[doc = " opaque request and response objects that are specific to the type action"]
- #[doc = " being undertaken."]
- pub async fn do_action(
- &mut self,
- request: impl tonic::IntoRequest<super::Action>,
- ) -> Result<tonic::Response<tonic::codec::Streaming<super::Result>>, tonic::Status>
- {
- self.inner.ready().await.map_err(|e| {
- tonic::Status::new(
- tonic::Code::Unknown,
- format!("Service was not ready: {}", e.into()),
- )
- })?;
- let codec = tonic::codec::ProstCodec::default();
- let path = http::uri::PathAndQuery::from_static(
- "/arrow.flight.protocol.FlightService/DoAction",
- );
- self.inner
- .server_streaming(request.into_request(), path, codec)
- .await
- }
- #[doc = ""]
- #[doc = " A flight service exposes all of the available action types that it has"]
- #[doc = " along with descriptions. This allows different flight consumers to"]
- #[doc = " understand the capabilities of the flight service."]
- pub async fn list_actions(
- &mut self,
- request: impl tonic::IntoRequest<super::Empty>,
- ) -> Result<
- tonic::Response<tonic::codec::Streaming<super::ActionType>>,
- tonic::Status,
- > {
- self.inner.ready().await.map_err(|e| {
- tonic::Status::new(
- tonic::Code::Unknown,
- format!("Service was not ready: {}", e.into()),
- )
- })?;
- let codec = tonic::codec::ProstCodec::default();
- let path = http::uri::PathAndQuery::from_static(
- "/arrow.flight.protocol.FlightService/ListActions",
- );
- self.inner
- .server_streaming(request.into_request(), path, codec)
- .await
- }
- }
- impl<T: Clone> Clone for FlightServiceClient<T> {
- fn clone(&self) -> Self {
- Self {
- inner: self.inner.clone(),
- }
- }
- }
- impl<T> std::fmt::Debug for FlightServiceClient<T> {
- fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
- write!(f, "FlightServiceClient {{ ... }}")
- }
- }
-}
-#[doc = r" Generated server implementations."]
-pub mod flight_service_server {
- #![allow(unused_variables, dead_code, missing_docs)]
- use tonic::codegen::*;
- #[doc = "Generated trait containing gRPC methods that should be implemented for use with FlightServiceServer."]
- #[async_trait]
- pub trait FlightService: Send + Sync + 'static {
- #[doc = "Server streaming response type for the Handshake method."]
- type HandshakeStream: futures_core::Stream<Item = Result<super::HandshakeResponse, tonic::Status>>
- + Send
- + Sync
- + 'static;
- #[doc = ""]
- #[doc = " Handshake between client and server. Depending on the server, the"]
- #[doc = " handshake may be required to determine the token that should be used for"]
- #[doc = " future operations. Both request and response are streams to allow multiple"]
- #[doc = " round-trips depending on auth mechanism."]
- async fn handshake(
- &self,
- request: tonic::Request<tonic::Streaming<super::HandshakeRequest>>,
- ) -> Result<tonic::Response<Self::HandshakeStream>, tonic::Status>;
- #[doc = "Server streaming response type for the ListFlights method."]
- type ListFlightsStream: futures_core::Stream<Item = Result<super::FlightInfo, tonic::Status>>
- + Send
- + Sync
- + 'static;
- #[doc = ""]
- #[doc = " Get a list of available streams given a particular criteria. Most flight"]
- #[doc = " services will expose one or more streams that are readily available for"]
- #[doc = " retrieval. This api allows listing the streams available for"]
- #[doc = " consumption. A user can also provide a criteria. The criteria can limit"]
- #[doc = " the subset of streams that can be listed via this interface. Each flight"]
- #[doc = " service allows its own definition of how to consume criteria."]
- async fn list_flights(
- &self,
- request: tonic::Request<super::Criteria>,
- ) -> Result<tonic::Response<Self::ListFlightsStream>, tonic::Status>;
- #[doc = ""]
- #[doc = " For a given FlightDescriptor, get information about how the flight can be"]
- #[doc = " consumed. This is a useful interface if the consumer of the interface"]
- #[doc = " already can identify the specific flight to consume. This interface can"]
- #[doc = " also allow a consumer to generate a flight stream through a specified"]
- #[doc = " descriptor. For example, a flight descriptor might be something that"]
- #[doc = " includes a SQL statement or a Pickled Python operation that will be"]
- #[doc = " executed. In those cases, the descriptor will not be previously available"]
- #[doc = " within the list of available streams provided by ListFlights but will be"]
- #[doc = " available for consumption for the duration defined by the specific flight"]
- #[doc = " service."]
- async fn get_flight_info(
- &self,
- request: tonic::Request<super::FlightDescriptor>,
- ) -> Result<tonic::Response<super::FlightInfo>, tonic::Status>;
- #[doc = ""]
- #[doc = " For a given FlightDescriptor, get the Schema as described in Schema.fbs::Schema"]
- #[doc = " This is used when a consumer needs the Schema of flight stream. Similar to"]
- #[doc = " GetFlightInfo this interface may generate a new flight that was not previously"]
- #[doc = " available in ListFlights."]
- async fn get_schema(
- &self,
- request: tonic::Request<super::FlightDescriptor>,
- ) -> Result<tonic::Response<super::SchemaResult>, tonic::Status>;
- #[doc = "Server streaming response type for the DoGet method."]
- type DoGetStream: futures_core::Stream<Item = Result<super::FlightData, tonic::Status>>
- + Send
- + Sync
- + 'static;
- #[doc = ""]
- #[doc = " Retrieve a single stream associated with a particular descriptor"]
- #[doc = " associated with the referenced ticket. A Flight can be composed of one or"]
- #[doc = " more streams where each stream can be retrieved using a separate opaque"]
- #[doc = " ticket that the flight service uses for managing a collection of streams."]
- async fn do_get(
- &self,
- request: tonic::Request<super::Ticket>,
- ) -> Result<tonic::Response<Self::DoGetStream>, tonic::Status>;
- #[doc = "Server streaming response type for the DoPut method."]
- type DoPutStream: futures_core::Stream<Item = Result<super::PutResult, tonic::Status>>
- + Send
- + Sync
- + 'static;
- #[doc = ""]
- #[doc = " Push a stream to the flight service associated with a particular"]
- #[doc = " flight stream. This allows a client of a flight service to upload a stream"]
- #[doc = " of data. Depending on the particular flight service, a client consumer"]
- #[doc = " could be allowed to upload a single stream per descriptor or an unlimited"]
- #[doc = " number. In the latter, the service might implement a 'seal' action that"]
- #[doc = " can be applied to a descriptor once all streams are uploaded."]
- async fn do_put(
- &self,
- request: tonic::Request<tonic::Streaming<super::FlightData>>,
- ) -> Result<tonic::Response<Self::DoPutStream>, tonic::Status>;
- #[doc = "Server streaming response type for the DoExchange method."]
- type DoExchangeStream: futures_core::Stream<Item = Result<super::FlightData, tonic::Status>>
- + Send
- + Sync
- + 'static;
- #[doc = ""]
- #[doc = " Open a bidirectional data channel for a given descriptor. This"]
- #[doc = " allows clients to send and receive arbitrary Arrow data and"]
- #[doc = " application-specific metadata in a single logical stream. In"]
- #[doc = " contrast to DoGet/DoPut, this is more suited for clients"]
- #[doc = " offloading computation (rather than storage) to a Flight service."]
- async fn do_exchange(
- &self,
- request: tonic::Request<tonic::Streaming<super::FlightData>>,
- ) -> Result<tonic::Response<Self::DoExchangeStream>, tonic::Status>;
- #[doc = "Server streaming response type for the DoAction method."]
- type DoActionStream: futures_core::Stream<Item = Result<super::Result, tonic::Status>>
- + Send
- + Sync
- + 'static;
- #[doc = ""]
- #[doc = " Flight services can support an arbitrary number of simple actions in"]
- #[doc = " addition to the possible ListFlights, GetFlightInfo, DoGet, DoPut"]
- #[doc = " operations that are potentially available. DoAction allows a flight client"]
- #[doc = " to do a specific action against a flight service. An action includes"]
- #[doc = " opaque request and response objects that are specific to the type action"]
- #[doc = " being undertaken."]
- async fn do_action(
- &self,
- request: tonic::Request<super::Action>,
- ) -> Result<tonic::Response<Self::DoActionStream>, tonic::Status>;
- #[doc = "Server streaming response type for the ListActions method."]
- type ListActionsStream: futures_core::Stream<Item = Result<super::ActionType, tonic::Status>>
- + Send
- + Sync
- + 'static;
- #[doc = ""]
- #[doc = " A flight service exposes all of the available action types that it has"]
- #[doc = " along with descriptions. This allows different flight consumers to"]
- #[doc = " understand the capabilities of the flight service."]
- async fn list_actions(
- &self,
- request: tonic::Request<super::Empty>,
- ) -> Result<tonic::Response<Self::ListActionsStream>, tonic::Status>;
- }
- #[doc = ""]
- #[doc = " A flight service is an endpoint for retrieving or storing Arrow data. A"]
- #[doc = " flight service can expose one or more predefined endpoints that can be"]
- #[doc = " accessed using the Arrow Flight Protocol. Additionally, a flight service"]
- #[doc = " can expose a set of actions that are available."]
- #[derive(Debug)]
- pub struct FlightServiceServer<T: FlightService> {
- inner: _Inner<T>,
- }
- struct _Inner<T>(Arc<T>, Option<tonic::Interceptor>);
- impl<T: FlightService> FlightServiceServer<T> {
- pub fn new(inner: T) -> Self {
- let inner = Arc::new(inner);
- let inner = _Inner(inner, None);
- Self { inner }
- }
- pub fn with_interceptor(
- inner: T,
- interceptor: impl Into<tonic::Interceptor>,
- ) -> Self {
- let inner = Arc::new(inner);
- let inner = _Inner(inner, Some(interceptor.into()));
- Self { inner }
- }
- }
- impl<T, B> Service<http::Request<B>> for FlightServiceServer<T>
- where
- T: FlightService,
- B: HttpBody + Send + Sync + 'static,
- B::Error: Into<StdError> + Send + 'static,
- {
- type Response = http::Response<tonic::body::BoxBody>;
- type Error = Never;
- type Future = BoxFuture<Self::Response, Self::Error>;
- fn poll_ready(&mut self, _cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
- Poll::Ready(Ok(()))
- }
- fn call(&mut self, req: http::Request<B>) -> Self::Future {
- let inner = self.inner.clone();
- match req.uri().path() {
- "/arrow.flight.protocol.FlightService/Handshake" => {
- #[allow(non_camel_case_types)]
- struct HandshakeSvc<T: FlightService>(pub Arc<T>);
- impl<T: FlightService>
- tonic::server::StreamingService<super::HandshakeRequest>
- for HandshakeSvc<T>
- {
- type Response = super::HandshakeResponse;
- type ResponseStream = T::HandshakeStream;
- type Future = BoxFuture<
- tonic::Response<Self::ResponseStream>,
- tonic::Status,
- >;
- fn call(
- &mut self,
- request: tonic::Request<
- tonic::Streaming<super::HandshakeRequest>,
- >,
- ) -> Self::Future {
- let inner = self.0.clone();
- let fut = async move { (*inner).handshake(request).await };
- Box::pin(fut)
- }
- }
- let inner = self.inner.clone();
- let fut = async move {
- let interceptor = inner.1;
- let inner = inner.0;
- let method = HandshakeSvc(inner);
- let codec = tonic::codec::ProstCodec::default();
- let mut grpc = if let Some(interceptor) = interceptor {
- tonic::server::Grpc::with_interceptor(codec, interceptor)
- } else {
- tonic::server::Grpc::new(codec)
- };
- let res = grpc.streaming(method, req).await;
- Ok(res)
- };
- Box::pin(fut)
- }
- "/arrow.flight.protocol.FlightService/ListFlights" => {
- #[allow(non_camel_case_types)]
- struct ListFlightsSvc<T: FlightService>(pub Arc<T>);
- impl<T: FlightService>
- tonic::server::ServerStreamingService<super::Criteria>
- for ListFlightsSvc<T>
- {
- type Response = super::FlightInfo;
- type ResponseStream = T::ListFlightsStream;
- type Future = BoxFuture<
- tonic::Response<Self::ResponseStream>,
- tonic::Status,
- >;
- fn call(
- &mut self,
- request: tonic::Request<super::Criteria>,
- ) -> Self::Future {
- let inner = self.0.clone();
- let fut = async move { (*inner).list_flights(request).await };
- Box::pin(fut)
- }
- }
- let inner = self.inner.clone();
- let fut = async move {
- let interceptor = inner.1;
- let inner = inner.0;
- let method = ListFlightsSvc(inner);
- let codec = tonic::codec::ProstCodec::default();
- let mut grpc = if let Some(interceptor) = interceptor {
- tonic::server::Grpc::with_interceptor(codec, interceptor)
- } else {
- tonic::server::Grpc::new(codec)
- };
- let res = grpc.server_streaming(method, req).await;
- Ok(res)
- };
- Box::pin(fut)
- }
- "/arrow.flight.protocol.FlightService/GetFlightInfo" => {
- #[allow(non_camel_case_types)]
- struct GetFlightInfoSvc<T: FlightService>(pub Arc<T>);
- impl<T: FlightService>
- tonic::server::UnaryService<super::FlightDescriptor>
- for GetFlightInfoSvc<T>
- {
- type Response = super::FlightInfo;
- type Future =
- BoxFuture<tonic::Response<Self::Response>, tonic::Status>;
- fn call(
- &mut self,
- request: tonic::Request<super::FlightDescriptor>,
- ) -> Self::Future {
- let inner = self.0.clone();
- let fut =
- async move { (*inner).get_flight_info(request).await };
- Box::pin(fut)
- }
- }
- let inner = self.inner.clone();
- let fut = async move {
- let interceptor = inner.1.clone();
- let inner = inner.0;
- let method = GetFlightInfoSvc(inner);
- let codec = tonic::codec::ProstCodec::default();
- let mut grpc = if let Some(interceptor) = interceptor {
- tonic::server::Grpc::with_interceptor(codec, interceptor)
- } else {
- tonic::server::Grpc::new(codec)
- };
- let res = grpc.unary(method, req).await;
- Ok(res)
- };
- Box::pin(fut)
- }
- "/arrow.flight.protocol.FlightService/GetSchema" => {
- #[allow(non_camel_case_types)]
- struct GetSchemaSvc<T: FlightService>(pub Arc<T>);
- impl<T: FlightService>
- tonic::server::UnaryService<super::FlightDescriptor>
- for GetSchemaSvc<T>
- {
- type Response = super::SchemaResult;
- type Future =
- BoxFuture<tonic::Response<Self::Response>, tonic::Status>;
- fn call(
- &mut self,
- request: tonic::Request<super::FlightDescriptor>,
- ) -> Self::Future {
- let inner = self.0.clone();
- let fut = async move { (*inner).get_schema(request).await };
- Box::pin(fut)
- }
- }
- let inner = self.inner.clone();
- let fut = async move {
- let interceptor = inner.1.clone();
- let inner = inner.0;
- let method = GetSchemaSvc(inner);
- let codec = tonic::codec::ProstCodec::default();
- let mut grpc = if let Some(interceptor) = interceptor {
- tonic::server::Grpc::with_interceptor(codec, interceptor)
- } else {
- tonic::server::Grpc::new(codec)
- };
- let res = grpc.unary(method, req).await;
- Ok(res)
- };
- Box::pin(fut)
- }
- "/arrow.flight.protocol.FlightService/DoGet" => {
- #[allow(non_camel_case_types)]
- struct DoGetSvc<T: FlightService>(pub Arc<T>);
- impl<T: FlightService>
- tonic::server::ServerStreamingService<super::Ticket>
- for DoGetSvc<T>
- {
- type Response = super::FlightData;
- type ResponseStream = T::DoGetStream;
- type Future = BoxFuture<
- tonic::Response<Self::ResponseStream>,
- tonic::Status,
- >;
- fn call(
- &mut self,
- request: tonic::Request<super::Ticket>,
- ) -> Self::Future {
- let inner = self.0.clone();
- let fut = async move { (*inner).do_get(request).await };
- Box::pin(fut)
- }
- }
- let inner = self.inner.clone();
- let fut = async move {
- let interceptor = inner.1;
- let inner = inner.0;
- let method = DoGetSvc(inner);
- let codec = tonic::codec::ProstCodec::default();
- let mut grpc = if let Some(interceptor) = interceptor {
- tonic::server::Grpc::with_interceptor(codec, interceptor)
- } else {
- tonic::server::Grpc::new(codec)
- };
- let res = grpc.server_streaming(method, req).await;
- Ok(res)
- };
- Box::pin(fut)
- }
- "/arrow.flight.protocol.FlightService/DoPut" => {
- #[allow(non_camel_case_types)]
- struct DoPutSvc<T: FlightService>(pub Arc<T>);
- impl<T: FlightService>
- tonic::server::StreamingService<super::FlightData>
- for DoPutSvc<T>
- {
- type Response = super::PutResult;
- type ResponseStream = T::DoPutStream;
- type Future = BoxFuture<
- tonic::Response<Self::ResponseStream>,
- tonic::Status,
- >;
- fn call(
- &mut self,
- request: tonic::Request<tonic::Streaming<super::FlightData>>,
- ) -> Self::Future {
- let inner = self.0.clone();
- let fut = async move { (*inner).do_put(request).await };
- Box::pin(fut)
- }
- }
- let inner = self.inner.clone();
- let fut = async move {
- let interceptor = inner.1;
- let inner = inner.0;
- let method = DoPutSvc(inner);
- let codec = tonic::codec::ProstCodec::default();
- let mut grpc = if let Some(interceptor) = interceptor {
- tonic::server::Grpc::with_interceptor(codec, interceptor)
- } else {
- tonic::server::Grpc::new(codec)
- };
- let res = grpc.streaming(method, req).await;
- Ok(res)
- };
- Box::pin(fut)
- }
- "/arrow.flight.protocol.FlightService/DoExchange" => {
- #[allow(non_camel_case_types)]
- struct DoExchangeSvc<T: FlightService>(pub Arc<T>);
- impl<T: FlightService>
- tonic::server::StreamingService<super::FlightData>
- for DoExchangeSvc<T>
- {
- type Response = super::FlightData;
- type ResponseStream = T::DoExchangeStream;
- type Future = BoxFuture<
- tonic::Response<Self::ResponseStream>,
- tonic::Status,
- >;
- fn call(
- &mut self,
- request: tonic::Request<tonic::Streaming<super::FlightData>>,
- ) -> Self::Future {
- let inner = self.0.clone();
- let fut = async move { (*inner).do_exchange(request).await };
- Box::pin(fut)
- }
- }
- let inner = self.inner.clone();
- let fut = async move {
- let interceptor = inner.1;
- let inner = inner.0;
- let method = DoExchangeSvc(inner);
- let codec = tonic::codec::ProstCodec::default();
- let mut grpc = if let Some(interceptor) = interceptor {
- tonic::server::Grpc::with_interceptor(codec, interceptor)
- } else {
- tonic::server::Grpc::new(codec)
- };
- let res = grpc.streaming(method, req).await;
- Ok(res)
- };
- Box::pin(fut)
- }
- "/arrow.flight.protocol.FlightService/DoAction" => {
- #[allow(non_camel_case_types)]
- struct DoActionSvc<T: FlightService>(pub Arc<T>);
- impl<T: FlightService>
- tonic::server::ServerStreamingService<super::Action>
- for DoActionSvc<T>
- {
- type Response = super::Result;
- type ResponseStream = T::DoActionStream;
- type Future = BoxFuture<
- tonic::Response<Self::ResponseStream>,
- tonic::Status,
- >;
- fn call(
- &mut self,
- request: tonic::Request<super::Action>,
- ) -> Self::Future {
- let inner = self.0.clone();
- let fut = async move { (*inner).do_action(request).await };
- Box::pin(fut)
- }
- }
- let inner = self.inner.clone();
- let fut = async move {
- let interceptor = inner.1;
- let inner = inner.0;
- let method = DoActionSvc(inner);
- let codec = tonic::codec::ProstCodec::default();
- let mut grpc = if let Some(interceptor) = interceptor {
- tonic::server::Grpc::with_interceptor(codec, interceptor)
- } else {
- tonic::server::Grpc::new(codec)
- };
- let res = grpc.server_streaming(method, req).await;
- Ok(res)
- };
- Box::pin(fut)
- }
- "/arrow.flight.protocol.FlightService/ListActions" => {
- #[allow(non_camel_case_types)]
- struct ListActionsSvc<T: FlightService>(pub Arc<T>);
- impl<T: FlightService>
- tonic::server::ServerStreamingService<super::Empty>
- for ListActionsSvc<T>
- {
- type Response = super::ActionType;
- type ResponseStream = T::ListActionsStream;
- type Future = BoxFuture<
- tonic::Response<Self::ResponseStream>,
- tonic::Status,
- >;
- fn call(
- &mut self,
- request: tonic::Request<super::Empty>,
- ) -> Self::Future {
- let inner = self.0.clone();
- let fut = async move { (*inner).list_actions(request).await };
- Box::pin(fut)
- }
- }
- let inner = self.inner.clone();
- let fut = async move {
- let interceptor = inner.1;
- let inner = inner.0;
- let method = ListActionsSvc(inner);
- let codec = tonic::codec::ProstCodec::default();
- let mut grpc = if let Some(interceptor) = interceptor {
- tonic::server::Grpc::with_interceptor(codec, interceptor)
- } else {
- tonic::server::Grpc::new(codec)
- };
- let res = grpc.server_streaming(method, req).await;
- Ok(res)
- };
- Box::pin(fut)
- }
- _ => Box::pin(async move {
- Ok(http::Response::builder()
- .status(200)
- .header("grpc-status", "12")
- .header("content-type", "application/grpc")
- .body(tonic::body::BoxBody::empty())
- .unwrap())
- }),
- }
- }
- }
- impl<T: FlightService> Clone for FlightServiceServer<T> {
- fn clone(&self) -> Self {
- let inner = self.inner.clone();
- Self { inner }
- }
- }
- impl<T: FlightService> Clone for _Inner<T> {
- fn clone(&self) -> Self {
- Self(self.0.clone(), self.1.clone())
- }
- }
- impl<T: std::fmt::Debug> std::fmt::Debug for _Inner<T> {
- fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
- write!(f, "{:?}", self.0)
- }
- }
- impl<T: FlightService> tonic::transport::NamedService for FlightServiceServer<T> {
- const NAME: &'static str = "arrow.flight.protocol.FlightService";
- }
-}
diff --git a/arrow-flight/src/lib.rs b/arrow-flight/src/lib.rs
deleted file mode 100644
index 6af2e74..0000000
--- a/arrow-flight/src/lib.rs
+++ /dev/null
@@ -1,20 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-include!("arrow.flight.protocol.rs");
-
-pub mod utils;
diff --git a/arrow-flight/src/utils.rs b/arrow-flight/src/utils.rs
deleted file mode 100644
index 659668c..0000000
--- a/arrow-flight/src/utils.rs
+++ /dev/null
@@ -1,167 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-//! Utilities to assist with reading and writing Arrow data as Flight messages
-
-use std::convert::TryFrom;
-
-use crate::{FlightData, SchemaResult};
-
-use arrow::array::ArrayRef;
-use arrow::datatypes::{Schema, SchemaRef};
-use arrow::error::{ArrowError, Result};
-use arrow::ipc::{convert, reader, writer, writer::EncodedData, writer::IpcWriteOptions};
-use arrow::record_batch::RecordBatch;
-
-/// Convert a `RecordBatch` to a vector of `FlightData` representing the bytes of the dictionaries
-/// and a `FlightData` representing the bytes of the batch's values
-pub fn flight_data_from_arrow_batch(
- batch: &RecordBatch,
- options: &IpcWriteOptions,
-) -> (Vec<FlightData>, FlightData) {
- let data_gen = writer::IpcDataGenerator::default();
- let mut dictionary_tracker = writer::DictionaryTracker::new(false);
-
- let (encoded_dictionaries, encoded_batch) = data_gen
- .encoded_batch(batch, &mut dictionary_tracker, &options)
- .expect("DictionaryTracker configured above to not error on replacement");
-
- let flight_dictionaries = encoded_dictionaries.into_iter().map(Into::into).collect();
- let flight_batch = encoded_batch.into();
-
- (flight_dictionaries, flight_batch)
-}
-
-impl From<EncodedData> for FlightData {
- fn from(data: EncodedData) -> Self {
- FlightData {
- data_header: data.ipc_message,
- data_body: data.arrow_data,
- ..Default::default()
- }
- }
-}
-
-/// Convert a `Schema` to `SchemaResult` by converting to an IPC message
-pub fn flight_schema_from_arrow_schema(
- schema: &Schema,
- options: &IpcWriteOptions,
-) -> SchemaResult {
- SchemaResult {
- schema: flight_schema_as_flatbuffer(schema, options),
- }
-}
-
-/// Convert a `Schema` to `FlightData` by converting to an IPC message
-pub fn flight_data_from_arrow_schema(
- schema: &Schema,
- options: &IpcWriteOptions,
-) -> FlightData {
- let data_header = flight_schema_as_flatbuffer(schema, options);
- FlightData {
- data_header,
- ..Default::default()
- }
-}
-
-/// Convert a `Schema` to bytes in the format expected in `FlightInfo.schema`
-pub fn ipc_message_from_arrow_schema(
- arrow_schema: &Schema,
- options: &IpcWriteOptions,
-) -> Result<Vec<u8>> {
- let encoded_data = flight_schema_as_encoded_data(arrow_schema, options);
-
- let mut schema = vec![];
- arrow::ipc::writer::write_message(&mut schema, encoded_data, options)?;
- Ok(schema)
-}
-
-fn flight_schema_as_flatbuffer(
- arrow_schema: &Schema,
- options: &IpcWriteOptions,
-) -> Vec<u8> {
- let encoded_data = flight_schema_as_encoded_data(arrow_schema, options);
- encoded_data.ipc_message
-}
-
-fn flight_schema_as_encoded_data(
- arrow_schema: &Schema,
- options: &IpcWriteOptions,
-) -> EncodedData {
- let data_gen = writer::IpcDataGenerator::default();
- data_gen.schema_to_bytes(arrow_schema, options)
-}
-
-/// Try convert `FlightData` into an Arrow Schema
-///
-/// Returns an error if the `FlightData` header is not a valid IPC schema
-impl TryFrom<&FlightData> for Schema {
- type Error = ArrowError;
- fn try_from(data: &FlightData) -> Result<Self> {
- convert::schema_from_bytes(&data.data_header[..]).map_err(|err| {
- ArrowError::ParseError(format!(
- "Unable to convert flight data to Arrow schema: {}",
- err
- ))
- })
- }
-}
-
-/// Try convert `SchemaResult` into an Arrow Schema
-///
-/// Returns an error if the `FlightData` header is not a valid IPC schema
-impl TryFrom<&SchemaResult> for Schema {
- type Error = ArrowError;
- fn try_from(data: &SchemaResult) -> Result<Self> {
- convert::schema_from_bytes(&data.schema[..]).map_err(|err| {
- ArrowError::ParseError(format!(
- "Unable to convert schema result to Arrow schema: {}",
- err
- ))
- })
- }
-}
-
-/// Convert a FlightData message to a RecordBatch
-pub fn flight_data_to_arrow_batch(
- data: &FlightData,
- schema: SchemaRef,
- dictionaries_by_field: &[Option<ArrayRef>],
-) -> Result<RecordBatch> {
- // check that the data_header is a record batch message
- let message = arrow::ipc::root_as_message(&data.data_header[..]).map_err(|err| {
- ArrowError::ParseError(format!("Unable to get root as message: {:?}", err))
- })?;
-
- message
- .header_as_record_batch()
- .ok_or_else(|| {
- ArrowError::ParseError(
- "Unable to convert flight data header to a record batch".to_string(),
- )
- })
- .map(|batch| {
- reader::read_record_batch(
- &data.data_body,
- batch,
- schema,
- &dictionaries_by_field,
- )
- })?
-}
-
-// TODO: add more explicit conversion that exposes flight descriptor and metadata options
diff --git a/arrow-pyarrow-integration-testing/.cargo/config b/arrow-pyarrow-integration-testing/.cargo/config
deleted file mode 100644
index a127967..0000000
--- a/arrow-pyarrow-integration-testing/.cargo/config
+++ /dev/null
@@ -1,22 +0,0 @@
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-
-[target.x86_64-apple-darwin]
-rustflags = [
- "-C", "link-arg=-undefined",
- "-C", "link-arg=dynamic_lookup",
-]
\ No newline at end of file
diff --git a/arrow-pyarrow-integration-testing/.gitignore b/arrow-pyarrow-integration-testing/.gitignore
deleted file mode 100644
index 82adb58..0000000
--- a/arrow-pyarrow-integration-testing/.gitignore
+++ /dev/null
@@ -1,2 +0,0 @@
-__pycache__
-venv
diff --git a/arrow-pyarrow-integration-testing/Cargo.toml b/arrow-pyarrow-integration-testing/Cargo.toml
deleted file mode 100644
index 34d2435..0000000
--- a/arrow-pyarrow-integration-testing/Cargo.toml
+++ /dev/null
@@ -1,38 +0,0 @@
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-
-[package]
-name = "arrow-pyarrow-integration-testing"
-description = ""
-version = "5.0.0-SNAPSHOT"
-homepage = "https://github.com/apache/arrow-rs"
-repository = "https://github.com/apache/arrow-rs"
-authors = ["Apache Arrow <dev@arrow.apache.org>"]
-license = "Apache-2.0"
-keywords = [ "arrow" ]
-edition = "2018"
-
-[lib]
-name = "arrow_pyarrow_integration_testing"
-crate-type = ["cdylib"]
-
-[dependencies]
-arrow = { path = "../arrow", version = "5.0.0-SNAPSHOT" }
-pyo3 = { version = "0.12.1", features = ["extension-module"] }
-
-[package.metadata.maturin]
-requires-dist = ["pyarrow>=1"]
diff --git a/arrow-pyarrow-integration-testing/README.md b/arrow-pyarrow-integration-testing/README.md
deleted file mode 100644
index 7e78aa9..0000000
--- a/arrow-pyarrow-integration-testing/README.md
+++ /dev/null
@@ -1,57 +0,0 @@
-<!---
- Licensed to the Apache Software Foundation (ASF) under one
- or more contributor license agreements. See the NOTICE file
- distributed with this work for additional information
- regarding copyright ownership. The ASF licenses this file
- to you under the Apache License, Version 2.0 (the
- "License"); you may not use this file except in compliance
- with the License. You may obtain a copy of the License at
-
- http://www.apache.org/licenses/LICENSE-2.0
-
- Unless required by applicable law or agreed to in writing,
- software distributed under the License is distributed on an
- "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
- KIND, either express or implied. See the License for the
- specific language governing permissions and limitations
- under the License.
--->
-
-# Arrow c integration
-
-This is a Rust crate that tests compatibility between Rust's Arrow implementation and PyArrow.
-
-Note that this crate uses two languages and an external ABI:
-* `Rust`
-* `Python`
-* C ABI privately exposed by `Pyarrow`.
-
-## Basic idea
-
-Pyarrow exposes a C ABI to convert arrow arrays from and to its C implementation, see [here](https://arrow.apache.org/docs/format/CDataInterface.html).
-
-This package uses the equivalent struct in Rust (`arrow::array::ArrowArray`), and verifies that
-we can use pyarrow's interface to move pointers from and to Rust.
-
-## Relevant literature
-
-* [Arrow's CDataInterface](https://arrow.apache.org/docs/format/CDataInterface.html)
-* [Rust's FFI](https://doc.rust-lang.org/nomicon/ffi.html)
-* [Pyarrow private binds](https://github.com/apache/arrow/blob/ae1d24efcc3f1ac2a876d8d9f544a34eb04ae874/python/pyarrow/array.pxi#L1226)
-* [PyO3](https://docs.rs/pyo3/0.12.1/pyo3/index.html)
-
-## How to develop
-
-```bash
-# prepare development environment (used to build wheel / install in development)
-python -m venv venv
-venv/bin/pip install maturin==0.8.2 toml==0.10.1 pyarrow==1.0.0
-```
-
-Whenever rust code changes (your changes or via git pull):
-
-```bash
-source venv/bin/activate
-maturin develop
-python -m unittest discover tests
-```
diff --git a/arrow-pyarrow-integration-testing/pyproject.toml b/arrow-pyarrow-integration-testing/pyproject.toml
deleted file mode 100644
index 2748069..0000000
--- a/arrow-pyarrow-integration-testing/pyproject.toml
+++ /dev/null
@@ -1,20 +0,0 @@
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-
-[build-system]
-requires = ["maturin"]
-build-backend = "maturin"
diff --git a/arrow-pyarrow-integration-testing/src/lib.rs b/arrow-pyarrow-integration-testing/src/lib.rs
deleted file mode 100644
index 5b5462d..0000000
--- a/arrow-pyarrow-integration-testing/src/lib.rs
+++ /dev/null
@@ -1,188 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-//! This library demonstrates a minimal usage of Rust's C data interface to pass
-//! arrays from and to Python.
-
-use std::error;
-use std::fmt;
-use std::sync::Arc;
-
-use pyo3::exceptions::PyOSError;
-use pyo3::wrap_pyfunction;
-use pyo3::{libc::uintptr_t, prelude::*};
-
-use arrow::array::{make_array_from_raw, ArrayRef, Int64Array};
-use arrow::compute::kernels;
-use arrow::error::ArrowError;
-use arrow::ffi;
-
-/// an error that bridges ArrowError with a Python error
-#[derive(Debug)]
-enum PyO3ArrowError {
- ArrowError(ArrowError),
-}
-
-impl fmt::Display for PyO3ArrowError {
- fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
- match *self {
- PyO3ArrowError::ArrowError(ref e) => e.fmt(f),
- }
- }
-}
-
-impl error::Error for PyO3ArrowError {
- fn source(&self) -> Option<&(dyn error::Error + 'static)> {
- match *self {
- // The cause is the underlying implementation error type. Is implicitly
- // cast to the trait object `&error::Error`. This works because the
- // underlying type already implements the `Error` trait.
- PyO3ArrowError::ArrowError(ref e) => Some(e),
- }
- }
-}
-
-impl From<ArrowError> for PyO3ArrowError {
- fn from(err: ArrowError) -> PyO3ArrowError {
- PyO3ArrowError::ArrowError(err)
- }
-}
-
-impl From<PyO3ArrowError> for PyErr {
- fn from(err: PyO3ArrowError) -> PyErr {
- PyOSError::new_err(err.to_string())
- }
-}
-
-fn to_rust(ob: PyObject, py: Python) -> PyResult<ArrayRef> {
- // prepare a pointer to receive the Array struct
- let (array_pointer, schema_pointer) =
- ffi::ArrowArray::into_raw(unsafe { ffi::ArrowArray::empty() });
-
- // make the conversion through PyArrow's private API
- // this changes the pointer's memory and is thus unsafe. In particular, `_export_to_c` can go out of bounds
- ob.call_method1(
- py,
- "_export_to_c",
- (array_pointer as uintptr_t, schema_pointer as uintptr_t),
- )?;
-
- let array = unsafe { make_array_from_raw(array_pointer, schema_pointer) }
- .map_err(|e| PyO3ArrowError::from(e))?;
- Ok(array)
-}
-
-fn to_py(array: ArrayRef, py: Python) -> PyResult<PyObject> {
- let (array_pointer, schema_pointer) =
- array.to_raw().map_err(|e| PyO3ArrowError::from(e))?;
-
- let pa = py.import("pyarrow")?;
-
- let array = pa.getattr("Array")?.call_method1(
- "_import_from_c",
- (array_pointer as uintptr_t, schema_pointer as uintptr_t),
- )?;
- Ok(array.to_object(py))
-}
-
-/// Returns `array + array` of an int64 array.
-#[pyfunction]
-fn double(array: PyObject, py: Python) -> PyResult<PyObject> {
- // import
- let array = to_rust(array, py)?;
-
- // perform some operation
- let array =
- array
- .as_any()
- .downcast_ref::<Int64Array>()
- .ok_or(PyO3ArrowError::ArrowError(ArrowError::ParseError(
- "Expects an int64".to_string(),
- )))?;
- let array =
- kernels::arithmetic::add(&array, &array).map_err(|e| PyO3ArrowError::from(e))?;
- let array = Arc::new(array);
-
- // export
- to_py(array, py)
-}
-
-/// calls a lambda function that receives and returns an array
-/// whose result must be the array multiplied by two
-#[pyfunction]
-fn double_py(lambda: PyObject, py: Python) -> PyResult<bool> {
- // create
- let array = Arc::new(Int64Array::from(vec![Some(1), None, Some(3)]));
- let expected = Arc::new(Int64Array::from(vec![Some(2), None, Some(6)])) as ArrayRef;
-
- // to py
- let array = to_py(array, py)?;
-
- let array = lambda.call1(py, (array,))?;
-
- let array = to_rust(array, py)?;
-
- Ok(array == expected)
-}
-
-/// Returns the substring
-#[pyfunction]
-fn substring(array: PyObject, start: i64, py: Python) -> PyResult<PyObject> {
- // import
- let array = to_rust(array, py)?;
-
- // substring
- let array = kernels::substring::substring(array.as_ref(), start, &None)
- .map_err(|e| PyO3ArrowError::from(e))?;
-
- // export
- to_py(array, py)
-}
-
-/// Returns the concatenate
-#[pyfunction]
-fn concatenate(array: PyObject, py: Python) -> PyResult<PyObject> {
- // import
- let array = to_rust(array, py)?;
-
- // concat
- let array = kernels::concat::concat(&[array.as_ref(), array.as_ref()])
- .map_err(|e| PyO3ArrowError::from(e))?;
-
- // export
- to_py(array, py)
-}
-
-/// Converts to rust and back to python
-#[pyfunction]
-fn round_trip(array: PyObject, py: Python) -> PyResult<PyObject> {
- // import
- let array = to_rust(array, py)?;
-
- // export
- to_py(array, py)
-}
-
-#[pymodule]
-fn arrow_pyarrow_integration_testing(_py: Python, m: &PyModule) -> PyResult<()> {
- m.add_wrapped(wrap_pyfunction!(double))?;
- m.add_wrapped(wrap_pyfunction!(double_py))?;
- m.add_wrapped(wrap_pyfunction!(substring))?;
- m.add_wrapped(wrap_pyfunction!(concatenate))?;
- m.add_wrapped(wrap_pyfunction!(round_trip))?;
- Ok(())
-}
diff --git a/arrow-pyarrow-integration-testing/tests/test_sql.py b/arrow-pyarrow-integration-testing/tests/test_sql.py
deleted file mode 100644
index c0de382..0000000
--- a/arrow-pyarrow-integration-testing/tests/test_sql.py
+++ /dev/null
@@ -1,99 +0,0 @@
-# -*- coding: utf-8 -*-
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-
-import unittest
-
-import pyarrow
-import arrow_pyarrow_integration_testing
-
-
-class TestCase(unittest.TestCase):
- def test_primitive_python(self):
- """
- Python -> Rust -> Python
- """
- old_allocated = pyarrow.total_allocated_bytes()
- a = pyarrow.array([1, 2, 3])
- b = arrow_pyarrow_integration_testing.double(a)
- self.assertEqual(b, pyarrow.array([2, 4, 6]))
- del a
- del b
- # No leak of C++ memory
- self.assertEqual(old_allocated, pyarrow.total_allocated_bytes())
-
- def test_primitive_rust(self):
- """
- Rust -> Python -> Rust
- """
- old_allocated = pyarrow.total_allocated_bytes()
-
- def double(array):
- array = array.to_pylist()
- return pyarrow.array([x * 2 if x is not None else None for x in array])
-
- is_correct = arrow_pyarrow_integration_testing.double_py(double)
- self.assertTrue(is_correct)
- # No leak of C++ memory
- self.assertEqual(old_allocated, pyarrow.total_allocated_bytes())
-
- def test_string_python(self):
- """
- Python -> Rust -> Python
- """
- old_allocated = pyarrow.total_allocated_bytes()
- a = pyarrow.array(["a", None, "ccc"])
- b = arrow_pyarrow_integration_testing.substring(a, 1)
- self.assertEqual(b, pyarrow.array(["", None, "cc"]))
- del a
- del b
- # No leak of C++ memory
- self.assertEqual(old_allocated, pyarrow.total_allocated_bytes())
-
- def test_time32_python(self):
- """
- Python -> Rust -> Python
- """
- old_allocated = pyarrow.total_allocated_bytes()
- a = pyarrow.array([None, 1, 2], pyarrow.time32('s'))
- b = arrow_pyarrow_integration_testing.concatenate(a)
- expected = pyarrow.array([None, 1, 2] + [None, 1, 2], pyarrow.time32('s'))
- self.assertEqual(b, expected)
- del a
- del b
- del expected
- # No leak of C++ memory
- self.assertEqual(old_allocated, pyarrow.total_allocated_bytes())
-
- def test_list_array(self):
- """
- Python -> Rust -> Python
- """
- old_allocated = pyarrow.total_allocated_bytes()
- a = pyarrow.array([[], None, [1, 2], [4, 5, 6]], pyarrow.list_(pyarrow.int64()))
- b = arrow_pyarrow_integration_testing.round_trip(a)
-
- b.validate(full=True)
- assert a.to_pylist() == b.to_pylist()
- assert a.type == b.type
- del a
- del b
- # No leak of C++ memory
- self.assertEqual(old_allocated, pyarrow.total_allocated_bytes())
-
-
-
diff --git a/arrow/Cargo.toml b/arrow/Cargo.toml
deleted file mode 100644
index c3ba8b1..0000000
--- a/arrow/Cargo.toml
+++ /dev/null
@@ -1,157 +0,0 @@
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-
-[package]
-name = "arrow"
-version = "5.0.0-SNAPSHOT"
-description = "Rust implementation of Apache Arrow"
-homepage = "https://github.com/apache/arrow-rs"
-repository = "https://github.com/apache/arrow-rs"
-authors = ["Apache Arrow <dev@arrow.apache.org>"]
-license = "Apache-2.0"
-keywords = [ "arrow" ]
-include = [
- "benches/*.rs",
- "src/**/*.rs",
- "Cargo.toml",
-]
-edition = "2018"
-
-[lib]
-name = "arrow"
-path = "src/lib.rs"
-
-[dependencies]
-serde = { version = "1.0", features = ["rc"] }
-serde_derive = "1.0"
-serde_json = { version = "1.0", features = ["preserve_order"] }
-indexmap = "1.6"
-rand = "0.7"
-num = "0.4"
-csv_crate = { version = "1.1", optional = true, package="csv" }
-regex = "1.3"
-lazy_static = "1.4"
-packed_simd = { version = "0.3.4", optional = true, package = "packed_simd_2" }
-chrono = "0.4"
-flatbuffers = { version = "=0.8.4", optional = true }
-hex = "0.4"
-prettytable-rs = { version = "0.8.0", optional = true }
-lexical-core = "^0.7"
-multiversion = "0.6.1"
-
-[features]
-default = ["csv", "ipc"]
-avx512 = []
-csv = ["csv_crate"]
-ipc = ["flatbuffers"]
-simd = ["packed_simd"]
-prettyprint = ["prettytable-rs"]
-# this is only intended to be used in single-threaded programs: it verifies that
-# all allocated memory is being released (no memory leaks).
-# See README for details
-memory-check = []
-
-[dev-dependencies]
-criterion = "0.3"
-flate2 = "1"
-tempfile = "3"
-
-[build-dependencies]
-
-[[bench]]
-name = "aggregate_kernels"
-harness = false
-
-[[bench]]
-name = "array_from_vec"
-harness = false
-
-[[bench]]
-name = "builder"
-harness = false
-
-[[bench]]
-name = "buffer_bit_ops"
-harness = false
-
-[[bench]]
-name = "boolean_kernels"
-harness = false
-
-[[bench]]
-name = "arithmetic_kernels"
-harness = false
-
-[[bench]]
-name = "cast_kernels"
-harness = false
-
-[[bench]]
-name = "comparison_kernels"
-harness = false
-
-[[bench]]
-name = "filter_kernels"
-harness = false
-
-[[bench]]
-name = "take_kernels"
-harness = false
-
-[[bench]]
-name = "length_kernel"
-harness = false
-
-[[bench]]
-name = "bit_length_kernel"
-harness = false
-
-[[bench]]
-name = "sort_kernel"
-harness = false
-
-[[bench]]
-name = "partition_kernels"
-harness = false
-
-[[bench]]
-name = "csv_writer"
-harness = false
-
-[[bench]]
-name = "json_reader"
-harness = false
-
-[[bench]]
-name = "equal"
-harness = false
-
-[[bench]]
-name = "array_slice"
-harness = false
-
-[[bench]]
-name = "concatenate_kernel"
-harness = false
-
-[[bench]]
-name = "mutable_array"
-harness = false
-
-[[bench]]
-name = "buffer_create"
-harness = false
diff --git a/arrow/README.md b/arrow/README.md
deleted file mode 100644
index dfd5926..0000000
--- a/arrow/README.md
+++ /dev/null
@@ -1,181 +0,0 @@
-<!---
- Licensed to the Apache Software Foundation (ASF) under one
- or more contributor license agreements. See the NOTICE file
- distributed with this work for additional information
- regarding copyright ownership. The ASF licenses this file
- to you under the Apache License, Version 2.0 (the
- "License"); you may not use this file except in compliance
- with the License. You may obtain a copy of the License at
-
- http://www.apache.org/licenses/LICENSE-2.0
-
- Unless required by applicable law or agreed to in writing,
- software distributed under the License is distributed on an
- "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
- KIND, either express or implied. See the License for the
- specific language governing permissions and limitations
- under the License.
--->
-
-# Native Rust implementation of Apache Arrow
-
-[](https://crates.io/crates/arrow)
-
-This crate contains a native Rust implementation of the [Arrow columnar format](https://arrow.apache.org/docs/format/Columnar.html).
-
-## Developer's guide
-
-Common information for all Rust libraries in this project, including
-testing, code formatting, and lints, can be found in the main Arrow
-Rust [README.md](../README.md).
-
-Please refer to [lib.rs](src/lib.rs) for an introduction to this
-specific crate and its current functionality.
-
-### How to check memory allocations
-
-This crate heavily uses `unsafe` due to how memory is allocated in cache lines.
-We have a small tool to verify that this crate does not leak memory (beyond what the compiler already does)
-
-Run it with
-
-```bash
-cargo test --features memory-check --lib -- --test-threads 1
-```
-
-This runs all unit-tests on a single thread and counts all allocations and de-allocations.
-
-## Examples
-
-The examples folder shows how to construct some different types of Arrow
-arrays, including dynamic arrays created at runtime.
-
-Examples can be run using the `cargo run --example` command. For example:
-
-```bash
-cargo run --example builders
-cargo run --example dynamic_types
-cargo run --example read_csv
-```
-
-## IPC
-
-The expected flatc version is 1.12.0+, built from [flatbuffers](https://github.com/google/flatbuffers)
-master at fixed commit ID, by regen.sh.
-
-The IPC flatbuffer code was generated by running this command from the root of the project:
-
-```bash
-./regen.sh
-```
-
-The above script will run the `flatc` compiler and perform some adjustments to the source code:
-
-- Replace `type__` with `type_`
-- Remove `org::apache::arrow::flatbuffers` namespace
-- Add includes to each generated file
-
-## Features
-
-Arrow uses the following features:
-
-- `simd` - Arrow uses the [packed_simd](https://crates.io/crates/packed_simd) crate to optimize many of the
- implementations in the [compute](https://github.com/apache/arrow/tree/master/rust/arrow/src/compute)
- module using SIMD intrinsics. These optimizations are turned _off_ by default.
- If the `simd` feature is enabled, an unstable version of Rust is required (we test with `nightly-2021-03-24`)
-- `flight` which contains useful functions to convert between the Flight wire format and Arrow data
-- `prettyprint` which is a utility for printing record batches
-
-Other than `simd` all the other features are enabled by default. Disabling `prettyprint` might be necessary in order to
-compile Arrow to the `wasm32-unknown-unknown` WASM target.
-
-## Guidelines in usage of `unsafe`
-
-[`unsafe`](https://doc.rust-lang.org/book/ch19-01-unsafe-rust.html) has a high maintenance cost because debugging and testing it is difficult, time consuming, often requires external tools (e.g. `valgrind`), and requires a higher-than-usual attention to details. Undefined behavior is particularly difficult to identify and test, and usage of `unsafe` is the [primary cause of undefined behavior](https://doc.rust-lang.org/reference/behavior-considered-undefined.html) in a program written in Rust.
-For two real world examples of where `unsafe` has consumed time in the past in this project see [#8545](https://github.com/apache/arrow/pull/8645) and [8829](https://github.com/apache/arrow/pull/8829)
-This crate only accepts the usage of `unsafe` code upon careful consideration, and strives to avoid it to the largest possible extent.
-
-### When can `unsafe` be used?
-
-Generally, `unsafe` should only be used when a `safe` counterpart is not available and there is no `safe` way to achieve additional performance in that area. The following is a summary of the current components of the crate that require `unsafe`:
-
-- alloc, dealloc and realloc of buffers along cache lines
-- Interpreting bytes as certain rust types, for access, representation and compute
-- Foreign interfaces (C data interface)
-- Inter-process communication (IPC)
-- SIMD
-- Performance (e.g. omit bounds checks, use of pointers to avoid bound checks)
-
-#### cache-line aligned memory management
-
-The arrow format recommends storing buffers aligned with cache lines, and this crate adopts this behavior.
-However, Rust's global allocator does not allocate memory aligned with cache-lines. As such, many of the low-level operations related to memory management require `unsafe`.
-
-#### Interpreting bytes
-
-The arrow format is specified in bytes (`u8`), which can be logically represented as certain types
-depending on the `DataType`.
-For many operations, such as access, representation, numerical computation and string manipulation,
-it is often necessary to interpret bytes as other physical types (e.g. `i32`).
-
-Usage of `unsafe` for the purpose of interpreting bytes in their corresponding type (according to the arrow specification) is allowed. Specifically, the pointer to the byte slice must be aligned to the type that it intends to represent and the length of the slice is a multiple of the size of the target type of the transmutation.
-
-#### FFI
-
-The arrow format declares an ABI for zero-copy from and to libraries that implement the specification
-(foreign interfaces). In Rust, receiving and sending pointers via FFI requires usage of `unsafe` due to
-the impossibility of the compiler to derive the invariants (such as lifetime, null pointers, and pointer alignment) from the source code alone as they are part of the FFI contract.
-
-#### IPC
-
-The arrow format declares a IPC protocol, which this crate supports. IPC is equivalent to a FFI in that the rust compiler can't reason about the contract's invariants.
-
-#### SIMD
-
-The API provided by the `packed_simd` library is currently `unsafe`. However, SIMD offers a significant performance improvement over non-SIMD operations.
-
-#### Performance
-
-Some operations are significantly faster when `unsafe` is used.
-
-A common usage of `unsafe` is to offer an API to access the `i`th element of an array (e.g. `UInt32Array`).
-This requires accessing the values buffer e.g. `array.buffers()[0]`, picking the slice
-`[i * size_of<i32>(), (i + 1) * size_of<i32>()]`, and then transmuting it to `i32`. In safe Rust,
-this operation requires boundary checks that are detrimental to performance.
-
-Usage of `unsafe` for performance reasons is justified only when all other alternatives have been exhausted and the performance benefits are sufficiently large (e.g. >~10%).
-
-### Considerations when introducing `unsafe`
-
-Usage of `unsafe` in this crate _must_:
-
-- not expose a public API as `safe` when there are necessary invariants for that API to be defined behavior.
-- have code documentation for why `safe` is not used / possible
-- have code documentation about which invariant the user needs to enforce to ensure [soundness](https://rust-lang.github.io/unsafe-code-guidelines/glossary.html#soundness-of-code--of-a-library), or which
-- invariant is being preserved.
-- if applicable, use `debug_assert`s to relevant invariants (e.g. bound checks)
-
-Example of code documentation:
-
-```rust
-// JUSTIFICATION
-// Benefit
-// Describe the benefit of using unsafe. E.g.
-// "30% performance degradation if the safe counterpart is used, see bench X."
-// Soundness
-// Describe why the code remains sound (according to the definition of rust's unsafe code guidelines). E.g.
-// "We bounded check these values at initialization and the array is immutable."
-let ... = unsafe { ... };
-```
-
-When adding this documentation to existing code that is not sound and cannot trivially be fixed, we should file
-specific JIRA issues and reference them in these code comments. For example:
-
-```rust
-// Soundness
-// This is not sound because .... see https://issues.apache.org/jira/browse/ARROW-nnnnn
-```
-
-# Releases and publishing to crates.io
-
-Please see the [release](../dev/release/README.md) for details on how to create arrow releases
diff --git a/arrow/benches/aggregate_kernels.rs b/arrow/benches/aggregate_kernels.rs
deleted file mode 100644
index 1724b73..0000000
--- a/arrow/benches/aggregate_kernels.rs
+++ /dev/null
@@ -1,67 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-#[macro_use]
-extern crate criterion;
-use criterion::Criterion;
-
-extern crate arrow;
-
-use arrow::compute::kernels::aggregate::*;
-use arrow::util::bench_util::*;
-use arrow::{array::*, datatypes::Float32Type};
-
-fn bench_sum(arr_a: &Float32Array) {
- criterion::black_box(sum(&arr_a).unwrap());
-}
-
-fn bench_min(arr_a: &Float32Array) {
- criterion::black_box(min(&arr_a).unwrap());
-}
-
-fn bench_max(arr_a: &Float32Array) {
- criterion::black_box(max(&arr_a).unwrap());
-}
-
-fn bench_min_string(arr_a: &StringArray) {
- criterion::black_box(min_string(&arr_a).unwrap());
-}
-
-fn add_benchmark(c: &mut Criterion) {
- let arr_a = create_primitive_array::<Float32Type>(512, 0.0);
-
- c.bench_function("sum 512", |b| b.iter(|| bench_sum(&arr_a)));
- c.bench_function("min 512", |b| b.iter(|| bench_min(&arr_a)));
- c.bench_function("max 512", |b| b.iter(|| bench_max(&arr_a)));
-
- let arr_a = create_primitive_array::<Float32Type>(512, 0.5);
-
- c.bench_function("sum nulls 512", |b| b.iter(|| bench_sum(&arr_a)));
- c.bench_function("min nulls 512", |b| b.iter(|| bench_min(&arr_a)));
- c.bench_function("max nulls 512", |b| b.iter(|| bench_max(&arr_a)));
-
- let arr_b = create_string_array::<i32>(512, 0.0);
- c.bench_function("min string 512", |b| b.iter(|| bench_min_string(&arr_b)));
-
- let arr_b = create_string_array::<i32>(512, 0.5);
- c.bench_function("min nulls string 512", |b| {
- b.iter(|| bench_min_string(&arr_b))
- });
-}
-
-criterion_group!(benches, add_benchmark);
-criterion_main!(benches);
diff --git a/arrow/benches/arithmetic_kernels.rs b/arrow/benches/arithmetic_kernels.rs
deleted file mode 100644
index 721157e..0000000
--- a/arrow/benches/arithmetic_kernels.rs
+++ /dev/null
@@ -1,103 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-#[macro_use]
-extern crate criterion;
-use criterion::Criterion;
-use rand::Rng;
-
-use std::sync::Arc;
-
-extern crate arrow;
-
-use arrow::compute::kernels::limit::*;
-use arrow::util::bench_util::*;
-use arrow::{array::*, datatypes::Float32Type};
-use arrow::{compute::kernels::arithmetic::*, util::test_util::seedable_rng};
-
-fn create_array(size: usize, with_nulls: bool) -> ArrayRef {
- let null_density = if with_nulls { 0.5 } else { 0.0 };
- let array = create_primitive_array::<Float32Type>(size, null_density);
- Arc::new(array)
-}
-
-fn bench_add(arr_a: &ArrayRef, arr_b: &ArrayRef) {
- let arr_a = arr_a.as_any().downcast_ref::<Float32Array>().unwrap();
- let arr_b = arr_b.as_any().downcast_ref::<Float32Array>().unwrap();
- criterion::black_box(add(arr_a, arr_b).unwrap());
-}
-
-fn bench_subtract(arr_a: &ArrayRef, arr_b: &ArrayRef) {
- let arr_a = arr_a.as_any().downcast_ref::<Float32Array>().unwrap();
- let arr_b = arr_b.as_any().downcast_ref::<Float32Array>().unwrap();
- criterion::black_box(subtract(&arr_a, &arr_b).unwrap());
-}
-
-fn bench_multiply(arr_a: &ArrayRef, arr_b: &ArrayRef) {
- let arr_a = arr_a.as_any().downcast_ref::<Float32Array>().unwrap();
- let arr_b = arr_b.as_any().downcast_ref::<Float32Array>().unwrap();
- criterion::black_box(multiply(&arr_a, &arr_b).unwrap());
-}
-
-fn bench_divide(arr_a: &ArrayRef, arr_b: &ArrayRef) {
- let arr_a = arr_a.as_any().downcast_ref::<Float32Array>().unwrap();
- let arr_b = arr_b.as_any().downcast_ref::<Float32Array>().unwrap();
- criterion::black_box(divide(&arr_a, &arr_b).unwrap());
-}
-
-fn bench_divide_scalar(array: &ArrayRef, divisor: f32) {
- let array = array.as_any().downcast_ref::<Float32Array>().unwrap();
- criterion::black_box(divide_scalar(&array, divisor).unwrap());
-}
-
-fn bench_limit(arr_a: &ArrayRef, max: usize) {
- criterion::black_box(limit(arr_a, max));
-}
-
-fn add_benchmark(c: &mut Criterion) {
- let arr_a = create_array(512, false);
- let arr_b = create_array(512, false);
- let scalar = seedable_rng().gen();
-
- c.bench_function("add 512", |b| b.iter(|| bench_add(&arr_a, &arr_b)));
- c.bench_function("subtract 512", |b| {
- b.iter(|| bench_subtract(&arr_a, &arr_b))
- });
- c.bench_function("multiply 512", |b| {
- b.iter(|| bench_multiply(&arr_a, &arr_b))
- });
- c.bench_function("divide 512", |b| b.iter(|| bench_divide(&arr_a, &arr_b)));
- c.bench_function("divide_scalar 512", |b| {
- b.iter(|| bench_divide_scalar(&arr_a, scalar))
- });
- c.bench_function("limit 512, 512", |b| b.iter(|| bench_limit(&arr_a, 512)));
-
- let arr_a_nulls = create_array(512, false);
- let arr_b_nulls = create_array(512, false);
- c.bench_function("add_nulls_512", |b| {
- b.iter(|| bench_add(&arr_a_nulls, &arr_b_nulls))
- });
- c.bench_function("divide_nulls_512", |b| {
- b.iter(|| bench_divide(&arr_a_nulls, &arr_b_nulls))
- });
- c.bench_function("divide_scalar_nulls_512", |b| {
- b.iter(|| bench_divide_scalar(&arr_a_nulls, scalar))
- });
-}
-
-criterion_group!(benches, add_benchmark);
-criterion_main!(benches);
diff --git a/arrow/benches/array_from_vec.rs b/arrow/benches/array_from_vec.rs
deleted file mode 100644
index 7740c6b..0000000
--- a/arrow/benches/array_from_vec.rs
+++ /dev/null
@@ -1,120 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-#[macro_use]
-extern crate criterion;
-use criterion::Criterion;
-
-extern crate arrow;
-
-use arrow::array::*;
-use arrow::buffer::Buffer;
-use arrow::datatypes::*;
-use std::{convert::TryFrom, sync::Arc};
-
-fn array_from_vec(n: usize) {
- let mut v: Vec<u8> = Vec::with_capacity(n);
- for i in 0..n {
- v.push((i & 0xffff) as u8);
- }
- let arr_data = ArrayDataBuilder::new(DataType::Int32)
- .add_buffer(Buffer::from(v))
- .build();
- criterion::black_box(Int32Array::from(arr_data));
-}
-
-fn array_string_from_vec(n: usize) {
- let mut v: Vec<Option<&str>> = Vec::with_capacity(n);
- for i in 0..n {
- if i % 2 == 0 {
- v.push(Some("hello world"));
- } else {
- v.push(None);
- }
- }
- criterion::black_box(StringArray::from(v));
-}
-
-fn struct_array_values(
- n: usize,
-) -> (
- &'static str,
- Vec<Option<&'static str>>,
- &'static str,
- Vec<Option<i32>>,
-) {
- let mut strings: Vec<Option<&str>> = Vec::with_capacity(n);
- let mut ints: Vec<Option<i32>> = Vec::with_capacity(n);
- for _ in 0..n / 4 {
- strings.extend_from_slice(&[Some("joe"), None, None, Some("mark")]);
- ints.extend_from_slice(&[Some(1), Some(2), None, Some(4)]);
- }
- ("f1", strings, "f2", ints)
-}
-
-fn struct_array_from_vec(
- field1: &str,
- strings: &[Option<&str>],
- field2: &str,
- ints: &[Option<i32>],
-) {
- let strings: ArrayRef = Arc::new(StringArray::from(strings.to_owned()));
- let ints: ArrayRef = Arc::new(Int32Array::from(ints.to_owned()));
-
- criterion::black_box(
- StructArray::try_from(vec![(field1, strings), (field2, ints)]).unwrap(),
- );
-}
-
-fn criterion_benchmark(c: &mut Criterion) {
- c.bench_function("array_from_vec 128", |b| b.iter(|| array_from_vec(128)));
- c.bench_function("array_from_vec 256", |b| b.iter(|| array_from_vec(256)));
- c.bench_function("array_from_vec 512", |b| b.iter(|| array_from_vec(512)));
-
- c.bench_function("array_string_from_vec 128", |b| {
- b.iter(|| array_string_from_vec(128))
- });
- c.bench_function("array_string_from_vec 256", |b| {
- b.iter(|| array_string_from_vec(256))
- });
- c.bench_function("array_string_from_vec 512", |b| {
- b.iter(|| array_string_from_vec(512))
- });
-
- let (field1, strings, field2, ints) = struct_array_values(128);
- c.bench_function("struct_array_from_vec 128", |b| {
- b.iter(|| struct_array_from_vec(&field1, &strings, &field2, &ints))
- });
-
- let (field1, strings, field2, ints) = struct_array_values(256);
- c.bench_function("struct_array_from_vec 256", |b| {
- b.iter(|| struct_array_from_vec(&field1, &strings, &field2, &ints))
- });
-
- let (field1, strings, field2, ints) = struct_array_values(512);
- c.bench_function("struct_array_from_vec 512", |b| {
- b.iter(|| struct_array_from_vec(&field1, &strings, &field2, &ints))
- });
-
- let (field1, strings, field2, ints) = struct_array_values(1024);
- c.bench_function("struct_array_from_vec 1024", |b| {
- b.iter(|| struct_array_from_vec(&field1, &strings, &field2, &ints))
- });
-}
-
-criterion_group!(benches, criterion_benchmark);
-criterion_main!(benches);
diff --git a/arrow/benches/array_slice.rs b/arrow/benches/array_slice.rs
deleted file mode 100644
index a535c80..0000000
--- a/arrow/benches/array_slice.rs
+++ /dev/null
@@ -1,52 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-#[macro_use]
-extern crate criterion;
-use criterion::Criterion;
-
-extern crate arrow;
-
-use arrow::array::*;
-use std::sync::Arc;
-
-fn create_array_slice(array: &ArrayRef, length: usize) -> ArrayRef {
- array.slice(0, length)
-}
-
-fn create_array_with_nulls(size: usize) -> ArrayRef {
- let array: Float64Array = (0..size)
- .map(|i| if i % 2 == 0 { Some(1.0) } else { None })
- .collect();
- Arc::new(array)
-}
-
-fn array_slice_benchmark(c: &mut Criterion) {
- let array = create_array_with_nulls(4096);
- c.bench_function("array_slice 128", |b| {
- b.iter(|| create_array_slice(&array, 128))
- });
- c.bench_function("array_slice 512", |b| {
- b.iter(|| create_array_slice(&array, 512))
- });
- c.bench_function("array_slice 2048", |b| {
- b.iter(|| create_array_slice(&array, 2048))
- });
-}
-
-criterion_group!(benches, array_slice_benchmark);
-criterion_main!(benches);
diff --git a/arrow/benches/bit_length_kernel.rs b/arrow/benches/bit_length_kernel.rs
deleted file mode 100644
index 51d3134..0000000
--- a/arrow/benches/bit_length_kernel.rs
+++ /dev/null
@@ -1,46 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-#[macro_use]
-extern crate criterion;
-use criterion::Criterion;
-
-extern crate arrow;
-
-use arrow::{array::*, compute::kernels::length::bit_length};
-
-fn bench_bit_length(array: &StringArray) {
- criterion::black_box(bit_length(array).unwrap());
-}
-
-fn add_benchmark(c: &mut Criterion) {
- fn double_vec<T: Clone>(v: Vec<T>) -> Vec<T> {
- [&v[..], &v[..]].concat()
- }
-
- // double ["hello", " ", "world", "!"] 10 times
- let mut values = vec!["one", "on", "o", ""];
- for _ in 0..10 {
- values = double_vec(values);
- }
- let array = StringArray::from(values);
-
- c.bench_function("bit_length", |b| b.iter(|| bench_bit_length(&array)));
-}
-
-criterion_group!(benches, add_benchmark);
-criterion_main!(benches);
diff --git a/arrow/benches/boolean_kernels.rs b/arrow/benches/boolean_kernels.rs
deleted file mode 100644
index e9394f7..0000000
--- a/arrow/benches/boolean_kernels.rs
+++ /dev/null
@@ -1,70 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-#[macro_use]
-extern crate criterion;
-use criterion::Criterion;
-
-use arrow::util::bench_util::create_boolean_array;
-
-extern crate arrow;
-
-use arrow::array::*;
-use arrow::compute::kernels::boolean as boolean_kernels;
-
-fn bench_and(lhs: &BooleanArray, rhs: &BooleanArray) {
- criterion::black_box(boolean_kernels::and(lhs, rhs).unwrap());
-}
-
-fn bench_or(lhs: &BooleanArray, rhs: &BooleanArray) {
- criterion::black_box(boolean_kernels::or(lhs, rhs).unwrap());
-}
-
-fn bench_not(array: &BooleanArray) {
- criterion::black_box(boolean_kernels::not(&array).unwrap());
-}
-
-fn add_benchmark(c: &mut Criterion) {
- let size = 2usize.pow(15);
- let array1 = create_boolean_array(size, 0.0, 0.5);
- let array2 = create_boolean_array(size, 0.0, 0.5);
- c.bench_function("and", |b| b.iter(|| bench_and(&array1, &array2)));
- c.bench_function("or", |b| b.iter(|| bench_or(&array1, &array2)));
- c.bench_function("not", |b| b.iter(|| bench_not(&array1)));
-
- let array1_slice = array1.slice(1, size - 1);
- let array1_slice = array1_slice
- .as_any()
- .downcast_ref::<BooleanArray>()
- .unwrap();
- let array2_slice = array2.slice(1, size - 1);
- let array2_slice = array2_slice
- .as_any()
- .downcast_ref::<BooleanArray>()
- .unwrap();
-
- c.bench_function("and_sliced", |b| {
- b.iter(|| bench_and(&array1_slice, &array2_slice))
- });
- c.bench_function("or_sliced", |b| {
- b.iter(|| bench_or(&array1_slice, &array2_slice))
- });
- c.bench_function("not_sliced", |b| b.iter(|| bench_not(&array1_slice)));
-}
-
-criterion_group!(benches, add_benchmark);
-criterion_main!(benches);
diff --git a/arrow/benches/buffer_bit_ops.rs b/arrow/benches/buffer_bit_ops.rs
deleted file mode 100644
index 063f39c..0000000
--- a/arrow/benches/buffer_bit_ops.rs
+++ /dev/null
@@ -1,59 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-#[macro_use]
-extern crate criterion;
-use criterion::Criterion;
-
-extern crate arrow;
-
-use arrow::buffer::{Buffer, MutableBuffer};
-
-/// Helper function to create arrays
-fn create_buffer(size: usize) -> Buffer {
- let mut result = MutableBuffer::new(size).with_bitset(size, false);
-
- for i in 0..size {
- result.as_slice_mut()[i] = 0b01010101 << i << (i % 4);
- }
-
- result.into()
-}
-
-fn bench_buffer_and(left: &Buffer, right: &Buffer) {
- criterion::black_box((left & right).unwrap());
-}
-
-fn bench_buffer_or(left: &Buffer, right: &Buffer) {
- criterion::black_box((left | right).unwrap());
-}
-
-fn bit_ops_benchmark(c: &mut Criterion) {
- let left = create_buffer(512 * 10);
- let right = create_buffer(512 * 10);
-
- c.bench_function("buffer_bit_ops and", |b| {
- b.iter(|| bench_buffer_and(&left, &right))
- });
-
- c.bench_function("buffer_bit_ops or", |b| {
- b.iter(|| bench_buffer_or(&left, &right))
- });
-}
-
-criterion_group!(benches, bit_ops_benchmark);
-criterion_main!(benches);
diff --git a/arrow/benches/buffer_create.rs b/arrow/benches/buffer_create.rs
deleted file mode 100644
index d628e03..0000000
--- a/arrow/benches/buffer_create.rs
+++ /dev/null
@@ -1,190 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-#[macro_use]
-extern crate criterion;
-use arrow::util::test_util::seedable_rng;
-use criterion::Criterion;
-use rand::distributions::Uniform;
-use rand::Rng;
-
-extern crate arrow;
-
-use arrow::{
- buffer::{Buffer, MutableBuffer},
- datatypes::ToByteSlice,
-};
-
-fn mutable_buffer_from_iter(data: &[Vec<bool>]) -> Vec<Buffer> {
- criterion::black_box(
- data.iter()
- .map(|vec| vec.iter().copied().collect::<MutableBuffer>().into())
- .collect::<Vec<_>>(),
- )
-}
-
-fn buffer_from_iter(data: &[Vec<bool>]) -> Vec<Buffer> {
- criterion::black_box(
- data.iter()
- .map(|vec| vec.iter().copied().collect::<Buffer>())
- .collect::<Vec<_>>(),
- )
-}
-
-fn mutable_buffer_iter_bitset(data: &[Vec<bool>]) -> Vec<Buffer> {
- criterion::black_box({
- data.iter()
- .map(|datum| {
- let mut result = MutableBuffer::new((data.len() + 7) / 8)
- .with_bitset(datum.len(), false);
- for (i, value) in datum.iter().enumerate() {
- if *value {
- unsafe {
- arrow::util::bit_util::set_bit_raw(result.as_mut_ptr(), i);
- }
- }
- }
- result.into()
- })
- .collect::<Vec<_>>()
- })
-}
-
-fn mutable_iter_extend_from_slice(data: &[Vec<u32>], capacity: usize) -> Buffer {
- criterion::black_box({
- let mut result = MutableBuffer::new(capacity);
-
- data.iter().for_each(|vec| {
- vec.iter()
- .for_each(|elem| result.extend_from_slice(elem.to_byte_slice()))
- });
-
- result.into()
- })
-}
-
-fn mutable_buffer(data: &[Vec<u32>], capacity: usize) -> Buffer {
- criterion::black_box({
- let mut result = MutableBuffer::new(capacity);
-
- data.iter().for_each(|vec| result.extend_from_slice(vec));
-
- result.into()
- })
-}
-
-fn mutable_buffer_extend(data: &[Vec<u32>], capacity: usize) -> Buffer {
- criterion::black_box({
- let mut result = MutableBuffer::new(capacity);
-
- data.iter()
- .for_each(|vec| result.extend(vec.iter().copied()));
-
- result.into()
- })
-}
-
-fn from_slice(data: &[Vec<u32>], capacity: usize) -> Buffer {
- criterion::black_box({
- let mut a = Vec::<u32>::with_capacity(capacity);
-
- data.iter().for_each(|vec| a.extend(vec));
-
- Buffer::from(a.to_byte_slice())
- })
-}
-
-fn create_data(size: usize) -> Vec<Vec<u32>> {
- let rng = &mut seedable_rng();
- let range = Uniform::new(0, 33);
-
- (0..size)
- .map(|_| {
- let size = rng.sample(range);
- seedable_rng()
- .sample_iter(&range)
- .take(size as usize)
- .collect()
- })
- .collect()
-}
-
-fn create_data_bool(size: usize) -> Vec<Vec<bool>> {
- let rng = &mut seedable_rng();
- let range = Uniform::new(0, 33);
-
- (0..size)
- .map(|_| {
- let size = rng.sample(range);
- seedable_rng()
- .sample_iter(&range)
- .take(size as usize)
- .map(|x| x > 15)
- .collect()
- })
- .collect()
-}
-fn benchmark(c: &mut Criterion) {
- let size = 2usize.pow(15);
- let data = create_data(size);
-
- let bool_data = create_data_bool(size);
- let cap = data.iter().map(|i| i.len()).sum();
- let byte_cap = cap * std::mem::size_of::<u32>();
-
- c.bench_function("mutable iter extend_from_slice", |b| {
- b.iter(|| {
- mutable_iter_extend_from_slice(
- criterion::black_box(&data),
- criterion::black_box(0),
- )
- })
- });
- c.bench_function("mutable", |b| {
- b.iter(|| mutable_buffer(criterion::black_box(&data), criterion::black_box(0)))
- });
-
- c.bench_function("mutable extend", |b| {
- b.iter(|| mutable_buffer_extend(&data, 0))
- });
-
- c.bench_function("mutable prepared", |b| {
- b.iter(|| {
- mutable_buffer(criterion::black_box(&data), criterion::black_box(byte_cap))
- })
- });
-
- c.bench_function("from_slice", |b| {
- b.iter(|| from_slice(criterion::black_box(&data), criterion::black_box(0)))
- });
- c.bench_function("from_slice prepared", |b| {
- b.iter(|| from_slice(criterion::black_box(&data), criterion::black_box(cap)))
- });
-
- c.bench_function("MutableBuffer iter bitset", |b| {
- b.iter(|| mutable_buffer_iter_bitset(criterion::black_box(&bool_data)))
- });
- c.bench_function("MutableBuffer::from_iter bool", |b| {
- b.iter(|| mutable_buffer_from_iter(criterion::black_box(&bool_data)))
- });
- c.bench_function("Buffer::from_iter bool", |b| {
- b.iter(|| buffer_from_iter(criterion::black_box(&bool_data)))
- });
-}
-
-criterion_group!(benches, benchmark);
-criterion_main!(benches);
diff --git a/arrow/benches/builder.rs b/arrow/benches/builder.rs
deleted file mode 100644
index fd9f319..0000000
--- a/arrow/benches/builder.rs
+++ /dev/null
@@ -1,116 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-extern crate arrow;
-extern crate criterion;
-extern crate rand;
-
-use std::mem::size_of;
-
-use criterion::*;
-use rand::distributions::Standard;
-
-use arrow::array::*;
-use arrow::util::test_util::seedable_rng;
-use rand::Rng;
-
-// Build arrays with 512k elements.
-const BATCH_SIZE: usize = 8 << 10;
-const NUM_BATCHES: usize = 64;
-
-fn bench_primitive(c: &mut Criterion) {
- let data: [i64; BATCH_SIZE] = [100; BATCH_SIZE];
-
- let mut group = c.benchmark_group("bench_primitive");
- group.throughput(Throughput::Bytes(
- ((data.len() * NUM_BATCHES * size_of::<i64>()) as u32).into(),
- ));
- group.bench_function("bench_primitive", |b| {
- b.iter(|| {
- let mut builder = Int64Builder::new(64);
- for _ in 0..NUM_BATCHES {
- let _ = black_box(builder.append_slice(&data[..]));
- }
- black_box(builder.finish());
- })
- });
- group.finish();
-}
-
-fn bench_primitive_nulls(c: &mut Criterion) {
- let mut group = c.benchmark_group("bench_primitive_nulls");
- group.bench_function("bench_primitive_nulls", |b| {
- b.iter(|| {
- let mut builder = UInt8Builder::new(64);
- for _ in 0..NUM_BATCHES * BATCH_SIZE {
- let _ = black_box(builder.append_null());
- }
- black_box(builder.finish());
- })
- });
- group.finish();
-}
-
-fn bench_bool(c: &mut Criterion) {
- let data: Vec<bool> = seedable_rng()
- .sample_iter(&Standard)
- .take(BATCH_SIZE)
- .collect();
- let data_len = data.len();
-
- let mut group = c.benchmark_group("bench_bool");
- group.throughput(Throughput::Bytes(
- ((data_len * NUM_BATCHES * size_of::<bool>()) as u32).into(),
- ));
- group.bench_function("bench_bool", |b| {
- b.iter(|| {
- let mut builder = BooleanBuilder::new(64);
- for _ in 0..NUM_BATCHES {
- let _ = black_box(builder.append_slice(&data[..]));
- }
- black_box(builder.finish());
- })
- });
- group.finish();
-}
-
-fn bench_string(c: &mut Criterion) {
- const SAMPLE_STRING: &str = "sample string";
- let mut group = c.benchmark_group("bench_primitive");
- group.throughput(Throughput::Bytes(
- ((BATCH_SIZE * NUM_BATCHES * SAMPLE_STRING.len()) as u32).into(),
- ));
- group.bench_function("bench_string", |b| {
- b.iter(|| {
- let mut builder = StringBuilder::new(64);
- for _ in 0..NUM_BATCHES * BATCH_SIZE {
- let _ = black_box(builder.append_value(SAMPLE_STRING));
- }
- black_box(builder.finish());
- })
- });
- group.finish();
-}
-
-criterion_group!(
- benches,
- bench_primitive,
- bench_primitive_nulls,
- bench_bool,
- bench_string
-);
-criterion_main!(benches);
diff --git a/arrow/benches/cast_kernels.rs b/arrow/benches/cast_kernels.rs
deleted file mode 100644
index d164e1f..0000000
--- a/arrow/benches/cast_kernels.rs
+++ /dev/null
@@ -1,185 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-#[macro_use]
-extern crate criterion;
-use criterion::Criterion;
-use rand::distributions::{Distribution, Standard, Uniform};
-use rand::Rng;
-
-use std::sync::Arc;
-
-extern crate arrow;
-
-use arrow::array::*;
-use arrow::compute::cast;
-use arrow::datatypes::*;
-use arrow::util::bench_util::*;
-use arrow::util::test_util::seedable_rng;
-
-fn build_array<T: ArrowPrimitiveType>(size: usize) -> ArrayRef
-where
- Standard: Distribution<T::Native>,
-{
- let array = create_primitive_array::<T>(size, 0.1);
- Arc::new(array)
-}
-
-fn build_utf8_date_array(size: usize, with_nulls: bool) -> ArrayRef {
- use chrono::NaiveDate;
-
- // use random numbers to avoid spurious compiler optimizations wrt to branching
- let mut rng = seedable_rng();
- let mut builder = StringBuilder::new(size);
- let range = Uniform::new(0, 737776);
-
- for _ in 0..size {
- if with_nulls && rng.gen::<f32>() > 0.8 {
- builder.append_null().unwrap();
- } else {
- let string = NaiveDate::from_num_days_from_ce(rng.sample(range))
- .format("%Y-%m-%d")
- .to_string();
- builder.append_value(&string).unwrap();
- }
- }
- Arc::new(builder.finish())
-}
-
-fn build_utf8_date_time_array(size: usize, with_nulls: bool) -> ArrayRef {
- use chrono::NaiveDateTime;
-
- // use random numbers to avoid spurious compiler optimizations wrt to branching
- let mut rng = seedable_rng();
- let mut builder = StringBuilder::new(size);
- let range = Uniform::new(0, 1608071414123);
-
- for _ in 0..size {
- if with_nulls && rng.gen::<f32>() > 0.8 {
- builder.append_null().unwrap();
- } else {
- let string = NaiveDateTime::from_timestamp(rng.sample(range), 0)
- .format("%Y-%m-%dT%H:%M:%S")
- .to_string();
- builder.append_value(&string).unwrap();
- }
- }
- Arc::new(builder.finish())
-}
-
-// cast array from specified primitive array type to desired data type
-fn cast_array(array: &ArrayRef, to_type: DataType) {
- criterion::black_box(cast(array, &to_type).unwrap());
-}
-
-fn add_benchmark(c: &mut Criterion) {
- let i32_array = build_array::<Int32Type>(512);
- let i64_array = build_array::<Int64Type>(512);
- let f32_array = build_array::<Float32Type>(512);
- let f32_utf8_array = cast(&build_array::<Float32Type>(512), &DataType::Utf8).unwrap();
-
- let f64_array = build_array::<Float64Type>(512);
- let date64_array = build_array::<Date64Type>(512);
- let date32_array = build_array::<Date32Type>(512);
- let time32s_array = build_array::<Time32SecondType>(512);
- let time64ns_array = build_array::<Time64NanosecondType>(512);
- let time_ns_array = build_array::<TimestampNanosecondType>(512);
- let time_ms_array = build_array::<TimestampMillisecondType>(512);
- let utf8_date_array = build_utf8_date_array(512, true);
- let utf8_date_time_array = build_utf8_date_time_array(512, true);
-
- c.bench_function("cast int32 to int32 512", |b| {
- b.iter(|| cast_array(&i32_array, DataType::Int32))
- });
- c.bench_function("cast int32 to uint32 512", |b| {
- b.iter(|| cast_array(&i32_array, DataType::UInt32))
- });
- c.bench_function("cast int32 to float32 512", |b| {
- b.iter(|| cast_array(&i32_array, DataType::Float32))
- });
- c.bench_function("cast int32 to float64 512", |b| {
- b.iter(|| cast_array(&i32_array, DataType::Float64))
- });
- c.bench_function("cast int32 to int64 512", |b| {
- b.iter(|| cast_array(&i32_array, DataType::Int64))
- });
- c.bench_function("cast float32 to int32 512", |b| {
- b.iter(|| cast_array(&f32_array, DataType::Int32))
- });
- c.bench_function("cast float64 to float32 512", |b| {
- b.iter(|| cast_array(&f64_array, DataType::Float32))
- });
- c.bench_function("cast float64 to uint64 512", |b| {
- b.iter(|| cast_array(&f64_array, DataType::UInt64))
- });
- c.bench_function("cast int64 to int32 512", |b| {
- b.iter(|| cast_array(&i64_array, DataType::Int32))
- });
- c.bench_function("cast date64 to date32 512", |b| {
- b.iter(|| cast_array(&date64_array, DataType::Date32))
- });
- c.bench_function("cast date32 to date64 512", |b| {
- b.iter(|| cast_array(&date32_array, DataType::Date64))
- });
- c.bench_function("cast time32s to time32ms 512", |b| {
- b.iter(|| cast_array(&time32s_array, DataType::Time32(TimeUnit::Millisecond)))
- });
- c.bench_function("cast time32s to time64us 512", |b| {
- b.iter(|| cast_array(&time32s_array, DataType::Time64(TimeUnit::Microsecond)))
- });
- c.bench_function("cast time64ns to time32s 512", |b| {
- b.iter(|| cast_array(&time64ns_array, DataType::Time32(TimeUnit::Second)))
- });
- c.bench_function("cast timestamp_ns to timestamp_s 512", |b| {
- b.iter(|| {
- cast_array(
- &time_ns_array,
- DataType::Timestamp(TimeUnit::Nanosecond, None),
- )
- })
- });
- c.bench_function("cast timestamp_ms to timestamp_ns 512", |b| {
- b.iter(|| {
- cast_array(
- &time_ms_array,
- DataType::Timestamp(TimeUnit::Nanosecond, None),
- )
- })
- });
- c.bench_function("cast utf8 to f32", |b| {
- b.iter(|| cast_array(&f32_utf8_array, DataType::Float32))
- });
- c.bench_function("cast i64 to string 512", |b| {
- b.iter(|| cast_array(&i64_array, DataType::Utf8))
- });
- c.bench_function("cast f32 to string 512", |b| {
- b.iter(|| cast_array(&f32_array, DataType::Utf8))
- });
-
- c.bench_function("cast timestamp_ms to i64 512", |b| {
- b.iter(|| cast_array(&time_ms_array, DataType::Int64))
- });
- c.bench_function("cast utf8 to date32 512", |b| {
- b.iter(|| cast_array(&utf8_date_array, DataType::Date32))
- });
- c.bench_function("cast utf8 to date64 512", |b| {
- b.iter(|| cast_array(&utf8_date_time_array, DataType::Date64))
- });
-}
-
-criterion_group!(benches, add_benchmark);
-criterion_main!(benches);
diff --git a/arrow/benches/comparison_kernels.rs b/arrow/benches/comparison_kernels.rs
deleted file mode 100644
index a3df556..0000000
--- a/arrow/benches/comparison_kernels.rs
+++ /dev/null
@@ -1,201 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-#[macro_use]
-extern crate criterion;
-use criterion::Criterion;
-
-extern crate arrow;
-
-use arrow::compute::*;
-use arrow::datatypes::ArrowNumericType;
-use arrow::util::bench_util::*;
-use arrow::{array::*, datatypes::Float32Type};
-
-fn bench_eq<T>(arr_a: &PrimitiveArray<T>, arr_b: &PrimitiveArray<T>)
-where
- T: ArrowNumericType,
-{
- eq(criterion::black_box(arr_a), criterion::black_box(arr_b)).unwrap();
-}
-
-fn bench_eq_scalar<T>(arr_a: &PrimitiveArray<T>, value_b: T::Native)
-where
- T: ArrowNumericType,
-{
- eq_scalar(criterion::black_box(arr_a), criterion::black_box(value_b)).unwrap();
-}
-
-fn bench_neq<T>(arr_a: &PrimitiveArray<T>, arr_b: &PrimitiveArray<T>)
-where
- T: ArrowNumericType,
-{
- neq(criterion::black_box(arr_a), criterion::black_box(arr_b)).unwrap();
-}
-
-fn bench_neq_scalar<T>(arr_a: &PrimitiveArray<T>, value_b: T::Native)
-where
- T: ArrowNumericType,
-{
- neq_scalar(criterion::black_box(arr_a), criterion::black_box(value_b)).unwrap();
-}
-
-fn bench_lt<T>(arr_a: &PrimitiveArray<T>, arr_b: &PrimitiveArray<T>)
-where
- T: ArrowNumericType,
-{
- lt(criterion::black_box(arr_a), criterion::black_box(arr_b)).unwrap();
-}
-
-fn bench_lt_scalar<T>(arr_a: &PrimitiveArray<T>, value_b: T::Native)
-where
- T: ArrowNumericType,
-{
- lt_scalar(criterion::black_box(arr_a), criterion::black_box(value_b)).unwrap();
-}
-
-fn bench_lt_eq<T>(arr_a: &PrimitiveArray<T>, arr_b: &PrimitiveArray<T>)
-where
- T: ArrowNumericType,
-{
- lt_eq(criterion::black_box(arr_a), criterion::black_box(arr_b)).unwrap();
-}
-
-fn bench_lt_eq_scalar<T>(arr_a: &PrimitiveArray<T>, value_b: T::Native)
-where
- T: ArrowNumericType,
-{
- lt_eq_scalar(criterion::black_box(arr_a), criterion::black_box(value_b)).unwrap();
-}
-
-fn bench_gt<T>(arr_a: &PrimitiveArray<T>, arr_b: &PrimitiveArray<T>)
-where
- T: ArrowNumericType,
-{
- gt(criterion::black_box(arr_a), criterion::black_box(arr_b)).unwrap();
-}
-
-fn bench_gt_scalar<T>(arr_a: &PrimitiveArray<T>, value_b: T::Native)
-where
- T: ArrowNumericType,
-{
- gt_scalar(criterion::black_box(arr_a), criterion::black_box(value_b)).unwrap();
-}
-
-fn bench_gt_eq<T>(arr_a: &PrimitiveArray<T>, arr_b: &PrimitiveArray<T>)
-where
- T: ArrowNumericType,
-{
- gt_eq(criterion::black_box(arr_a), criterion::black_box(arr_b)).unwrap();
-}
-
-fn bench_gt_eq_scalar<T>(arr_a: &PrimitiveArray<T>, value_b: T::Native)
-where
- T: ArrowNumericType,
-{
- gt_eq_scalar(criterion::black_box(arr_a), criterion::black_box(value_b)).unwrap();
-}
-
-fn bench_like_utf8_scalar(arr_a: &StringArray, value_b: &str) {
- like_utf8_scalar(criterion::black_box(arr_a), criterion::black_box(value_b)).unwrap();
-}
-
-fn bench_nlike_utf8_scalar(arr_a: &StringArray, value_b: &str) {
- nlike_utf8_scalar(criterion::black_box(arr_a), criterion::black_box(value_b))
- .unwrap();
-}
-
-fn add_benchmark(c: &mut Criterion) {
- let size = 65536;
- let arr_a = create_primitive_array_with_seed::<Float32Type>(size, 0.0, 42);
- let arr_b = create_primitive_array_with_seed::<Float32Type>(size, 0.0, 43);
-
- let arr_string = create_string_array::<i32>(size, 0.0);
-
- c.bench_function("eq Float32", |b| b.iter(|| bench_eq(&arr_a, &arr_b)));
- c.bench_function("eq scalar Float32", |b| {
- b.iter(|| bench_eq_scalar(&arr_a, 1.0))
- });
-
- c.bench_function("neq Float32", |b| b.iter(|| bench_neq(&arr_a, &arr_b)));
- c.bench_function("neq scalar Float32", |b| {
- b.iter(|| bench_neq_scalar(&arr_a, 1.0))
- });
-
- c.bench_function("lt Float32", |b| b.iter(|| bench_lt(&arr_a, &arr_b)));
- c.bench_function("lt scalar Float32", |b| {
- b.iter(|| bench_lt_scalar(&arr_a, 1.0))
- });
-
- c.bench_function("lt_eq Float32", |b| b.iter(|| bench_lt_eq(&arr_a, &arr_b)));
- c.bench_function("lt_eq scalar Float32", |b| {
- b.iter(|| bench_lt_eq_scalar(&arr_a, 1.0))
- });
-
- c.bench_function("gt Float32", |b| b.iter(|| bench_gt(&arr_a, &arr_b)));
- c.bench_function("gt scalar Float32", |b| {
- b.iter(|| bench_gt_scalar(&arr_a, 1.0))
- });
-
- c.bench_function("gt_eq Float32", |b| b.iter(|| bench_gt_eq(&arr_a, &arr_b)));
- c.bench_function("gt_eq scalar Float32", |b| {
- b.iter(|| bench_gt_eq_scalar(&arr_a, 1.0))
- });
-
- c.bench_function("like_utf8 scalar equals", |b| {
- b.iter(|| bench_like_utf8_scalar(&arr_string, "xxxx"))
- });
-
- c.bench_function("like_utf8 scalar contains", |b| {
- b.iter(|| bench_like_utf8_scalar(&arr_string, "%xxxx%"))
- });
-
- c.bench_function("like_utf8 scalar ends with", |b| {
- b.iter(|| bench_like_utf8_scalar(&arr_string, "xxxx%"))
- });
-
- c.bench_function("like_utf8 scalar starts with", |b| {
- b.iter(|| bench_like_utf8_scalar(&arr_string, "%xxxx"))
- });
-
- c.bench_function("like_utf8 scalar complex", |b| {
- b.iter(|| bench_like_utf8_scalar(&arr_string, "%xx_xx%xxx"))
- });
-
- c.bench_function("nlike_utf8 scalar equals", |b| {
- b.iter(|| bench_nlike_utf8_scalar(&arr_string, "xxxx"))
- });
-
- c.bench_function("nlike_utf8 scalar contains", |b| {
- b.iter(|| bench_nlike_utf8_scalar(&arr_string, "%xxxx%"))
- });
-
- c.bench_function("nlike_utf8 scalar ends with", |b| {
- b.iter(|| bench_nlike_utf8_scalar(&arr_string, "xxxx%"))
- });
-
- c.bench_function("nlike_utf8 scalar starts with", |b| {
- b.iter(|| bench_nlike_utf8_scalar(&arr_string, "%xxxx"))
- });
-
- c.bench_function("nlike_utf8 scalar complex", |b| {
- b.iter(|| bench_nlike_utf8_scalar(&arr_string, "%xx_xx%xxx"))
- });
-}
-
-criterion_group!(benches, add_benchmark);
-criterion_main!(benches);
diff --git a/arrow/benches/concatenate_kernel.rs b/arrow/benches/concatenate_kernel.rs
deleted file mode 100644
index 3fff2ab..0000000
--- a/arrow/benches/concatenate_kernel.rs
+++ /dev/null
@@ -1,66 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-#[macro_use]
-extern crate criterion;
-use criterion::Criterion;
-
-extern crate arrow;
-
-use arrow::array::*;
-use arrow::compute::concat;
-use arrow::datatypes::*;
-use arrow::util::bench_util::*;
-
-fn bench_concat(v1: &dyn Array, v2: &dyn Array) {
- criterion::black_box(concat(&[v1, v2]).unwrap());
-}
-
-fn bench_concat_arrays(arrays: &[&dyn Array]) {
- criterion::black_box(concat(arrays).unwrap());
-}
-
-fn add_benchmark(c: &mut Criterion) {
- let v1 = create_primitive_array::<Int32Type>(1024, 0.0);
- let v2 = create_primitive_array::<Int32Type>(1024, 0.0);
- c.bench_function("concat i32 1024", |b| b.iter(|| bench_concat(&v1, &v2)));
-
- let v1 = create_primitive_array::<Int32Type>(1024, 0.5);
- let v2 = create_primitive_array::<Int32Type>(1024, 0.5);
- c.bench_function("concat i32 nulls 1024", |b| {
- b.iter(|| bench_concat(&v1, &v2))
- });
-
- let small_array = create_primitive_array::<Int32Type>(4, 0.0);
- let arrays: Vec<_> = (0..1024).map(|_| &small_array as &dyn Array).collect();
- c.bench_function("concat 1024 arrays i32 4", |b| {
- b.iter(|| bench_concat_arrays(&arrays))
- });
-
- let v1 = create_string_array::<i32>(1024, 0.0);
- let v2 = create_string_array::<i32>(1024, 0.0);
- c.bench_function("concat str 1024", |b| b.iter(|| bench_concat(&v1, &v2)));
-
- let v1 = create_string_array::<i32>(1024, 0.5);
- let v2 = create_string_array::<i32>(1024, 0.5);
- c.bench_function("concat str nulls 1024", |b| {
- b.iter(|| bench_concat(&v1, &v2))
- });
-}
-
-criterion_group!(benches, add_benchmark);
-criterion_main!(benches);
diff --git a/arrow/benches/csv_writer.rs b/arrow/benches/csv_writer.rs
deleted file mode 100644
index 62c5da9..0000000
--- a/arrow/benches/csv_writer.rs
+++ /dev/null
@@ -1,75 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-extern crate arrow;
-extern crate criterion;
-
-use criterion::*;
-
-use arrow::array::*;
-#[cfg(feature = "csv")]
-use arrow::csv;
-use arrow::datatypes::*;
-use arrow::record_batch::RecordBatch;
-use std::fs::File;
-use std::sync::Arc;
-
-fn criterion_benchmark(c: &mut Criterion) {
- #[cfg(feature = "csv")]
- {
- let schema = Schema::new(vec![
- Field::new("c1", DataType::Utf8, false),
- Field::new("c2", DataType::Float64, true),
- Field::new("c3", DataType::UInt32, false),
- Field::new("c4", DataType::Boolean, true),
- ]);
-
- let c1 = StringArray::from(vec![
- "Lorem ipsum dolor sit amet",
- "consectetur adipiscing elit",
- "sed do eiusmod tempor",
- ]);
- let c2 = PrimitiveArray::<Float64Type>::from(vec![
- Some(123.564532),
- None,
- Some(-556132.25),
- ]);
- let c3 = PrimitiveArray::<UInt32Type>::from(vec![3, 2, 1]);
- let c4 = BooleanArray::from(vec![Some(true), Some(false), None]);
-
- let b = RecordBatch::try_new(
- Arc::new(schema),
- vec![Arc::new(c1), Arc::new(c2), Arc::new(c3), Arc::new(c4)],
- )
- .unwrap();
- let file = File::create("target/bench_write_csv.csv").unwrap();
- let mut writer = csv::Writer::new(file);
- let batches = vec![&b, &b, &b, &b, &b, &b, &b, &b, &b, &b, &b];
-
- c.bench_function("record_batches_to_csv", |b| {
- b.iter(|| {
- #[allow(clippy::unit_arg)]
- criterion::black_box(for batch in &batches {
- writer.write(batch).unwrap()
- });
- });
- });
- }
-}
-
-criterion_group!(benches, criterion_benchmark);
-criterion_main!(benches);
diff --git a/arrow/benches/equal.rs b/arrow/benches/equal.rs
deleted file mode 100644
index af53550..0000000
--- a/arrow/benches/equal.rs
+++ /dev/null
@@ -1,57 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-// Allowed because we use `arr == arr` in benchmarks
-#![allow(clippy::eq_op)]
-
-#[macro_use]
-extern crate criterion;
-use criterion::Criterion;
-
-extern crate arrow;
-
-use arrow::util::bench_util::*;
-use arrow::{array::*, datatypes::Float32Type};
-
-fn bench_equal<A: Array + PartialEq<A>>(arr_a: &A) {
- criterion::black_box(arr_a == arr_a);
-}
-
-fn add_benchmark(c: &mut Criterion) {
- let arr_a = create_primitive_array::<Float32Type>(512, 0.0);
- c.bench_function("equal_512", |b| b.iter(|| bench_equal(&arr_a)));
-
- let arr_a_nulls = create_primitive_array::<Float32Type>(512, 0.5);
- c.bench_function("equal_nulls_512", |b| b.iter(|| bench_equal(&arr_a_nulls)));
-
- let arr_a = create_string_array::<i32>(512, 0.0);
- c.bench_function("equal_string_512", |b| b.iter(|| bench_equal(&arr_a)));
-
- let arr_a_nulls = create_string_array::<i32>(512, 0.5);
- c.bench_function("equal_string_nulls_512", |b| {
- b.iter(|| bench_equal(&arr_a_nulls))
- });
-
- let arr_a = create_boolean_array(512, 0.0, 0.5);
- c.bench_function("equal_bool_512", |b| b.iter(|| bench_equal(&arr_a)));
-
- let arr_a = create_boolean_array(513, 0.0, 0.5);
- c.bench_function("equal_bool_513", |b| b.iter(|| bench_equal(&arr_a)));
-}
-
-criterion_group!(benches, add_benchmark);
-criterion_main!(benches);
diff --git a/arrow/benches/filter_kernels.rs b/arrow/benches/filter_kernels.rs
deleted file mode 100644
index ca317b4..0000000
--- a/arrow/benches/filter_kernels.rs
+++ /dev/null
@@ -1,106 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-extern crate arrow;
-
-use arrow::compute::Filter;
-use arrow::util::bench_util::*;
-
-use arrow::array::*;
-use arrow::compute::{build_filter, filter};
-use arrow::datatypes::{Float32Type, UInt8Type};
-
-use criterion::{criterion_group, criterion_main, Criterion};
-
-fn bench_filter(data_array: &dyn Array, filter_array: &BooleanArray) {
- criterion::black_box(filter(data_array, filter_array).unwrap());
-}
-
-fn bench_built_filter<'a>(filter: &Filter<'a>, data: &impl Array) {
- criterion::black_box(filter(&data.data()));
-}
-
-fn add_benchmark(c: &mut Criterion) {
- let size = 65536;
- let filter_array = create_boolean_array(size, 0.0, 0.5);
- let dense_filter_array = create_boolean_array(size, 0.0, 1.0 - 1.0 / 1024.0);
- let sparse_filter_array = create_boolean_array(size, 0.0, 1.0 / 1024.0);
-
- let filter = build_filter(&filter_array).unwrap();
- let dense_filter = build_filter(&dense_filter_array).unwrap();
- let sparse_filter = build_filter(&sparse_filter_array).unwrap();
-
- let data_array = create_primitive_array::<UInt8Type>(size, 0.0);
-
- c.bench_function("filter u8", |b| {
- b.iter(|| bench_filter(&data_array, &filter_array))
- });
- c.bench_function("filter u8 high selectivity", |b| {
- b.iter(|| bench_filter(&data_array, &dense_filter_array))
- });
- c.bench_function("filter u8 low selectivity", |b| {
- b.iter(|| bench_filter(&data_array, &sparse_filter_array))
- });
-
- c.bench_function("filter context u8", |b| {
- b.iter(|| bench_built_filter(&filter, &data_array))
- });
- c.bench_function("filter context u8 high selectivity", |b| {
- b.iter(|| bench_built_filter(&dense_filter, &data_array))
- });
- c.bench_function("filter context u8 low selectivity", |b| {
- b.iter(|| bench_built_filter(&sparse_filter, &data_array))
- });
-
- let data_array = create_primitive_array::<UInt8Type>(size, 0.5);
- c.bench_function("filter context u8 w NULLs", |b| {
- b.iter(|| bench_built_filter(&filter, &data_array))
- });
- c.bench_function("filter context u8 w NULLs high selectivity", |b| {
- b.iter(|| bench_built_filter(&dense_filter, &data_array))
- });
- c.bench_function("filter context u8 w NULLs low selectivity", |b| {
- b.iter(|| bench_built_filter(&sparse_filter, &data_array))
- });
-
- let data_array = create_primitive_array::<Float32Type>(size, 0.5);
- c.bench_function("filter f32", |b| {
- b.iter(|| bench_filter(&data_array, &filter_array))
- });
- c.bench_function("filter context f32", |b| {
- b.iter(|| bench_built_filter(&filter, &data_array))
- });
- c.bench_function("filter context f32 high selectivity", |b| {
- b.iter(|| bench_built_filter(&dense_filter, &data_array))
- });
- c.bench_function("filter context f32 low selectivity", |b| {
- b.iter(|| bench_built_filter(&sparse_filter, &data_array))
- });
-
- let data_array = create_string_array::<i32>(size, 0.5);
- c.bench_function("filter context string", |b| {
- b.iter(|| bench_built_filter(&filter, &data_array))
- });
- c.bench_function("filter context string high selectivity", |b| {
- b.iter(|| bench_built_filter(&dense_filter, &data_array))
- });
- c.bench_function("filter context string low selectivity", |b| {
- b.iter(|| bench_built_filter(&sparse_filter, &data_array))
- });
-}
-
-criterion_group!(benches, add_benchmark);
-criterion_main!(benches);
diff --git a/arrow/benches/json_reader.rs b/arrow/benches/json_reader.rs
deleted file mode 100644
index ef3ddf0..0000000
--- a/arrow/benches/json_reader.rs
+++ /dev/null
@@ -1,112 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-extern crate arrow;
-extern crate criterion;
-
-use criterion::*;
-
-use arrow::datatypes::*;
-use arrow::json::ReaderBuilder;
-use std::io::Cursor;
-use std::sync::Arc;
-
-fn json_primitive_to_record_batch() {
- let schema = Arc::new(Schema::new(vec![
- Field::new("c1", DataType::Utf8, true),
- Field::new("c2", DataType::Float64, true),
- Field::new("c3", DataType::UInt32, true),
- Field::new("c4", DataType::Boolean, true),
- ]));
- let builder = ReaderBuilder::new().with_schema(schema).with_batch_size(64);
- let json_content = r#"
- {"c1": "eleven", "c2": 6.2222222225, "c3": 5.0, "c4": false}
- {"c1": "twelve", "c2": -55555555555555.2, "c3": 3}
- {"c1": null, "c2": 3, "c3": 125, "c4": null}
- {"c2": -35, "c3": 100.0, "c4": true}
- {"c1": "fifteen", "c2": null, "c4": true}
- {"c1": "eleven", "c2": 6.2222222225, "c3": 5.0, "c4": false}
- {"c1": "twelve", "c2": -55555555555555.2, "c3": 3}
- {"c1": null, "c2": 3, "c3": 125, "c4": null}
- {"c2": -35, "c3": 100.0, "c4": true}
- {"c1": "fifteen", "c2": null, "c4": true}
- "#;
- let cursor = Cursor::new(json_content);
- let mut reader = builder.build(cursor).unwrap();
- #[allow(clippy::unit_arg)]
- criterion::black_box({
- reader.next().unwrap();
- });
-}
-
-fn json_list_primitive_to_record_batch() {
- let schema = Arc::new(Schema::new(vec![
- Field::new(
- "c1",
- DataType::List(Box::new(Field::new("item", DataType::Utf8, true))),
- true,
- ),
- Field::new(
- "c2",
- DataType::List(Box::new(Field::new("item", DataType::Float64, true))),
- true,
- ),
- Field::new(
- "c3",
- DataType::List(Box::new(Field::new("item", DataType::UInt32, true))),
- true,
- ),
- Field::new(
- "c4",
- DataType::List(Box::new(Field::new("item", DataType::Boolean, true))),
- true,
- ),
- ]));
- let builder = ReaderBuilder::new().with_schema(schema).with_batch_size(64);
- let json_content = r#"
- {"c1": ["eleven"], "c2": [6.2222222225, -3.2, null], "c3": [5.0, 6], "c4": [false, true]}
- {"c1": ["twelve"], "c2": [-55555555555555.2, 12500000.0], "c3": [3, 4, 5]}
- {"c1": null, "c2": [3], "c3": [125, 127, 129], "c4": [null, false, true]}
- {"c2": [-35], "c3": [100.0, 200.0], "c4": null}
- {"c1": ["fifteen"], "c2": [null, 2.1, 1.5, -3], "c4": [true, false, null]}
- {"c1": ["fifteen"], "c2": [], "c4": [true, false, null]}
- {"c1": ["eleven"], "c2": [6.2222222225, -3.2, null], "c3": [5.0, 6], "c4": [false, true]}
- {"c1": ["twelve"], "c2": [-55555555555555.2, 12500000.0], "c3": [3, 4, 5]}
- {"c1": null, "c2": [3], "c3": [125, 127, 129], "c4": [null, false, true]}
- {"c2": [-35], "c3": [100.0, 200.0], "c4": null}
- {"c1": ["fifteen"], "c2": [null, 2.1, 1.5, -3], "c4": [true, false, null]}
- {"c1": ["fifteen"], "c2": [], "c4": [true, false, null]}
- "#;
- let cursor = Cursor::new(json_content);
- let mut reader = builder.build(cursor).unwrap();
- #[allow(clippy::unit_arg)]
- criterion::black_box({
- reader.next().unwrap();
- });
-}
-
-fn criterion_benchmark(c: &mut Criterion) {
- c.bench_function("json_primitive_to_record_batch", |b| {
- b.iter(json_primitive_to_record_batch)
- });
- c.bench_function("json_list_primitive_to_record_batch", |b| {
- b.iter(json_list_primitive_to_record_batch)
- });
-}
-
-criterion_group!(benches, criterion_benchmark);
-criterion_main!(benches);
diff --git a/arrow/benches/length_kernel.rs b/arrow/benches/length_kernel.rs
deleted file mode 100644
index b70f637..0000000
--- a/arrow/benches/length_kernel.rs
+++ /dev/null
@@ -1,47 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-#[macro_use]
-extern crate criterion;
-use criterion::Criterion;
-
-extern crate arrow;
-
-use arrow::array::*;
-use arrow::compute::kernels::length::length;
-
-fn bench_length(array: &StringArray) {
- criterion::black_box(length(array).unwrap());
-}
-
-fn add_benchmark(c: &mut Criterion) {
- fn double_vec<T: Clone>(v: Vec<T>) -> Vec<T> {
- [&v[..], &v[..]].concat()
- }
-
- // double ["hello", " ", "world", "!"] 10 times
- let mut values = vec!["one", "on", "o", ""];
- for _ in 0..10 {
- values = double_vec(values);
- }
- let array = StringArray::from(values);
-
- c.bench_function("length", |b| b.iter(|| bench_length(&array)));
-}
-
-criterion_group!(benches, add_benchmark);
-criterion_main!(benches);
diff --git a/arrow/benches/mutable_array.rs b/arrow/benches/mutable_array.rs
deleted file mode 100644
index 52da38a..0000000
--- a/arrow/benches/mutable_array.rs
+++ /dev/null
@@ -1,60 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-#[macro_use]
-extern crate criterion;
-use criterion::Criterion;
-
-use rand::Rng;
-
-extern crate arrow;
-
-use arrow::util::test_util::seedable_rng;
-use arrow::{array::*, util::bench_util::create_string_array};
-
-fn create_slices(size: usize) -> Vec<(usize, usize)> {
- let rng = &mut seedable_rng();
-
- (0..size)
- .map(|_| {
- let start = rng.gen_range(0, size / 2);
- let end = rng.gen_range(start + 1, size);
- (start, end)
- })
- .collect()
-}
-
-fn bench<T: Array>(v1: &T, slices: &[(usize, usize)]) {
- let mut mutable = MutableArrayData::new(vec![v1.data_ref()], false, 5);
- for (start, end) in slices {
- mutable.extend(0, *start, *end)
- }
- mutable.freeze();
-}
-
-fn add_benchmark(c: &mut Criterion) {
- let v1 = create_string_array::<i32>(1024, 0.0);
- let v2 = create_slices(1024);
- c.bench_function("mutable str 1024", |b| b.iter(|| bench(&v1, &v2)));
-
- let v1 = create_string_array::<i32>(1024, 0.5);
- let v2 = create_slices(1024);
- c.bench_function("mutable str nulls 1024", |b| b.iter(|| bench(&v1, &v2)));
-}
-
-criterion_group!(benches, add_benchmark);
-criterion_main!(benches);
diff --git a/arrow/benches/partition_kernels.rs b/arrow/benches/partition_kernels.rs
deleted file mode 100644
index ae55fbd..0000000
--- a/arrow/benches/partition_kernels.rs
+++ /dev/null
@@ -1,146 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-#[macro_use]
-extern crate criterion;
-use criterion::Criterion;
-use std::sync::Arc;
-extern crate arrow;
-use arrow::compute::kernels::partition::lexicographical_partition_ranges;
-use arrow::compute::kernels::sort::{lexsort, SortColumn};
-use arrow::util::bench_util::*;
-use arrow::{
- array::*,
- datatypes::{ArrowPrimitiveType, Float64Type, UInt8Type},
-};
-use rand::distributions::{Distribution, Standard};
-use std::iter;
-
-fn create_array<T: ArrowPrimitiveType>(size: usize, with_nulls: bool) -> ArrayRef
-where
- Standard: Distribution<T::Native>,
-{
- let null_density = if with_nulls { 0.5 } else { 0.0 };
- let array = create_primitive_array::<T>(size, null_density);
- Arc::new(array)
-}
-
-fn bench_partition(sorted_columns: &[ArrayRef]) {
- let columns = sorted_columns
- .iter()
- .map(|arr| SortColumn {
- values: arr.clone(),
- options: None,
- })
- .collect::<Vec<_>>();
-
- criterion::black_box(
- lexicographical_partition_ranges(&columns)
- .unwrap()
- .collect::<Vec<_>>(),
- );
-}
-
-fn create_sorted_low_cardinality_data(length: usize) -> Vec<ArrayRef> {
- let arr = Int64Array::from_iter_values(
- iter::repeat(1)
- .take(length / 4)
- .chain(iter::repeat(2).take(length / 4))
- .chain(iter::repeat(3).take(length / 4))
- .chain(iter::repeat(4).take(length / 4)),
- );
- lexsort(
- &[SortColumn {
- values: Arc::new(arr),
- options: None,
- }],
- None,
- )
- .unwrap()
-}
-
-fn create_sorted_float_data(pow: u32, with_nulls: bool) -> Vec<ArrayRef> {
- lexsort(
- &[
- SortColumn {
- values: create_array::<Float64Type>(2u64.pow(pow) as usize, with_nulls),
- options: None,
- },
- SortColumn {
- values: create_array::<Float64Type>(2u64.pow(pow) as usize, with_nulls),
- options: None,
- },
- ],
- None,
- )
- .unwrap()
-}
-
-fn create_sorted_data(pow: u32, with_nulls: bool) -> Vec<ArrayRef> {
- lexsort(
- &[
- SortColumn {
- values: create_array::<UInt8Type>(2u64.pow(pow) as usize, with_nulls),
- options: None,
- },
- SortColumn {
- values: create_array::<UInt8Type>(2u64.pow(pow) as usize, with_nulls),
- options: None,
- },
- ],
- None,
- )
- .unwrap()
-}
-
-fn add_benchmark(c: &mut Criterion) {
- let sorted_columns = create_sorted_data(10, false);
- c.bench_function("lexicographical_partition_ranges(u8) 2^10", |b| {
- b.iter(|| bench_partition(&sorted_columns))
- });
-
- let sorted_columns = create_sorted_data(12, false);
- c.bench_function("lexicographical_partition_ranges(u8) 2^12", |b| {
- b.iter(|| bench_partition(&sorted_columns))
- });
-
- let sorted_columns = create_sorted_data(10, true);
- c.bench_function(
- "lexicographical_partition_ranges(u8) 2^10 with nulls",
- |b| b.iter(|| bench_partition(&sorted_columns)),
- );
-
- let sorted_columns = create_sorted_data(12, true);
- c.bench_function(
- "lexicographical_partition_ranges(u8) 2^12 with nulls",
- |b| b.iter(|| bench_partition(&sorted_columns)),
- );
-
- let sorted_columns = create_sorted_float_data(10, false);
- c.bench_function("lexicographical_partition_ranges(f64) 2^10", |b| {
- b.iter(|| bench_partition(&sorted_columns))
- });
-
- let sorted_columns = create_sorted_low_cardinality_data(1024);
- c.bench_function(
- "lexicographical_partition_ranges(low cardinality) 1024",
- |b| b.iter(|| bench_partition(&sorted_columns)),
- );
-}
-
-criterion_group!(benches, add_benchmark);
-criterion_main!(benches);
diff --git a/arrow/benches/sort_kernel.rs b/arrow/benches/sort_kernel.rs
deleted file mode 100644
index 8467b50..0000000
--- a/arrow/benches/sort_kernel.rs
+++ /dev/null
@@ -1,121 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-#[macro_use]
-extern crate criterion;
-use criterion::Criterion;
-
-use std::sync::Arc;
-
-extern crate arrow;
-
-use arrow::compute::kernels::sort::{lexsort, SortColumn};
-use arrow::util::bench_util::*;
-use arrow::{array::*, datatypes::Float32Type};
-
-fn create_array(size: usize, with_nulls: bool) -> ArrayRef {
- let null_density = if with_nulls { 0.5 } else { 0.0 };
- let array = create_primitive_array::<Float32Type>(size, null_density);
- Arc::new(array)
-}
-
-fn bench_sort(array_a: &ArrayRef, array_b: &ArrayRef, limit: Option<usize>) {
- let columns = vec![
- SortColumn {
- values: array_a.clone(),
- options: None,
- },
- SortColumn {
- values: array_b.clone(),
- options: None,
- },
- ];
-
- criterion::black_box(lexsort(&columns, limit).unwrap());
-}
-
-fn add_benchmark(c: &mut Criterion) {
- let arr_a = create_array(2u64.pow(10) as usize, false);
- let arr_b = create_array(2u64.pow(10) as usize, false);
-
- c.bench_function("sort 2^10", |b| b.iter(|| bench_sort(&arr_a, &arr_b, None)));
-
- let arr_a = create_array(2u64.pow(12) as usize, false);
- let arr_b = create_array(2u64.pow(12) as usize, false);
-
- c.bench_function("sort 2^12", |b| b.iter(|| bench_sort(&arr_a, &arr_b, None)));
-
- let arr_a = create_array(2u64.pow(10) as usize, true);
- let arr_b = create_array(2u64.pow(10) as usize, true);
-
- c.bench_function("sort nulls 2^10", |b| {
- b.iter(|| bench_sort(&arr_a, &arr_b, None))
- });
-
- let arr_a = create_array(2u64.pow(12) as usize, true);
- let arr_b = create_array(2u64.pow(12) as usize, true);
-
- c.bench_function("sort nulls 2^12", |b| {
- b.iter(|| bench_sort(&arr_a, &arr_b, None))
- });
-
- // with limit
- {
- let arr_a = create_array(2u64.pow(12) as usize, false);
- let arr_b = create_array(2u64.pow(12) as usize, false);
- c.bench_function("sort 2^12 limit 10", |b| {
- b.iter(|| bench_sort(&arr_a, &arr_b, Some(10)))
- });
-
- let arr_a = create_array(2u64.pow(12) as usize, false);
- let arr_b = create_array(2u64.pow(12) as usize, false);
- c.bench_function("sort 2^12 limit 100", |b| {
- b.iter(|| bench_sort(&arr_a, &arr_b, Some(100)))
- });
-
- let arr_a = create_array(2u64.pow(12) as usize, false);
- let arr_b = create_array(2u64.pow(12) as usize, false);
- c.bench_function("sort 2^12 limit 1000", |b| {
- b.iter(|| bench_sort(&arr_a, &arr_b, Some(1000)))
- });
-
- let arr_a = create_array(2u64.pow(12) as usize, false);
- let arr_b = create_array(2u64.pow(12) as usize, false);
- c.bench_function("sort 2^12 limit 2^12", |b| {
- b.iter(|| bench_sort(&arr_a, &arr_b, Some(2u64.pow(12) as usize)))
- });
-
- let arr_a = create_array(2u64.pow(12) as usize, true);
- let arr_b = create_array(2u64.pow(12) as usize, true);
-
- c.bench_function("sort nulls 2^12 limit 10", |b| {
- b.iter(|| bench_sort(&arr_a, &arr_b, Some(10)))
- });
- c.bench_function("sort nulls 2^12 limit 100", |b| {
- b.iter(|| bench_sort(&arr_a, &arr_b, Some(100)))
- });
- c.bench_function("sort nulls 2^12 limit 1000", |b| {
- b.iter(|| bench_sort(&arr_a, &arr_b, Some(1000)))
- });
- c.bench_function("sort nulls 2^12 limit 2^12", |b| {
- b.iter(|| bench_sort(&arr_a, &arr_b, Some(2u64.pow(12) as usize)))
- });
- }
-}
-
-criterion_group!(benches, add_benchmark);
-criterion_main!(benches);
diff --git a/arrow/benches/take_kernels.rs b/arrow/benches/take_kernels.rs
deleted file mode 100644
index b1d03d7..0000000
--- a/arrow/benches/take_kernels.rs
+++ /dev/null
@@ -1,145 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-#[macro_use]
-extern crate criterion;
-use criterion::Criterion;
-
-use rand::Rng;
-
-extern crate arrow;
-
-use arrow::compute::{take, TakeOptions};
-use arrow::datatypes::*;
-use arrow::util::test_util::seedable_rng;
-use arrow::{array::*, util::bench_util::*};
-
-fn create_random_index(size: usize, null_density: f32) -> UInt32Array {
- let mut rng = seedable_rng();
- let mut builder = UInt32Builder::new(size);
- for _ in 0..size {
- if rng.gen::<f32>() < null_density {
- builder.append_null().unwrap()
- } else {
- let value = rng.gen_range::<u32, _, _>(0u32, size as u32);
- builder.append_value(value).unwrap();
- }
- }
- builder.finish()
-}
-
-fn bench_take(values: &dyn Array, indices: &UInt32Array) {
- criterion::black_box(take(values, &indices, None).unwrap());
-}
-
-fn bench_take_bounds_check(values: &dyn Array, indices: &UInt32Array) {
- criterion::black_box(
- take(values, &indices, Some(TakeOptions { check_bounds: true })).unwrap(),
- );
-}
-
-fn add_benchmark(c: &mut Criterion) {
- let values = create_primitive_array::<Int32Type>(512, 0.0);
- let indices = create_random_index(512, 0.0);
- c.bench_function("take i32 512", |b| b.iter(|| bench_take(&values, &indices)));
- let values = create_primitive_array::<Int32Type>(1024, 0.0);
- let indices = create_random_index(1024, 0.0);
- c.bench_function("take i32 1024", |b| {
- b.iter(|| bench_take(&values, &indices))
- });
-
- let values = create_primitive_array::<Int32Type>(512, 0.0);
- let indices = create_random_index(512, 0.0);
- c.bench_function("take check bounds i32 512", |b| {
- b.iter(|| bench_take_bounds_check(&values, &indices))
- });
- let values = create_primitive_array::<Int32Type>(1024, 0.0);
- let indices = create_random_index(1024, 0.0);
- c.bench_function("take check bounds i32 1024", |b| {
- b.iter(|| bench_take_bounds_check(&values, &indices))
- });
-
- let indices = create_random_index(512, 0.5);
- c.bench_function("take i32 nulls 512", |b| {
- b.iter(|| bench_take(&values, &indices))
- });
- let values = create_primitive_array::<Int32Type>(1024, 0.0);
- let indices = create_random_index(1024, 0.5);
- c.bench_function("take i32 nulls 1024", |b| {
- b.iter(|| bench_take(&values, &indices))
- });
-
- let values = create_boolean_array(512, 0.0, 0.5);
- let indices = create_random_index(512, 0.0);
- c.bench_function("take bool 512", |b| {
- b.iter(|| bench_take(&values, &indices))
- });
- let values = create_boolean_array(1024, 0.0, 0.5);
- let indices = create_random_index(1024, 0.0);
- c.bench_function("take bool 1024", |b| {
- b.iter(|| bench_take(&values, &indices))
- });
-
- let values = create_boolean_array(512, 0.0, 0.5);
- let indices = create_random_index(512, 0.5);
- c.bench_function("take bool nulls 512", |b| {
- b.iter(|| bench_take(&values, &indices))
- });
- let values = create_boolean_array(1024, 0.0, 0.5);
- let indices = create_random_index(1024, 0.5);
- c.bench_function("take bool nulls 1024", |b| {
- b.iter(|| bench_take(&values, &indices))
- });
-
- let values = create_string_array::<i32>(512, 0.0);
- let indices = create_random_index(512, 0.0);
- c.bench_function("take str 512", |b| b.iter(|| bench_take(&values, &indices)));
-
- let values = create_string_array::<i32>(1024, 0.0);
- let indices = create_random_index(1024, 0.0);
- c.bench_function("take str 1024", |b| {
- b.iter(|| bench_take(&values, &indices))
- });
-
- let values = create_string_array::<i32>(512, 0.0);
- let indices = create_random_index(512, 0.5);
- c.bench_function("take str null indices 512", |b| {
- b.iter(|| bench_take(&values, &indices))
- });
-
- let values = create_string_array::<i32>(1024, 0.0);
- let indices = create_random_index(1024, 0.5);
- c.bench_function("take str null indices 1024", |b| {
- b.iter(|| bench_take(&values, &indices))
- });
-
- let values = create_string_array::<i32>(1024, 0.5);
-
- let indices = create_random_index(1024, 0.0);
- c.bench_function("take str null values 1024", |b| {
- b.iter(|| bench_take(&values, &indices))
- });
-
- let values = create_string_array::<i32>(1024, 0.5);
- let indices = create_random_index(1024, 0.5);
- c.bench_function("take str null values null indices 1024", |b| {
- b.iter(|| bench_take(&values, &indices))
- });
-}
-
-criterion_group!(benches, add_benchmark);
-criterion_main!(benches);
diff --git a/arrow/examples/builders.rs b/arrow/examples/builders.rs
deleted file mode 100644
index 61cce0e..0000000
--- a/arrow/examples/builders.rs
+++ /dev/null
@@ -1,131 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-///! Many builders are available to easily create different types of arrow arrays
-extern crate arrow;
-
-use std::sync::Arc;
-
-use arrow::array::{
- Array, ArrayData, BooleanArray, Int32Array, Int32Builder, ListArray, PrimitiveArray,
- StringArray, StructArray,
-};
-use arrow::buffer::Buffer;
-use arrow::datatypes::{DataType, Date64Type, Field, Time64NanosecondType, ToByteSlice};
-
-fn main() {
- // Primitive Arrays
- //
- // Primitive arrays are arrays of fixed-width primitive types (bool, u8, u16, u32,
- // u64, i8, i16, i32, i64, f32, f64)
-
- // Create a new builder with a capacity of 100
- let mut primitive_array_builder = Int32Builder::new(100);
-
- // Append an individual primitive value
- primitive_array_builder.append_value(55).unwrap();
-
- // Append a null value
- primitive_array_builder.append_null().unwrap();
-
- // Append a slice of primitive values
- primitive_array_builder.append_slice(&[39, 89, 12]).unwrap();
-
- // Append lots of values
- primitive_array_builder.append_null().unwrap();
- primitive_array_builder
- .append_slice(&(25..50).collect::<Vec<i32>>())
- .unwrap();
-
- // Build the `PrimitiveArray`
- let primitive_array = primitive_array_builder.finish();
- // Long arrays will have an ellipsis printed in the middle
- println!("{:?}", primitive_array);
-
- // Arrays can also be built from `Vec<Option<T>>`. `None`
- // represents a null value in the array.
- let date_array: PrimitiveArray<Date64Type> =
- vec![Some(1550902545147), None, Some(1550902545147)].into();
- println!("{:?}", date_array);
-
- let time_array: PrimitiveArray<Time64NanosecondType> =
- (0..100).collect::<Vec<i64>>().into();
- println!("{:?}", time_array);
-
- // We can build arrays directly from the underlying buffers.
-
- // BinaryArrays are arrays of byte arrays, where each byte array
- // is a slice of an underlying buffer.
-
- // Array data: ["hello", null, "parquet"]
- let values: [u8; 12] = [
- b'h', b'e', b'l', b'l', b'o', b'p', b'a', b'r', b'q', b'u', b'e', b't',
- ];
- let offsets: [i32; 4] = [0, 5, 5, 12];
-
- let array_data = ArrayData::builder(DataType::Utf8)
- .len(3)
- .add_buffer(Buffer::from(offsets.to_byte_slice()))
- .add_buffer(Buffer::from(&values[..]))
- .null_bit_buffer(Buffer::from([0b00000101]))
- .build();
- let binary_array = StringArray::from(array_data);
- println!("{:?}", binary_array);
-
- // ListArrays are similar to ByteArrays: they are arrays of other
- // arrays, where each child array is a slice of the underlying
- // buffer.
- let value_data = ArrayData::builder(DataType::Int32)
- .len(8)
- .add_buffer(Buffer::from(&[0, 1, 2, 3, 4, 5, 6, 7].to_byte_slice()))
- .build();
-
- // Construct a buffer for value offsets, for the nested array:
- // [[0, 1, 2], [3, 4, 5], [6, 7]]
- let value_offsets = Buffer::from(&[0, 3, 6, 8].to_byte_slice());
-
- // Construct a list array from the above two
- let list_data_type =
- DataType::List(Box::new(Field::new("item", DataType::Int32, false)));
- let list_data = ArrayData::builder(list_data_type)
- .len(3)
- .add_buffer(value_offsets)
- .add_child_data(value_data)
- .build();
- let list_array = ListArray::from(list_data);
-
- println!("{:?}", list_array);
-
- // StructArrays are arrays of tuples, where each tuple element is
- // from a child array. (In other words, they're like zipping
- // multiple columns into one and giving each subcolumn a label.)
-
- // StructArrays can be constructed using the StructArray::from
- // helper, which takes the underlying arrays and field types.
- let struct_array = StructArray::from(vec![
- (
- Field::new("b", DataType::Boolean, false),
- Arc::new(BooleanArray::from(vec![false, false, true, true]))
- as Arc<dyn Array>,
- ),
- (
- Field::new("c", DataType::Int32, false),
- Arc::new(Int32Array::from(vec![42, 28, 19, 31])),
- ),
- ]);
- println!("{:?}", struct_array);
-}
diff --git a/arrow/examples/dynamic_types.rs b/arrow/examples/dynamic_types.rs
deleted file mode 100644
index 58e4156..0000000
--- a/arrow/examples/dynamic_types.rs
+++ /dev/null
@@ -1,101 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-///! This example demonstrates dealing with mixed types dynamically at runtime
-use std::sync::Arc;
-
-extern crate arrow;
-
-use arrow::array::*;
-use arrow::datatypes::*;
-use arrow::error::Result;
-use arrow::record_batch::*;
-
-fn main() -> Result<()> {
- // define schema
- let schema = Schema::new(vec![
- Field::new("id", DataType::Int32, false),
- Field::new(
- "nested",
- DataType::Struct(vec![
- Field::new("a", DataType::Utf8, false),
- Field::new("b", DataType::Float64, false),
- Field::new("c", DataType::Float64, false),
- ]),
- false,
- ),
- ]);
-
- // create some data
- let id = Int32Array::from(vec![1, 2, 3, 4, 5]);
-
- let nested = StructArray::from(vec![
- (
- Field::new("a", DataType::Utf8, false),
- Arc::new(StringArray::from(vec!["a", "b", "c", "d", "e"])) as Arc<dyn Array>,
- ),
- (
- Field::new("b", DataType::Float64, false),
- Arc::new(Float64Array::from(vec![1.1, 2.2, 3.3, 4.4, 5.5])),
- ),
- (
- Field::new("c", DataType::Float64, false),
- Arc::new(Float64Array::from(vec![2.2, 3.3, 4.4, 5.5, 6.6])),
- ),
- ]);
-
- // build a record batch
- let batch =
- RecordBatch::try_new(Arc::new(schema), vec![Arc::new(id), Arc::new(nested)])?;
-
- process(&batch);
- Ok(())
-}
-
-/// Create a new batch by performing a projection of id, nested.c
-fn process(batch: &RecordBatch) {
- let id = batch.column(0);
- let nested = batch
- .column(1)
- .as_any()
- .downcast_ref::<StructArray>()
- .unwrap();
-
- let _nested_b = nested
- .column(1)
- .as_any()
- .downcast_ref::<Float64Array>()
- .unwrap();
- let nested_c: &Float64Array = nested
- .column(2)
- .as_any()
- .downcast_ref::<Float64Array>()
- .unwrap();
-
- let projected_schema = Schema::new(vec![
- Field::new("id", DataType::Int32, false),
- Field::new("sum", DataType::Float64, false),
- ]);
-
- let _ = RecordBatch::try_new(
- Arc::new(projected_schema),
- vec![
- id.clone(), // NOTE: this is cloning the Arc not the array data
- Arc::new(Float64Array::from(nested_c.data().clone())),
- ],
- );
-}
diff --git a/arrow/examples/read_csv.rs b/arrow/examples/read_csv.rs
deleted file mode 100644
index 506b898..0000000
--- a/arrow/examples/read_csv.rs
+++ /dev/null
@@ -1,48 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-extern crate arrow;
-
-use std::fs::File;
-use std::sync::Arc;
-
-#[cfg(feature = "csv")]
-use arrow::csv;
-use arrow::datatypes::{DataType, Field, Schema};
-#[cfg(feature = "prettyprint")]
-use arrow::util::pretty::print_batches;
-
-fn main() {
- #[cfg(feature = "csv")]
- {
- let schema = Schema::new(vec![
- Field::new("city", DataType::Utf8, false),
- Field::new("lat", DataType::Float64, false),
- Field::new("lng", DataType::Float64, false),
- ]);
-
- let file = File::open("test/data/uk_cities.csv").unwrap();
-
- let mut csv =
- csv::Reader::new(file, Arc::new(schema), false, None, 1024, None, None);
- let _batch = csv.next().unwrap().unwrap();
- #[cfg(feature = "prettyprint")]
- {
- print_batches(&[_batch]).unwrap();
- }
- }
-}
diff --git a/arrow/examples/read_csv_infer_schema.rs b/arrow/examples/read_csv_infer_schema.rs
deleted file mode 100644
index 11f8cfb..0000000
--- a/arrow/examples/read_csv_infer_schema.rs
+++ /dev/null
@@ -1,40 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-extern crate arrow;
-
-#[cfg(feature = "csv")]
-use arrow::csv;
-#[cfg(feature = "prettyprint")]
-use arrow::util::pretty::print_batches;
-use std::fs::File;
-
-fn main() {
- #[cfg(feature = "csv")]
- {
- let file = File::open("test/data/uk_cities_with_headers.csv").unwrap();
- let builder = csv::ReaderBuilder::new()
- .has_header(true)
- .infer_schema(Some(100));
- let mut csv = builder.build(file).unwrap();
- let _batch = csv.next().unwrap().unwrap();
- #[cfg(feature = "prettyprint")]
- {
- print_batches(&[_batch]).unwrap();
- }
- }
-}
diff --git a/arrow/examples/tensor_builder.rs b/arrow/examples/tensor_builder.rs
deleted file mode 100644
index 1ef5392..0000000
--- a/arrow/examples/tensor_builder.rs
+++ /dev/null
@@ -1,67 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-///! Tensor builder example
-extern crate arrow;
-
-use arrow::array::*; //{Int32BufferBuilder, Float32BufferBuilder};
-use arrow::buffer::Buffer;
-use arrow::datatypes::ToByteSlice;
-use arrow::error::Result;
-use arrow::tensor::{Float32Tensor, Int32Tensor};
-
-fn main() -> Result<()> {
- // Building a tensor using the buffer builder for Int32
- // The buffer builder will pad the appended numbers
- // to match the required size for each buffer
- let mut builder = Int32BufferBuilder::new(16);
- for i in 0..16 {
- builder.append(i);
- }
- let buf = builder.finish();
-
- // When building a tensor the buffer and shape are required
- // The new function will estimate the expected stride for the
- // storage data
- let tensor = Int32Tensor::try_new(buf, Some(vec![2, 8]), None, None)?;
- println!("Int32 Tensor");
- println!("{:?}", tensor);
-
- // Creating a tensor using float type buffer builder
- let mut builder = Float32BufferBuilder::new(4);
- builder.append(1.0);
- builder.append(2.0);
- builder.append(3.0);
- builder.append(4.0);
- let buf = builder.finish();
-
- // When building the tensor the buffer and shape are necessary
- // The new function will estimate the expected stride for the
- // storage data
- let tensor = Float32Tensor::try_new(buf, Some(vec![2, 2]), None, None)?;
- println!("\nFloat32 Tensor");
- println!("{:?}", tensor);
-
- // In order to build a tensor from an array the function to_byte_slice add the
- // required padding to the elements in the array.
- let buf = Buffer::from(&[0, 1, 2, 3, 4, 5, 6, 7, 9, 10].to_byte_slice());
- let tensor = Int32Tensor::try_new(buf, Some(vec![2, 5]), None, None)?;
- println!("\nInt32 Tensor");
- println!("{:?}", tensor);
-
- Ok(())
-}
diff --git a/arrow/format-0ed34c83.patch b/arrow/format-0ed34c83.patch
deleted file mode 100644
index 5da0a0c..0000000
--- a/arrow/format-0ed34c83.patch
+++ /dev/null
@@ -1,220 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-diff --git a/format/Message.fbs b/format/Message.fbs
-index 1a7e0dfff..f1c18d765 100644
---- a/format/Message.fbs
-+++ b/format/Message.fbs
-@@ -28,7 +28,7 @@ namespace org.apache.arrow.flatbuf;
- /// Metadata about a field at some level of a nested type tree (but not
- /// its children).
- ///
--/// For example, a List<Int16> with values [[1, 2, 3], null, [4], [5, 6], null]
-+/// For example, a List<Int16> 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 {
-diff --git a/format/Schema.fbs b/format/Schema.fbs
-index 3b37e5d85..3b00dd478 100644
---- a/format/Schema.fbs
-+++ b/format/Schema.fbs
-@@ -110,10 +110,11 @@ table FixedSizeList {
- /// 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
-@@ -129,7 +130,7 @@ 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
-+/// 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.
-diff --git a/format/SparseTensor.fbs b/format/SparseTensor.fbs
-index 3fe8a7582..a6fd2f9e7 100644
---- a/format/SparseTensor.fbs
-+++ b/format/SparseTensor.fbs
-@@ -37,21 +37,21 @@ namespace org.apache.arrow.flatbuf;
- ///
- /// 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.
-@@ -86,26 +86,27 @@ table SparseMatrixIndexCSX {
-
- /// 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 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
-@@ -116,9 +117,9 @@ table SparseMatrixIndexCSX {
- /// 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);
- }
-@@ -126,7 +127,7 @@ table SparseMatrixIndexCSX {
- /// 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
-+ /// 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
-@@ -135,7 +136,7 @@ table SparseTensorIndexCSF {
- ///
- /// 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
-@@ -144,9 +145,9 @@ table SparseTensorIndexCSF {
- /// 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
-@@ -154,24 +155,24 @@ table SparseTensorIndexCSF {
- /// 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.
-+ /// 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
-@@ -180,22 +181,22 @@ table SparseTensorIndexCSF {
- /// 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);
- }
-
diff --git a/arrow/regen.sh b/arrow/regen.sh
deleted file mode 100755
index 9d384b6..0000000
--- a/arrow/regen.sh
+++ /dev/null
@@ -1,157 +0,0 @@
-#!/bin/bash -e
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-
-DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" >/dev/null 2>&1 && pwd )"
-
-# Change to the toplevel Rust directory
-pushd $DIR/../../
-
-echo "Build flatc from source ..."
-
-FB_URL="https://github.com/google/flatbuffers"
-# https://github.com/google/flatbuffers/pull/6393
-FB_COMMIT="408cf5802415e1dea65fef7489a6c2f3740fb381"
-FB_DIR="rust/arrow/.flatbuffers"
-FLATC="$FB_DIR/bazel-bin/flatc"
-
-if [ -z $(which bazel) ]; then
- echo "bazel is required to build flatc"
- exit 1
-fi
-
-echo "Bazel version: $(bazel version | head -1 | awk -F':' '{print $2}')"
-
-if [ ! -e $FB_DIR ]; then
- echo "git clone $FB_URL ..."
- git clone -b master --no-tag --depth 1 $FB_URL $FB_DIR
-else
- echo "git pull $FB_URL ..."
- git -C $FB_DIR pull
-fi
-
-echo "hard reset to $FB_COMMIT"
-git -C $FB_DIR reset --hard $FB_COMMIT
-
-pushd $FB_DIR
-echo "run: bazel build :flatc ..."
-bazel build :flatc
-popd
-
-FB_PATCH="rust/arrow/format-0ed34c83.patch"
-echo "Patch flatbuffer files with ${FB_PATCH} for cargo doc"
-echo "NOTE: the patch MAY need update in case of changes in format/*.fbs"
-git apply --check ${FB_PATCH} && git apply ${FB_PATCH}
-
-# Execute the code generation:
-$FLATC --filename-suffix "" --rust -o rust/arrow/src/ipc/gen/ format/*.fbs
-
-# Reset changes to format/
-git checkout -- format
-
-# Now the files are wrongly named so we have to change that.
-popd
-pushd $DIR/src/ipc/gen
-
-PREFIX=$(cat <<'HEREDOC'
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-#![allow(dead_code)]
-#![allow(unused_imports)]
-
-use std::{cmp::Ordering, mem};
-use flatbuffers::EndianScalar;
-
-HEREDOC
-)
-
-SCHEMA_IMPORT="\nuse crate::ipc::gen::Schema::*;"
-SPARSE_TENSOR_IMPORT="\nuse crate::ipc::gen::SparseTensor::*;"
-TENSOR_IMPORT="\nuse crate::ipc::gen::Tensor::*;"
-
-# For flatbuffer(1.12.0+), remove: use crate::${name}::\*;
-names=("File" "Message" "Schema" "SparseTensor" "Tensor")
-
-# Remove all generated lines we don't need
-for f in `ls *.rs`; do
- if [[ $f == "mod.rs" ]]; then
- continue
- fi
-
- echo "Modifying: $f"
- sed -i '' '/extern crate flatbuffers;/d' $f
- sed -i '' '/use self::flatbuffers::EndianScalar;/d' $f
- sed -i '' '/\#\[allow(unused_imports, dead_code)\]/d' $f
- sed -i '' '/pub mod org {/d' $f
- sed -i '' '/pub mod apache {/d' $f
- sed -i '' '/pub mod arrow {/d' $f
- sed -i '' '/pub mod flatbuf {/d' $f
- sed -i '' '/} \/\/ pub mod flatbuf/d' $f
- sed -i '' '/} \/\/ pub mod arrow/d' $f
- sed -i '' '/} \/\/ pub mod apache/d' $f
- sed -i '' '/} \/\/ pub mod org/d' $f
- sed -i '' '/use std::mem;/d' $f
- sed -i '' '/use std::cmp::Ordering;/d' $f
-
- # required by flatc 1.12.0+
- sed -i '' "/\#\!\[allow(unused_imports, dead_code)\]/d" $f
- for name in ${names[@]}; do
- sed -i '' "/use crate::${name}::\*;/d" $f
- sed -i '' "s/use self::flatbuffers::Verifiable;/use flatbuffers::Verifiable;/g" $f
- done
-
- # Replace all occurrences of "type__" with "type_", "TYPE__" with "TYPE_".
- sed -i '' 's/type__/type_/g' $f
- sed -i '' 's/TYPE__/TYPE_/g' $f
-
- # Some files need prefixes
- if [[ $f == "File.rs" ]]; then
- # Now prefix the file with the static contents
- echo -e "${PREFIX}" "${SCHEMA_IMPORT}" | cat - $f > temp && mv temp $f
- elif [[ $f == "Message.rs" ]]; then
- echo -e "${PREFIX}" "${SCHEMA_IMPORT}" "${SPARSE_TENSOR_IMPORT}" "${TENSOR_IMPORT}" | cat - $f > temp && mv temp $f
- elif [[ $f == "SparseTensor.rs" ]]; then
- echo -e "${PREFIX}" "${SCHEMA_IMPORT}" "${TENSOR_IMPORT}" | cat - $f > temp && mv temp $f
- elif [[ $f == "Tensor.rs" ]]; then
- echo -e "${PREFIX}" "${SCHEMA_IMPORT}" | cat - $f > temp && mv temp $f
- else
- echo "${PREFIX}" | cat - $f > temp && mv temp $f
- fi
-done
-
-# Return back to base directory
-popd
-cargo +stable fmt -- src/ipc/gen/*
-
-echo "DONE!"
-echo "Please run 'cargo doc' and 'cargo test' with nightly and stable, "
-echo "and fix possible errors or warnings!"
diff --git a/arrow/src/alloc/alignment.rs b/arrow/src/alloc/alignment.rs
deleted file mode 100644
index dbf4602..0000000
--- a/arrow/src/alloc/alignment.rs
+++ /dev/null
@@ -1,119 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-// NOTE: Below code is written for spatial/temporal prefetcher optimizations. Memory allocation
-// should align well with usage pattern of cache access and block sizes on layers of storage levels from
-// registers to non-volatile memory. These alignments are all cache aware alignments incorporated
-// from [cuneiform](https://crates.io/crates/cuneiform) crate. This approach mimicks Intel TBB's
-// cache_aligned_allocator which exploits cache locality and minimizes prefetch signals
-// resulting in less round trip time between the layers of storage.
-// For further info: https://software.intel.com/en-us/node/506094
-
-// 32-bit architecture and things other than netburst microarchitecture are using 64 bytes.
-/// Cache and allocation multiple alignment size
-#[cfg(target_arch = "x86")]
-pub const ALIGNMENT: usize = 1 << 6;
-
-// Intel x86_64:
-// L2D streamer from L1:
-// Loads data or instructions from memory to the second-level cache. To use the streamer,
-// organize the data or instructions in blocks of 128 bytes, aligned on 128 bytes.
-// - https://www.intel.com/content/dam/www/public/us/en/documents/manuals/64-ia-32-architectures-optimization-manual.pdf
-/// Cache and allocation multiple alignment size
-#[cfg(target_arch = "x86_64")]
-pub const ALIGNMENT: usize = 1 << 7;
-
-// 24Kc:
-// Data Line Size
-// - https://s3-eu-west-1.amazonaws.com/downloads-mips/documents/MD00346-2B-24K-DTS-04.00.pdf
-// - https://gitlab.e.foundation/e/devices/samsung/n7100/stable_android_kernel_samsung_smdk4412/commit/2dbac10263b2f3c561de68b4c369bc679352ccee
-/// Cache and allocation multiple alignment size
-#[cfg(target_arch = "mips")]
-pub const ALIGNMENT: usize = 1 << 5;
-/// Cache and allocation multiple alignment size
-#[cfg(target_arch = "mips64")]
-pub const ALIGNMENT: usize = 1 << 5;
-
-// Defaults for powerpc
-/// Cache and allocation multiple alignment size
-#[cfg(target_arch = "powerpc")]
-pub const ALIGNMENT: usize = 1 << 5;
-
-// Defaults for the ppc 64
-/// Cache and allocation multiple alignment size
-#[cfg(target_arch = "powerpc64")]
-pub const ALIGNMENT: usize = 1 << 6;
-
-// e.g.: sifive
-// - https://github.com/torvalds/linux/blob/master/Documentation/devicetree/bindings/riscv/sifive-l2-cache.txt#L41
-// in general all of them are the same.
-/// Cache and allocation multiple alignment size
-#[cfg(target_arch = "riscv")]
-pub const ALIGNMENT: usize = 1 << 6;
-
-// This size is same across all hardware for this architecture.
-// - https://docs.huihoo.com/doxygen/linux/kernel/3.7/arch_2s390_2include_2asm_2cache_8h.html
-/// Cache and allocation multiple alignment size
-#[cfg(target_arch = "s390x")]
-pub const ALIGNMENT: usize = 1 << 8;
-
-// This size is same across all hardware for this architecture.
-// - https://docs.huihoo.com/doxygen/linux/kernel/3.7/arch_2sparc_2include_2asm_2cache_8h.html#a9400cc2ba37e33279bdbc510a6311fb4
-/// Cache and allocation multiple alignment size
-#[cfg(target_arch = "sparc")]
-pub const ALIGNMENT: usize = 1 << 5;
-/// Cache and allocation multiple alignment size
-#[cfg(target_arch = "sparc64")]
-pub const ALIGNMENT: usize = 1 << 6;
-
-// On ARM cache line sizes are fixed. both v6 and v7.
-// Need to add board specific or platform specific things later.
-/// Cache and allocation multiple alignment size
-#[cfg(target_arch = "thumbv6")]
-pub const ALIGNMENT: usize = 1 << 5;
-/// Cache and allocation multiple alignment size
-#[cfg(target_arch = "thumbv7")]
-pub const ALIGNMENT: usize = 1 << 5;
-
-// Operating Systems cache size determines this.
-// Currently no way to determine this without runtime inference.
-/// Cache and allocation multiple alignment size
-#[cfg(target_arch = "wasm32")]
-pub const ALIGNMENT: usize = 1 << 6;
-
-// Same as v6 and v7.
-// List goes like that:
-// Cortex A, M, R, ARM v7, v7-M, Krait and NeoverseN uses this size.
-/// Cache and allocation multiple alignment size
-#[cfg(target_arch = "arm")]
-pub const ALIGNMENT: usize = 1 << 5;
-
-// Combined from 4 sectors. Volta says 128.
-// Prevent chunk optimizations better to go to the default size.
-// If you have smaller data with less padded functionality then use 32 with force option.
-// - https://devtalk.nvidia.com/default/topic/803600/variable-cache-line-width-/
-/// Cache and allocation multiple alignment size
-#[cfg(target_arch = "nvptx")]
-pub const ALIGNMENT: usize = 1 << 7;
-/// Cache and allocation multiple alignment size
-#[cfg(target_arch = "nvptx64")]
-pub const ALIGNMENT: usize = 1 << 7;
-
-// This size is same across all hardware for this architecture.
-/// Cache and allocation multiple alignment size
-#[cfg(target_arch = "aarch64")]
-pub const ALIGNMENT: usize = 1 << 6;
diff --git a/arrow/src/alloc/mod.rs b/arrow/src/alloc/mod.rs
deleted file mode 100644
index a225d32..0000000
--- a/arrow/src/alloc/mod.rs
+++ /dev/null
@@ -1,136 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-//! Defines memory-related functions, such as allocate/deallocate/reallocate memory
-//! regions, cache and allocation alignments.
-
-use std::mem::size_of;
-use std::ptr::NonNull;
-use std::{
- alloc::{handle_alloc_error, Layout},
- sync::atomic::AtomicIsize,
-};
-
-mod alignment;
-mod types;
-
-pub use alignment::ALIGNMENT;
-pub use types::NativeType;
-
-// If this number is not zero after all objects have been `drop`, there is a memory leak
-pub static mut ALLOCATIONS: AtomicIsize = AtomicIsize::new(0);
-
-#[inline]
-unsafe fn null_pointer<T: NativeType>() -> NonNull<T> {
- NonNull::new_unchecked(ALIGNMENT as *mut T)
-}
-
-/// Allocates a cache-aligned memory region of `size` bytes with uninitialized values.
-/// This is more performant than using [allocate_aligned_zeroed] when all bytes will have
-/// an unknown or non-zero value and is semantically similar to `malloc`.
-pub fn allocate_aligned<T: NativeType>(size: usize) -> NonNull<T> {
- unsafe {
- if size == 0 {
- null_pointer()
- } else {
- let size = size * size_of::<T>();
- ALLOCATIONS.fetch_add(size as isize, std::sync::atomic::Ordering::SeqCst);
-
- let layout = Layout::from_size_align_unchecked(size, ALIGNMENT);
- let raw_ptr = std::alloc::alloc(layout) as *mut T;
- NonNull::new(raw_ptr).unwrap_or_else(|| handle_alloc_error(layout))
- }
- }
-}
-
-/// Allocates a cache-aligned memory region of `size` bytes with `0` on all of them.
-/// This is more performant than using [allocate_aligned] and setting all bytes to zero
-/// and is semantically similar to `calloc`.
-pub fn allocate_aligned_zeroed<T: NativeType>(size: usize) -> NonNull<T> {
- unsafe {
- if size == 0 {
- null_pointer()
- } else {
- let size = size * size_of::<T>();
- ALLOCATIONS.fetch_add(size as isize, std::sync::atomic::Ordering::SeqCst);
-
- let layout = Layout::from_size_align_unchecked(size, ALIGNMENT);
- let raw_ptr = std::alloc::alloc_zeroed(layout) as *mut T;
- NonNull::new(raw_ptr).unwrap_or_else(|| handle_alloc_error(layout))
- }
- }
-}
-
-/// # Safety
-///
-/// This function is unsafe because undefined behavior can result if the caller does not ensure all
-/// of the following:
-///
-/// * ptr must denote a block of memory currently allocated via this allocator,
-///
-/// * size must be the same size that was used to allocate that block of memory,
-pub unsafe fn free_aligned<T: NativeType>(ptr: NonNull<T>, size: usize) {
- if ptr != null_pointer() {
- let size = size * size_of::<T>();
- ALLOCATIONS.fetch_sub(size as isize, std::sync::atomic::Ordering::SeqCst);
- std::alloc::dealloc(
- ptr.as_ptr() as *mut u8,
- Layout::from_size_align_unchecked(size, ALIGNMENT),
- );
- }
-}
-
-/// # Safety
-///
-/// This function is unsafe because undefined behavior can result if the caller does not ensure all
-/// of the following:
-///
-/// * ptr must be currently allocated via this allocator,
-///
-/// * new_size must be greater than zero.
-///
-/// * new_size, when rounded up to the nearest multiple of [ALIGNMENT], must not overflow (i.e.,
-/// the rounded value must be less than usize::MAX).
-pub unsafe fn reallocate<T: NativeType>(
- ptr: NonNull<T>,
- old_size: usize,
- new_size: usize,
-) -> NonNull<T> {
- let old_size = old_size * size_of::<T>();
- let new_size = new_size * size_of::<T>();
- if ptr == null_pointer() {
- return allocate_aligned(new_size);
- }
-
- if new_size == 0 {
- free_aligned(ptr, old_size);
- return null_pointer();
- }
-
- ALLOCATIONS.fetch_add(
- new_size as isize - old_size as isize,
- std::sync::atomic::Ordering::SeqCst,
- );
- let raw_ptr = std::alloc::realloc(
- ptr.as_ptr() as *mut u8,
- Layout::from_size_align_unchecked(old_size, ALIGNMENT),
- new_size,
- ) as *mut T;
- NonNull::new(raw_ptr).unwrap_or_else(|| {
- handle_alloc_error(Layout::from_size_align_unchecked(new_size, ALIGNMENT))
- })
-}
diff --git a/arrow/src/alloc/types.rs b/arrow/src/alloc/types.rs
deleted file mode 100644
index c1f0ef9..0000000
--- a/arrow/src/alloc/types.rs
+++ /dev/null
@@ -1,71 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-use crate::datatypes::DataType;
-
-/// A type that Rust's custom allocator knows how to allocate and deallocate.
-/// This is implemented for all Arrow's physical types whose in-memory representation
-/// matches Rust's physical types. Consider this trait sealed.
-/// # Safety
-/// Do not implement this trait.
-pub unsafe trait NativeType:
- Sized + Copy + std::fmt::Debug + std::fmt::Display + PartialEq + Default + Sized + 'static
-{
- type Bytes: AsRef<[u8]>;
-
- /// Whether a DataType is a valid type for this physical representation.
- fn is_valid(data_type: &DataType) -> bool;
-
- /// How this type represents itself as bytes in little endianess.
- /// This is used for IPC, where data is communicated with a specific endianess.
- fn to_le_bytes(&self) -> Self::Bytes;
-}
-
-macro_rules! create_native {
- ($native_ty:ty,$($impl_pattern:pat)|+) => {
- unsafe impl NativeType for $native_ty {
- type Bytes = [u8; std::mem::size_of::<Self>()];
-
- #[inline]
- fn to_le_bytes(&self) -> Self::Bytes {
- Self::to_le_bytes(*self)
- }
-
- #[inline]
- fn is_valid(data_type: &DataType) -> bool {
- matches!(data_type, $($impl_pattern)|+)
- }
- }
- };
-}
-
-create_native!(u8, DataType::UInt8);
-create_native!(u16, DataType::UInt16);
-create_native!(u32, DataType::UInt32);
-create_native!(u64, DataType::UInt64);
-create_native!(i8, DataType::Int8);
-create_native!(i16, DataType::Int16);
-create_native!(
- i32,
- DataType::Int32 | DataType::Date32 | DataType::Time32(_)
-);
-create_native!(
- i64,
- DataType::Int64 | DataType::Date64 | DataType::Time64(_) | DataType::Timestamp(_, _)
-);
-create_native!(f32, DataType::Float32);
-create_native!(f64, DataType::Float64);
diff --git a/arrow/src/arch/avx512.rs b/arrow/src/arch/avx512.rs
deleted file mode 100644
index 264532f..0000000
--- a/arrow/src/arch/avx512.rs
+++ /dev/null
@@ -1,73 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-pub(crate) const AVX512_U8X64_LANES: usize = 64;
-
-#[target_feature(enable = "avx512f")]
-pub(crate) unsafe fn avx512_bin_and(left: &[u8], right: &[u8], res: &mut [u8]) {
- use core::arch::x86_64::{__m512i, _mm512_and_si512, _mm512_loadu_epi64};
-
- let l: __m512i = _mm512_loadu_epi64(left.as_ptr() as *const _);
- let r: __m512i = _mm512_loadu_epi64(right.as_ptr() as *const _);
- let f = _mm512_and_si512(l, r);
- let s = &f as *const __m512i as *const u8;
- let d = res.get_unchecked_mut(0) as *mut _ as *mut u8;
- std::ptr::copy_nonoverlapping(s, d, std::mem::size_of::<__m512i>());
-}
-
-#[target_feature(enable = "avx512f")]
-pub(crate) unsafe fn avx512_bin_or(left: &[u8], right: &[u8], res: &mut [u8]) {
- use core::arch::x86_64::{__m512i, _mm512_loadu_epi64, _mm512_or_si512};
-
- let l: __m512i = _mm512_loadu_epi64(left.as_ptr() as *const _);
- let r: __m512i = _mm512_loadu_epi64(right.as_ptr() as *const _);
- let f = _mm512_or_si512(l, r);
- let s = &f as *const __m512i as *const u8;
- let d = res.get_unchecked_mut(0) as *mut _ as *mut u8;
- std::ptr::copy_nonoverlapping(s, d, std::mem::size_of::<__m512i>());
-}
-
-#[cfg(test)]
-mod tests {
- use super::*;
-
- #[test]
- fn test_bitwise_and_avx512() {
- let buf1 = [0b00110011u8; 64];
- let buf2 = [0b11110000u8; 64];
- let mut buf3 = [0b00000000; 64];
- unsafe {
- avx512_bin_and(&buf1, &buf2, &mut buf3);
- };
- for i in buf3.iter() {
- assert_eq!(&0b00110000u8, i);
- }
- }
-
- #[test]
- fn test_bitwise_or_avx512() {
- let buf1 = [0b00010011u8; 64];
- let buf2 = [0b11100000u8; 64];
- let mut buf3 = [0b00000000; 64];
- unsafe {
- avx512_bin_or(&buf1, &buf2, &mut buf3);
- };
- for i in buf3.iter() {
- assert_eq!(&0b11110011u8, i);
- }
- }
-}
diff --git a/arrow/src/arch/mod.rs b/arrow/src/arch/mod.rs
deleted file mode 100644
index 56d8f4c..0000000
--- a/arrow/src/arch/mod.rs
+++ /dev/null
@@ -1,22 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-///
-/// Arch module contains architecture specific code.
-/// Be aware that not all machines have these specific operations available.
-#[cfg(all(target_arch = "x86_64", feature = "avx512"))]
-pub(crate) mod avx512;
diff --git a/arrow/src/array/array.rs b/arrow/src/array/array.rs
deleted file mode 100644
index de87c3d..0000000
--- a/arrow/src/array/array.rs
+++ /dev/null
@@ -1,664 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-use std::fmt;
-use std::sync::Arc;
-use std::{any::Any, convert::TryFrom};
-
-use super::*;
-use crate::array::equal_json::JsonEqual;
-use crate::buffer::{Buffer, MutableBuffer};
-use crate::error::Result;
-use crate::ffi;
-
-/// Trait for dealing with different types of array at runtime when the type of the
-/// array is not known in advance.
-pub trait Array: fmt::Debug + Send + Sync + JsonEqual {
- /// Returns the array as [`Any`](std::any::Any) so that it can be
- /// downcasted to a specific implementation.
- ///
- /// # Example:
- ///
- /// ```
- /// use std::sync::Arc;
- /// use arrow::array::Int32Array;
- /// use arrow::datatypes::{Schema, Field, DataType};
- /// use arrow::record_batch::RecordBatch;
- ///
- /// # fn main() -> arrow::error::Result<()> {
- /// let id = Int32Array::from(vec![1, 2, 3, 4, 5]);
- /// let batch = RecordBatch::try_new(
- /// Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)])),
- /// vec![Arc::new(id)]
- /// )?;
- ///
- /// let int32array = batch
- /// .column(0)
- /// .as_any()
- /// .downcast_ref::<Int32Array>()
- /// .expect("Failed to downcast");
- /// # Ok(())
- /// # }
- /// ```
- fn as_any(&self) -> &Any;
-
- /// Returns a reference to the underlying data of this array.
- fn data(&self) -> &ArrayData;
-
- /// Returns a reference-counted pointer to the underlying data of this array.
- fn data_ref(&self) -> &ArrayData {
- self.data()
- }
-
- /// Returns a reference to the [`DataType`](crate::datatypes::DataType) of this array.
- ///
- /// # Example:
- ///
- /// ```
- /// use arrow::datatypes::DataType;
- /// use arrow::array::{Array, Int32Array};
- ///
- /// let array = Int32Array::from(vec![1, 2, 3, 4, 5]);
- ///
- /// assert_eq!(*array.data_type(), DataType::Int32);
- /// ```
- fn data_type(&self) -> &DataType {
- self.data_ref().data_type()
- }
-
- /// Returns a zero-copy slice of this array with the indicated offset and length.
- ///
- /// # Example:
- ///
- /// ```
- /// use arrow::array::{Array, Int32Array};
- ///
- /// let array = Int32Array::from(vec![1, 2, 3, 4, 5]);
- /// // Make slice over the values [2, 3, 4]
- /// let array_slice = array.slice(1, 3);
- ///
- /// assert_eq!(array_slice.as_ref(), &Int32Array::from(vec![2, 3, 4]));
- /// ```
- fn slice(&self, offset: usize, length: usize) -> ArrayRef {
- make_array(self.data_ref().slice(offset, length))
- }
-
- /// Returns the length (i.e., number of elements) of this array.
- ///
- /// # Example:
- ///
- /// ```
- /// use arrow::array::{Array, Int32Array};
- ///
- /// let array = Int32Array::from(vec![1, 2, 3, 4, 5]);
- ///
- /// assert_eq!(array.len(), 5);
- /// ```
- fn len(&self) -> usize {
- self.data_ref().len()
- }
-
- /// Returns whether this array is empty.
- ///
- /// # Example:
- ///
- /// ```
- /// use arrow::array::{Array, Int32Array};
- ///
- /// let array = Int32Array::from(vec![1, 2, 3, 4, 5]);
- ///
- /// assert_eq!(array.is_empty(), false);
- /// ```
- fn is_empty(&self) -> bool {
- self.data_ref().is_empty()
- }
-
- /// Returns the offset into the underlying data used by this array(-slice).
- /// Note that the underlying data can be shared by many arrays.
- /// This defaults to `0`.
- ///
- /// # Example:
- ///
- /// ```
- /// use arrow::array::{Array, Int32Array};
- ///
- /// let array = Int32Array::from(vec![1, 2, 3, 4, 5]);
- /// // Make slice over the values [2, 3, 4]
- /// let array_slice = array.slice(1, 3);
- ///
- /// assert_eq!(array.offset(), 0);
- /// assert_eq!(array_slice.offset(), 1);
- /// ```
- fn offset(&self) -> usize {
- self.data_ref().offset()
- }
-
- /// Returns whether the element at `index` is null.
- /// When using this function on a slice, the index is relative to the slice.
- ///
- /// # Example:
- ///
- /// ```
- /// use arrow::array::{Array, Int32Array};
- ///
- /// let array = Int32Array::from(vec![Some(1), None]);
- ///
- /// assert_eq!(array.is_null(0), false);
- /// assert_eq!(array.is_null(1), true);
- /// ```
- fn is_null(&self, index: usize) -> bool {
- self.data_ref().is_null(index)
- }
-
- /// Returns whether the element at `index` is not null.
- /// When using this function on a slice, the index is relative to the slice.
- ///
- /// # Example:
- ///
- /// ```
- /// use arrow::array::{Array, Int32Array};
- ///
- /// let array = Int32Array::from(vec![Some(1), None]);
- ///
- /// assert_eq!(array.is_valid(0), true);
- /// assert_eq!(array.is_valid(1), false);
- /// ```
- fn is_valid(&self, index: usize) -> bool {
- self.data_ref().is_valid(index)
- }
-
- /// Returns the total number of null values in this array.
- ///
- /// # Example:
- ///
- /// ```
- /// use arrow::array::{Array, Int32Array};
- ///
- /// // Construct an array with values [1, NULL, NULL]
- /// let array = Int32Array::from(vec![Some(1), None, None]);
- ///
- /// assert_eq!(array.null_count(), 2);
- /// ```
- fn null_count(&self) -> usize {
- self.data_ref().null_count()
- }
-
- /// Returns the total number of bytes of memory occupied by the buffers owned by this array.
- fn get_buffer_memory_size(&self) -> usize;
-
- /// Returns the total number of bytes of memory occupied physically by this array.
- fn get_array_memory_size(&self) -> usize;
-
- /// returns two pointers that represent this array in the C Data Interface (FFI)
- fn to_raw(
- &self,
- ) -> Result<(*const ffi::FFI_ArrowArray, *const ffi::FFI_ArrowSchema)> {
- let data = self.data().clone();
- let array = ffi::ArrowArray::try_from(data)?;
- Ok(ffi::ArrowArray::into_raw(array))
- }
-}
-
-/// A reference-counted reference to a generic `Array`.
-pub type ArrayRef = Arc<Array>;
-
-/// Constructs an array using the input `data`.
-/// Returns a reference-counted `Array` instance.
-pub fn make_array(data: ArrayData) -> ArrayRef {
- match data.data_type() {
- DataType::Boolean => Arc::new(BooleanArray::from(data)) as ArrayRef,
- DataType::Int8 => Arc::new(Int8Array::from(data)) as ArrayRef,
- DataType::Int16 => Arc::new(Int16Array::from(data)) as ArrayRef,
- DataType::Int32 => Arc::new(Int32Array::from(data)) as ArrayRef,
- DataType::Int64 => Arc::new(Int64Array::from(data)) as ArrayRef,
- DataType::UInt8 => Arc::new(UInt8Array::from(data)) as ArrayRef,
- DataType::UInt16 => Arc::new(UInt16Array::from(data)) as ArrayRef,
- DataType::UInt32 => Arc::new(UInt32Array::from(data)) as ArrayRef,
- DataType::UInt64 => Arc::new(UInt64Array::from(data)) as ArrayRef,
- DataType::Float16 => panic!("Float16 datatype not supported"),
- DataType::Float32 => Arc::new(Float32Array::from(data)) as ArrayRef,
- DataType::Float64 => Arc::new(Float64Array::from(data)) as ArrayRef,
- DataType::Date32 => Arc::new(Date32Array::from(data)) as ArrayRef,
- DataType::Date64 => Arc::new(Date64Array::from(data)) as ArrayRef,
- DataType::Time32(TimeUnit::Second) => {
- Arc::new(Time32SecondArray::from(data)) as ArrayRef
- }
- DataType::Time32(TimeUnit::Millisecond) => {
- Arc::new(Time32MillisecondArray::from(data)) as ArrayRef
- }
- DataType::Time64(TimeUnit::Microsecond) => {
- Arc::new(Time64MicrosecondArray::from(data)) as ArrayRef
- }
- DataType::Time64(TimeUnit::Nanosecond) => {
- Arc::new(Time64NanosecondArray::from(data)) as ArrayRef
- }
- DataType::Timestamp(TimeUnit::Second, _) => {
- Arc::new(TimestampSecondArray::from(data)) as ArrayRef
- }
- DataType::Timestamp(TimeUnit::Millisecond, _) => {
- Arc::new(TimestampMillisecondArray::from(data)) as ArrayRef
- }
- DataType::Timestamp(TimeUnit::Microsecond, _) => {
- Arc::new(TimestampMicrosecondArray::from(data)) as ArrayRef
- }
- DataType::Timestamp(TimeUnit::Nanosecond, _) => {
- Arc::new(TimestampNanosecondArray::from(data)) as ArrayRef
- }
- DataType::Interval(IntervalUnit::YearMonth) => {
- Arc::new(IntervalYearMonthArray::from(data)) as ArrayRef
- }
- DataType::Interval(IntervalUnit::DayTime) => {
- Arc::new(IntervalDayTimeArray::from(data)) as ArrayRef
- }
- DataType::Duration(TimeUnit::Second) => {
- Arc::new(DurationSecondArray::from(data)) as ArrayRef
- }
- DataType::Duration(TimeUnit::Millisecond) => {
- Arc::new(DurationMillisecondArray::from(data)) as ArrayRef
- }
- DataType::Duration(TimeUnit::Microsecond) => {
- Arc::new(DurationMicrosecondArray::from(data)) as ArrayRef
- }
- DataType::Duration(TimeUnit::Nanosecond) => {
- Arc::new(DurationNanosecondArray::from(data)) as ArrayRef
- }
- DataType::Binary => Arc::new(BinaryArray::from(data)) as ArrayRef,
- DataType::LargeBinary => Arc::new(LargeBinaryArray::from(data)) as ArrayRef,
- DataType::FixedSizeBinary(_) => {
- Arc::new(FixedSizeBinaryArray::from(data)) as ArrayRef
- }
- DataType::Utf8 => Arc::new(StringArray::from(data)) as ArrayRef,
- DataType::LargeUtf8 => Arc::new(LargeStringArray::from(data)) as ArrayRef,
- DataType::List(_) => Arc::new(ListArray::from(data)) as ArrayRef,
- DataType::LargeList(_) => Arc::new(LargeListArray::from(data)) as ArrayRef,
- DataType::Struct(_) => Arc::new(StructArray::from(data)) as ArrayRef,
- DataType::Union(_) => Arc::new(UnionArray::from(data)) as ArrayRef,
- DataType::FixedSizeList(_, _) => {
- Arc::new(FixedSizeListArray::from(data)) as ArrayRef
- }
- DataType::Dictionary(ref key_type, _) => match key_type.as_ref() {
- DataType::Int8 => {
- Arc::new(DictionaryArray::<Int8Type>::from(data)) as ArrayRef
- }
- DataType::Int16 => {
- Arc::new(DictionaryArray::<Int16Type>::from(data)) as ArrayRef
- }
- DataType::Int32 => {
- Arc::new(DictionaryArray::<Int32Type>::from(data)) as ArrayRef
- }
- DataType::Int64 => {
- Arc::new(DictionaryArray::<Int64Type>::from(data)) as ArrayRef
- }
- DataType::UInt8 => {
- Arc::new(DictionaryArray::<UInt8Type>::from(data)) as ArrayRef
- }
- DataType::UInt16 => {
- Arc::new(DictionaryArray::<UInt16Type>::from(data)) as ArrayRef
- }
- DataType::UInt32 => {
- Arc::new(DictionaryArray::<UInt32Type>::from(data)) as ArrayRef
- }
- DataType::UInt64 => {
- Arc::new(DictionaryArray::<UInt64Type>::from(data)) as ArrayRef
- }
- dt => panic!("Unexpected dictionary key type {:?}", dt),
- },
- DataType::Null => Arc::new(NullArray::from(data)) as ArrayRef,
- DataType::Decimal(_, _) => Arc::new(DecimalArray::from(data)) as ArrayRef,
- dt => panic!("Unexpected data type {:?}", dt),
- }
-}
-
-/// Creates a new empty array
-///
-/// ```
-/// use std::sync::Arc;
-/// use arrow::datatypes::DataType;
-/// use arrow::array::{ArrayRef, Int32Array, new_empty_array};
-///
-/// let empty_array = new_empty_array(&DataType::Int32);
-/// let array: ArrayRef = Arc::new(Int32Array::from(vec![] as Vec<i32>));
-///
-/// assert_eq!(&array, &empty_array);
-/// ```
-pub fn new_empty_array(data_type: &DataType) -> ArrayRef {
- let data = ArrayData::new_empty(data_type);
- make_array(data)
-}
-
-/// Creates a new array of `data_type` of length `length` filled
-/// entirely of `NULL` values
-///
-/// ```
-/// use std::sync::Arc;
-/// use arrow::datatypes::DataType;
-/// use arrow::array::{ArrayRef, Int32Array, new_null_array};
-///
-/// let null_array = new_null_array(&DataType::Int32, 3);
-/// let array: ArrayRef = Arc::new(Int32Array::from(vec![None, None, None]));
-///
-/// assert_eq!(&array, &null_array);
-/// ```
-pub fn new_null_array(data_type: &DataType, length: usize) -> ArrayRef {
- // context: https://github.com/apache/arrow/pull/9469#discussion_r574761687
- match data_type {
- DataType::Null => Arc::new(NullArray::new(length)),
- DataType::Boolean => {
- let null_buf: Buffer = MutableBuffer::new_null(length).into();
- make_array(ArrayData::new(
- data_type.clone(),
- length,
- Some(length),
- Some(null_buf.clone()),
- 0,
- vec![null_buf],
- vec![],
- ))
- }
- DataType::Int8 => new_null_sized_array::<Int8Type>(data_type, length),
- DataType::UInt8 => new_null_sized_array::<UInt8Type>(data_type, length),
- DataType::Int16 => new_null_sized_array::<Int16Type>(data_type, length),
- DataType::UInt16 => new_null_sized_array::<UInt16Type>(data_type, length),
- DataType::Float16 => unreachable!(),
- DataType::Int32 => new_null_sized_array::<Int32Type>(data_type, length),
- DataType::UInt32 => new_null_sized_array::<UInt32Type>(data_type, length),
- DataType::Float32 => new_null_sized_array::<Float32Type>(data_type, length),
- DataType::Date32 => new_null_sized_array::<Date32Type>(data_type, length),
- // expanding this into Date23{unit}Type results in needless branching
- DataType::Time32(_) => new_null_sized_array::<Int32Type>(data_type, length),
- DataType::Int64 => new_null_sized_array::<Int64Type>(data_type, length),
- DataType::UInt64 => new_null_sized_array::<UInt64Type>(data_type, length),
- DataType::Float64 => new_null_sized_array::<Float64Type>(data_type, length),
- DataType::Date64 => new_null_sized_array::<Date64Type>(data_type, length),
- // expanding this into Timestamp{unit}Type results in needless branching
- DataType::Timestamp(_, _) => new_null_sized_array::<Int64Type>(data_type, length),
- DataType::Time64(_) => new_null_sized_array::<Int64Type>(data_type, length),
- DataType::Duration(_) => new_null_sized_array::<Int64Type>(data_type, length),
- DataType::Interval(unit) => match unit {
- IntervalUnit::YearMonth => {
- new_null_sized_array::<IntervalYearMonthType>(data_type, length)
- }
- IntervalUnit::DayTime => {
- new_null_sized_array::<IntervalDayTimeType>(data_type, length)
- }
- },
- DataType::FixedSizeBinary(value_len) => make_array(ArrayData::new(
- data_type.clone(),
- length,
- Some(length),
- Some(MutableBuffer::new_null(length).into()),
- 0,
- vec![Buffer::from(vec![0u8; *value_len as usize * length])],
- vec![],
- )),
- DataType::Binary | DataType::Utf8 => {
- new_null_binary_array::<i32>(data_type, length)
- }
- DataType::LargeBinary | DataType::LargeUtf8 => {
- new_null_binary_array::<i64>(data_type, length)
- }
- DataType::List(field) => {
- new_null_list_array::<i32>(data_type, field.data_type(), length)
- }
- DataType::LargeList(field) => {
- new_null_list_array::<i64>(data_type, field.data_type(), length)
- }
- DataType::FixedSizeList(field, value_len) => make_array(ArrayData::new(
- data_type.clone(),
- length,
- Some(length),
- Some(MutableBuffer::new_null(length).into()),
- 0,
- vec![],
- vec![
- new_null_array(field.data_type(), *value_len as usize * length)
- .data()
- .clone(),
- ],
- )),
- DataType::Struct(fields) => make_array(ArrayData::new(
- data_type.clone(),
- length,
- Some(length),
- Some(MutableBuffer::new_null(length).into()),
- 0,
- vec![],
- fields
- .iter()
- .map(|field| ArrayData::new_empty(field.data_type()))
- .collect(),
- )),
- DataType::Union(_) => {
- unimplemented!("Creating null Union array not yet supported")
- }
- DataType::Dictionary(key, value) => {
- let keys = new_null_array(key, length);
- let keys = keys.data();
-
- make_array(ArrayData::new(
- data_type.clone(),
- length,
- Some(length),
- keys.null_buffer().cloned(),
- 0,
- keys.buffers().into(),
- vec![new_empty_array(value.as_ref()).data().clone()],
- ))
- }
- DataType::Decimal(_, _) => {
- unimplemented!("Creating null Decimal array not yet supported")
- }
- }
-}
-
-#[inline]
-fn new_null_list_array<OffsetSize: OffsetSizeTrait>(
- data_type: &DataType,
- child_data_type: &DataType,
- length: usize,
-) -> ArrayRef {
- make_array(ArrayData::new(
- data_type.clone(),
- length,
- Some(length),
- Some(MutableBuffer::new_null(length).into()),
- 0,
- vec![Buffer::from(
- vec![OffsetSize::zero(); length + 1].to_byte_slice(),
- )],
- vec![ArrayData::new_empty(child_data_type)],
- ))
-}
-
-#[inline]
-fn new_null_binary_array<OffsetSize: OffsetSizeTrait>(
- data_type: &DataType,
- length: usize,
-) -> ArrayRef {
- make_array(ArrayData::new(
- data_type.clone(),
- length,
- Some(length),
- Some(MutableBuffer::new_null(length).into()),
- 0,
- vec![
- Buffer::from(vec![OffsetSize::zero(); length + 1].to_byte_slice()),
- MutableBuffer::new(0).into(),
- ],
- vec![],
- ))
-}
-
-#[inline]
-fn new_null_sized_array<T: ArrowPrimitiveType>(
- data_type: &DataType,
- length: usize,
-) -> ArrayRef {
- make_array(ArrayData::new(
- data_type.clone(),
- length,
- Some(length),
- Some(MutableBuffer::new_null(length).into()),
- 0,
- vec![Buffer::from(vec![0u8; length * T::get_byte_width()])],
- vec![],
- ))
-}
-
-/// Creates a new array from two FFI pointers. Used to import arrays from the C Data Interface
-/// # Safety
-/// Assumes that these pointers represent valid C Data Interfaces, both in memory
-/// representation and lifetime via the `release` mechanism.
-pub unsafe fn make_array_from_raw(
- array: *const ffi::FFI_ArrowArray,
- schema: *const ffi::FFI_ArrowSchema,
-) -> Result<ArrayRef> {
- let array = ffi::ArrowArray::try_from_raw(array, schema)?;
- let data = ArrayData::try_from(array)?;
- Ok(make_array(data))
-}
-// Helper function for printing potentially long arrays.
-pub(super) fn print_long_array<A, F>(
- array: &A,
- f: &mut fmt::Formatter,
- print_item: F,
-) -> fmt::Result
-where
- A: Array,
- F: Fn(&A, usize, &mut fmt::Formatter) -> fmt::Result,
-{
- let head = std::cmp::min(10, array.len());
-
- for i in 0..head {
- if array.is_null(i) {
- writeln!(f, " null,")?;
- } else {
- write!(f, " ")?;
- print_item(&array, i, f)?;
- writeln!(f, ",")?;
- }
- }
- if array.len() > 10 {
- if array.len() > 20 {
- writeln!(f, " ...{} elements...,", array.len() - 20)?;
- }
-
- let tail = std::cmp::max(head, array.len() - 10);
-
- for i in tail..array.len() {
- if array.is_null(i) {
- writeln!(f, " null,")?;
- } else {
- write!(f, " ")?;
- print_item(&array, i, f)?;
- writeln!(f, ",")?;
- }
- }
- }
- Ok(())
-}
-
-#[cfg(test)]
-mod tests {
- use super::*;
- #[test]
- fn test_empty_primitive() {
- let array = new_empty_array(&DataType::Int32);
- let a = array.as_any().downcast_ref::<Int32Array>().unwrap();
- assert_eq!(a.len(), 0);
- let expected: &[i32] = &[];
- assert_eq!(a.values(), expected);
- }
-
- #[test]
- fn test_empty_variable_sized() {
- let array = new_empty_array(&DataType::Utf8);
- let a = array.as_any().downcast_ref::<StringArray>().unwrap();
- assert_eq!(a.len(), 0);
- assert_eq!(a.value_offsets()[0], 0i32);
- }
-
- #[test]
- fn test_empty_list_primitive() {
- let data_type =
- DataType::List(Box::new(Field::new("item", DataType::Int32, false)));
- let array = new_empty_array(&data_type);
- let a = array.as_any().downcast_ref::<ListArray>().unwrap();
- assert_eq!(a.len(), 0);
- assert_eq!(a.value_offsets()[0], 0i32);
- }
-
- #[test]
- fn test_null_boolean() {
- let array = new_null_array(&DataType::Boolean, 9);
- let a = array.as_any().downcast_ref::<BooleanArray>().unwrap();
- assert_eq!(a.len(), 9);
- for i in 0..9 {
- assert!(a.is_null(i));
- }
- }
-
- #[test]
- fn test_null_primitive() {
- let array = new_null_array(&DataType::Int32, 9);
- let a = array.as_any().downcast_ref::<Int32Array>().unwrap();
- assert_eq!(a.len(), 9);
- for i in 0..9 {
- assert!(a.is_null(i));
- }
- }
-
- #[test]
- fn test_null_variable_sized() {
- let array = new_null_array(&DataType::Utf8, 9);
- let a = array.as_any().downcast_ref::<StringArray>().unwrap();
- assert_eq!(a.len(), 9);
- assert_eq!(a.value_offsets()[9], 0i32);
- for i in 0..9 {
- assert!(a.is_null(i));
- }
- }
-
- #[test]
- fn test_null_list_primitive() {
- let data_type =
- DataType::List(Box::new(Field::new("item", DataType::Int32, true)));
- let array = new_null_array(&data_type, 9);
- let a = array.as_any().downcast_ref::<ListArray>().unwrap();
- assert_eq!(a.len(), 9);
- assert_eq!(a.value_offsets()[9], 0i32);
- for i in 0..9 {
- assert!(a.is_null(i));
- }
- }
-
- #[test]
- fn test_null_dictionary() {
- let values = vec![None, None, None, None, None, None, None, None, None]
- as Vec<Option<&str>>;
-
- let array: DictionaryArray<Int8Type> = values.into_iter().collect();
- let array = Arc::new(array) as ArrayRef;
-
- let null_array = new_null_array(array.data_type(), 9);
- assert_eq!(&array, &null_array);
- assert_eq!(
- array.data().buffers()[0].len(),
- null_array.data().buffers()[0].len()
- );
- }
-}
diff --git a/arrow/src/array/array_binary.rs b/arrow/src/array/array_binary.rs
deleted file mode 100644
index 014d1bf..0000000
--- a/arrow/src/array/array_binary.rs
+++ /dev/null
@@ -1,1240 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-use std::convert::{From, TryInto};
-use std::fmt;
-use std::mem;
-use std::{any::Any, iter::FromIterator};
-
-use super::{
- array::print_long_array, raw_pointer::RawPtrBox, Array, ArrayData,
- FixedSizeListArray, GenericBinaryIter, GenericListArray, OffsetSizeTrait,
-};
-use crate::buffer::Buffer;
-use crate::error::ArrowError;
-use crate::util::bit_util;
-use crate::{buffer::MutableBuffer, datatypes::DataType};
-
-/// Like OffsetSizeTrait, but specialized for Binary
-// This allow us to expose a constant datatype for the GenericBinaryArray
-pub trait BinaryOffsetSizeTrait: OffsetSizeTrait {
- const DATA_TYPE: DataType;
-}
-
-impl BinaryOffsetSizeTrait for i32 {
- const DATA_TYPE: DataType = DataType::Binary;
-}
-
-impl BinaryOffsetSizeTrait for i64 {
- const DATA_TYPE: DataType = DataType::LargeBinary;
-}
-
-pub struct GenericBinaryArray<OffsetSize: BinaryOffsetSizeTrait> {
- data: ArrayData,
- value_offsets: RawPtrBox<OffsetSize>,
- value_data: RawPtrBox<u8>,
-}
-
-impl<OffsetSize: BinaryOffsetSizeTrait> GenericBinaryArray<OffsetSize> {
- /// Returns the length for value at index `i`.
- #[inline]
- pub fn value_length(&self, i: usize) -> OffsetSize {
- let offsets = self.value_offsets();
- offsets[i + 1] - offsets[i]
- }
-
- /// Returns a clone of the value data buffer
- pub fn value_data(&self) -> Buffer {
- self.data.buffers()[1].clone()
- }
-
- /// Returns the offset values in the offsets buffer
- #[inline]
- pub fn value_offsets(&self) -> &[OffsetSize] {
- // Soundness
- // pointer alignment & location is ensured by RawPtrBox
- // buffer bounds/offset is ensured by the ArrayData instance.
- unsafe {
- std::slice::from_raw_parts(
- self.value_offsets.as_ptr().add(self.data.offset()),
- self.len() + 1,
- )
- }
- }
-
- /// Returns the element at index `i` as bytes slice
- /// # Safety
- /// Caller is responsible for ensuring that the index is within the bounds of the array
- pub unsafe fn value_unchecked(&self, i: usize) -> &[u8] {
- let end = *self.value_offsets().get_unchecked(i + 1);
- let start = *self.value_offsets().get_unchecked(i);
-
- // Soundness
- // pointer alignment & location is ensured by RawPtrBox
- // buffer bounds/offset is ensured by the value_offset invariants
-
- // Safety of `to_isize().unwrap()`
- // `start` and `end` are &OffsetSize, which is a generic type that implements the
- // OffsetSizeTrait. Currently, only i32 and i64 implement OffsetSizeTrait,
- // both of which should cleanly cast to isize on an architecture that supports
- // 32/64-bit offsets
- std::slice::from_raw_parts(
- self.value_data.as_ptr().offset(start.to_isize().unwrap()),
- (end - start).to_usize().unwrap(),
- )
- }
-
- /// Returns the element at index `i` as bytes slice
- pub fn value(&self, i: usize) -> &[u8] {
- assert!(i < self.data.len(), "BinaryArray out of bounds access");
- //Soundness: length checked above, offset buffer length is 1 larger than logical array length
- let end = unsafe { self.value_offsets().get_unchecked(i + 1) };
- let start = unsafe { self.value_offsets().get_unchecked(i) };
-
- // Soundness
- // pointer alignment & location is ensured by RawPtrBox
- // buffer bounds/offset is ensured by the value_offset invariants
-
- // Safety of `to_isize().unwrap()`
- // `start` and `end` are &OffsetSize, which is a generic type that implements the
- // OffsetSizeTrait. Currently, only i32 and i64 implement OffsetSizeTrait,
- // both of which should cleanly cast to isize on an architecture that supports
- // 32/64-bit offsets
- unsafe {
- std::slice::from_raw_parts(
- self.value_data.as_ptr().offset(start.to_isize().unwrap()),
- (*end - *start).to_usize().unwrap(),
- )
- }
- }
-
- /// Creates a [GenericBinaryArray] from a vector of byte slices
- pub fn from_vec(v: Vec<&[u8]>) -> Self {
- let mut offsets = Vec::with_capacity(v.len() + 1);
- let mut values = Vec::new();
- let mut length_so_far: OffsetSize = OffsetSize::zero();
- offsets.push(length_so_far);
- for s in &v {
- length_so_far += OffsetSize::from_usize(s.len()).unwrap();
- offsets.push(length_so_far);
- values.extend_from_slice(s);
- }
- let array_data = ArrayData::builder(OffsetSize::DATA_TYPE)
- .len(v.len())
- .add_buffer(Buffer::from_slice_ref(&offsets))
- .add_buffer(Buffer::from_slice_ref(&values))
- .build();
- GenericBinaryArray::<OffsetSize>::from(array_data)
- }
-
- /// Creates a [GenericBinaryArray] from a vector of Optional (null) byte slices
- pub fn from_opt_vec(v: Vec<Option<&[u8]>>) -> Self {
- v.into_iter().collect()
- }
-
- fn from_list(v: GenericListArray<OffsetSize>) -> Self {
- assert_eq!(
- v.data_ref().child_data()[0].child_data().len(),
- 0,
- "BinaryArray can only be created from list array of u8 values \
- (i.e. List<PrimitiveArray<u8>>)."
- );
- assert_eq!(
- v.data_ref().child_data()[0].data_type(),
- &DataType::UInt8,
- "BinaryArray can only be created from List<u8> arrays, mismatched data types."
- );
-
- let mut builder = ArrayData::builder(OffsetSize::DATA_TYPE)
- .len(v.len())
- .add_buffer(v.data_ref().buffers()[0].clone())
- .add_buffer(v.data_ref().child_data()[0].buffers()[0].clone());
- if let Some(bitmap) = v.data_ref().null_bitmap() {
- builder = builder.null_bit_buffer(bitmap.bits.clone())
- }
-
- let data = builder.build();
- Self::from(data)
- }
-}
-
-impl<'a, T: BinaryOffsetSizeTrait> GenericBinaryArray<T> {
- /// constructs a new iterator
- pub fn iter(&'a self) -> GenericBinaryIter<'a, T> {
- GenericBinaryIter::<'a, T>::new(&self)
- }
-}
-
-impl<OffsetSize: BinaryOffsetSizeTrait> fmt::Debug for GenericBinaryArray<OffsetSize> {
- fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
- let prefix = if OffsetSize::is_large() { "Large" } else { "" };
-
- write!(f, "{}BinaryArray\n[\n", prefix)?;
- print_long_array(self, f, |array, index, f| {
- fmt::Debug::fmt(&array.value(index), f)
- })?;
- write!(f, "]")
- }
-}
-
-impl<OffsetSize: BinaryOffsetSizeTrait> Array for GenericBinaryArray<OffsetSize> {
- fn as_any(&self) -> &Any {
- self
- }
-
- fn data(&self) -> &ArrayData {
- &self.data
- }
-
- /// Returns the total number of bytes of memory occupied by the buffers owned by this [$name].
- fn get_buffer_memory_size(&self) -> usize {
- self.data.get_buffer_memory_size()
- }
-
- /// Returns the total number of bytes of memory occupied physically by this [$name].
- fn get_array_memory_size(&self) -> usize {
- self.data.get_array_memory_size() + mem::size_of_val(self)
- }
-}
-
-impl<OffsetSize: BinaryOffsetSizeTrait> From<ArrayData>
- for GenericBinaryArray<OffsetSize>
-{
- fn from(data: ArrayData) -> Self {
- assert_eq!(
- data.data_type(),
- &<OffsetSize as BinaryOffsetSizeTrait>::DATA_TYPE,
- "[Large]BinaryArray expects Datatype::[Large]Binary"
- );
- assert_eq!(
- data.buffers().len(),
- 2,
- "BinaryArray data should contain 2 buffers only (offsets and values)"
- );
- let offsets = data.buffers()[0].as_ptr();
- let values = data.buffers()[1].as_ptr();
- Self {
- data,
- value_offsets: unsafe { RawPtrBox::new(offsets) },
- value_data: unsafe { RawPtrBox::new(values) },
- }
- }
-}
-
-impl<Ptr, OffsetSize: BinaryOffsetSizeTrait> FromIterator<Option<Ptr>>
- for GenericBinaryArray<OffsetSize>
-where
- Ptr: AsRef<[u8]>,
-{
- fn from_iter<I: IntoIterator<Item = Option<Ptr>>>(iter: I) -> Self {
- let iter = iter.into_iter();
- let (_, data_len) = iter.size_hint();
- let data_len = data_len.expect("Iterator must be sized"); // panic if no upper bound.
-
- let mut offsets = Vec::with_capacity(data_len + 1);
- let mut values = Vec::new();
- let mut null_buf = MutableBuffer::new_null(data_len);
- let mut length_so_far: OffsetSize = OffsetSize::zero();
- offsets.push(length_so_far);
-
- {
- let null_slice = null_buf.as_slice_mut();
-
- for (i, s) in iter.enumerate() {
- if let Some(s) = s {
- let s = s.as_ref();
- bit_util::set_bit(null_slice, i);
- length_so_far += OffsetSize::from_usize(s.len()).unwrap();
- values.extend_from_slice(s);
- }
- // always add an element in offsets
- offsets.push(length_so_far);
- }
- }
-
- // calculate actual data_len, which may be different from the iterator's upper bound
- let data_len = offsets.len() - 1;
- let array_data = ArrayData::builder(OffsetSize::DATA_TYPE)
- .len(data_len)
- .add_buffer(Buffer::from_slice_ref(&offsets))
- .add_buffer(Buffer::from_slice_ref(&values))
- .null_bit_buffer(null_buf.into())
- .build();
- Self::from(array_data)
- }
-}
-
-/// An array where each element is a byte whose maximum length is represented by a i32.
-pub type BinaryArray = GenericBinaryArray<i32>;
-
-/// An array where each element is a byte whose maximum length is represented by a i64.
-pub type LargeBinaryArray = GenericBinaryArray<i64>;
-
-impl<'a, T: BinaryOffsetSizeTrait> IntoIterator for &'a GenericBinaryArray<T> {
- type Item = Option<&'a [u8]>;
- type IntoIter = GenericBinaryIter<'a, T>;
-
- fn into_iter(self) -> Self::IntoIter {
- GenericBinaryIter::<'a, T>::new(self)
- }
-}
-
-impl<OffsetSize: BinaryOffsetSizeTrait> From<Vec<Option<&[u8]>>>
- for GenericBinaryArray<OffsetSize>
-{
- fn from(v: Vec<Option<&[u8]>>) -> Self {
- GenericBinaryArray::<OffsetSize>::from_opt_vec(v)
- }
-}
-
-impl<OffsetSize: BinaryOffsetSizeTrait> From<Vec<&[u8]>>
- for GenericBinaryArray<OffsetSize>
-{
- fn from(v: Vec<&[u8]>) -> Self {
- GenericBinaryArray::<OffsetSize>::from_vec(v)
- }
-}
-
-impl<T: BinaryOffsetSizeTrait> From<GenericListArray<T>> for GenericBinaryArray<T> {
- fn from(v: GenericListArray<T>) -> Self {
- GenericBinaryArray::<T>::from_list(v)
- }
-}
-
-/// A type of `FixedSizeListArray` whose elements are binaries.
-///
-/// # Examples
-///
-/// Create an array from an iterable argument of byte slices.
-///
-/// ```
-/// use arrow::array::{Array, FixedSizeBinaryArray};
-/// let input_arg = vec![ vec![1, 2], vec![3, 4], vec![5, 6] ];
-/// let arr = FixedSizeBinaryArray::try_from_iter(input_arg.into_iter()).unwrap();
-///
-/// assert_eq!(3, arr.len());
-///
-/// ```
-/// Create an array from an iterable argument of sparse byte slices.
-/// Sparsity means that the input argument can contain `None` items.
-/// ```
-/// use arrow::array::{Array, FixedSizeBinaryArray};
-/// let input_arg = vec![ None, Some(vec![7, 8]), Some(vec![9, 10]), None, Some(vec![13, 14]) ];
-/// let arr = FixedSizeBinaryArray::try_from_sparse_iter(input_arg.into_iter()).unwrap();
-/// assert_eq!(5, arr.len())
-///
-/// ```
-///
-pub struct FixedSizeBinaryArray {
- data: ArrayData,
- value_data: RawPtrBox<u8>,
- length: i32,
-}
-
-impl FixedSizeBinaryArray {
- /// Returns the element at index `i` as a byte slice.
- pub fn value(&self, i: usize) -> &[u8] {
- assert!(
- i < self.data.len(),
- "FixedSizeBinaryArray out of bounds access"
- );
- let offset = i.checked_add(self.data.offset()).unwrap();
- unsafe {
- let pos = self.value_offset_at(offset);
- std::slice::from_raw_parts(
- self.value_data.as_ptr().offset(pos as isize),
- (self.value_offset_at(offset + 1) - pos) as usize,
- )
- }
- }
-
- /// Returns the offset for the element at index `i`.
- ///
- /// Note this doesn't do any bound checking, for performance reason.
- #[inline]
- pub fn value_offset(&self, i: usize) -> i32 {
- self.value_offset_at(self.data.offset() + i)
- }
-
- /// Returns the length for an element.
- ///
- /// All elements have the same length as the array is a fixed size.
- #[inline]
- pub fn value_length(&self) -> i32 {
- self.length
- }
-
- /// Returns a clone of the value data buffer
- pub fn value_data(&self) -> Buffer {
- self.data.buffers()[0].clone()
- }
-
- /// Create an array from an iterable argument of sparse byte slices.
- /// Sparsity means that items returned by the iterator are optional, i.e input argument can
- /// contain `None` items.
- ///
- /// # Examples
- ///
- /// ```
- /// use arrow::array::FixedSizeBinaryArray;
- /// let input_arg = vec![
- /// None,
- /// Some(vec![7, 8]),
- /// Some(vec![9, 10]),
- /// None,
- /// Some(vec![13, 14]),
- /// None,
- /// ];
- /// let array = FixedSizeBinaryArray::try_from_sparse_iter(input_arg.into_iter()).unwrap();
- /// ```
- ///
- /// # Errors
- ///
- /// Returns error if argument has length zero, or sizes of nested slices don't match.
- pub fn try_from_sparse_iter<T, U>(mut iter: T) -> Result<Self, ArrowError>
- where
- T: Iterator<Item = Option<U>>,
- U: AsRef<[u8]>,
- {
- let mut len = 0;
- let mut size = None;
- let mut byte = 0;
- let mut null_buf = MutableBuffer::from_len_zeroed(0);
- let mut buffer = MutableBuffer::from_len_zeroed(0);
- let mut prepend = 0;
- iter.try_for_each(|item| -> Result<(), ArrowError> {
- // extend null bitmask by one byte per each 8 items
- if byte == 0 {
- null_buf.push(0u8);
- byte = 8;
- }
- byte -= 1;
-
- if let Some(slice) = item {
- let slice = slice.as_ref();
- if let Some(size) = size {
- if size != slice.len() {
- return Err(ArrowError::InvalidArgumentError(format!(
- "Nested array size mismatch: one is {}, and the other is {}",
- size,
- slice.len()
- )));
- }
- } else {
- size = Some(slice.len());
- buffer.extend_zeros(slice.len() * prepend);
- }
- bit_util::set_bit(null_buf.as_slice_mut(), len);
- buffer.extend_from_slice(slice);
- } else if let Some(size) = size {
- buffer.extend_zeros(size);
- } else {
- prepend += 1;
- }
-
- len += 1;
-
- Ok(())
- })?;
-
- if len == 0 {
- return Err(ArrowError::InvalidArgumentError(
- "Input iterable argument has no data".to_owned(),
- ));
- }
-
- let size = size.unwrap_or(0);
- let array_data = ArrayData::new(
- DataType::FixedSizeBinary(size as i32),
- len,
- None,
- Some(null_buf.into()),
- 0,
- vec![buffer.into()],
- vec![],
- );
- Ok(FixedSizeBinaryArray::from(array_data))
- }
-
- /// Create an array from an iterable argument of byte slices.
- ///
- /// # Examples
- ///
- /// ```
- /// use arrow::array::FixedSizeBinaryArray;
- /// let input_arg = vec![
- /// vec![1, 2],
- /// vec![3, 4],
- /// vec![5, 6],
- /// ];
- /// let array = FixedSizeBinaryArray::try_from_iter(input_arg.into_iter()).unwrap();
- /// ```
- ///
- /// # Errors
- ///
- /// Returns error if argument has length zero, or sizes of nested slices don't match.
- pub fn try_from_iter<T, U>(mut iter: T) -> Result<Self, ArrowError>
- where
- T: Iterator<Item = U>,
- U: AsRef<[u8]>,
- {
- let mut len = 0;
- let mut size = None;
- let mut buffer = MutableBuffer::from_len_zeroed(0);
- iter.try_for_each(|item| -> Result<(), ArrowError> {
- let slice = item.as_ref();
- if let Some(size) = size {
- if size != slice.len() {
- return Err(ArrowError::InvalidArgumentError(format!(
- "Nested array size mismatch: one is {}, and the other is {}",
- size,
- slice.len()
- )));
- }
- } else {
- size = Some(slice.len());
- }
- buffer.extend_from_slice(slice);
-
- len += 1;
-
- Ok(())
- })?;
-
- if len == 0 {
- return Err(ArrowError::InvalidArgumentError(
- "Input iterable argument has no data".to_owned(),
- ));
- }
-
- let size = size.unwrap_or(0);
- let array_data = ArrayData::builder(DataType::FixedSizeBinary(size as i32))
- .len(len)
- .add_buffer(buffer.into())
- .build();
- Ok(FixedSizeBinaryArray::from(array_data))
- }
-
- #[inline]
- fn value_offset_at(&self, i: usize) -> i32 {
- self.length * i as i32
- }
-}
-
-impl From<ArrayData> for FixedSizeBinaryArray {
- fn from(data: ArrayData) -> Self {
- assert_eq!(
- data.buffers().len(),
- 1,
- "FixedSizeBinaryArray data should contain 1 buffer only (values)"
- );
- let value_data = data.buffers()[0].as_ptr();
- let length = match data.data_type() {
- DataType::FixedSizeBinary(len) => *len,
- _ => panic!("Expected data type to be FixedSizeBinary"),
- };
- Self {
- data,
- value_data: unsafe { RawPtrBox::new(value_data) },
- length,
- }
- }
-}
-
-/// Creates a `FixedSizeBinaryArray` from `FixedSizeList<u8>` array
-impl From<FixedSizeListArray> for FixedSizeBinaryArray {
- fn from(v: FixedSizeListArray) -> Self {
- assert_eq!(
- v.data_ref().child_data()[0].child_data().len(),
- 0,
- "FixedSizeBinaryArray can only be created from list array of u8 values \
- (i.e. FixedSizeList<PrimitiveArray<u8>>)."
- );
- assert_eq!(
- v.data_ref().child_data()[0].data_type(),
- &DataType::UInt8,
- "FixedSizeBinaryArray can only be created from FixedSizeList<u8> arrays, mismatched data types."
- );
-
- let mut builder = ArrayData::builder(DataType::FixedSizeBinary(v.value_length()))
- .len(v.len())
- .add_buffer(v.data_ref().child_data()[0].buffers()[0].clone());
- if let Some(bitmap) = v.data_ref().null_bitmap() {
- builder = builder.null_bit_buffer(bitmap.bits.clone())
- }
-
- let data = builder.build();
- Self::from(data)
- }
-}
-
-impl fmt::Debug for FixedSizeBinaryArray {
- fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
- write!(f, "FixedSizeBinaryArray<{}>\n[\n", self.value_length())?;
- print_long_array(self, f, |array, index, f| {
- fmt::Debug::fmt(&array.value(index), f)
- })?;
- write!(f, "]")
- }
-}
-
-impl Array for FixedSizeBinaryArray {
- fn as_any(&self) -> &Any {
- self
- }
-
- fn data(&self) -> &ArrayData {
- &self.data
- }
-
- /// Returns the total number of bytes of memory occupied by the buffers owned by this [FixedSizeBinaryArray].
- fn get_buffer_memory_size(&self) -> usize {
- self.data.get_buffer_memory_size()
- }
-
- /// Returns the total number of bytes of memory occupied physically by this [FixedSizeBinaryArray].
- fn get_array_memory_size(&self) -> usize {
- self.data.get_array_memory_size() + mem::size_of_val(self)
- }
-}
-
-/// A type of `DecimalArray` whose elements are binaries.
-///
-/// # Examples
-///
-/// ```
-/// use arrow::array::{Array, DecimalArray, DecimalBuilder};
-/// use arrow::datatypes::DataType;
-/// let mut builder = DecimalBuilder::new(30, 23, 6);
-///
-/// builder.append_value(8_887_000_000).unwrap();
-/// builder.append_null().unwrap();
-/// builder.append_value(-8_887_000_000).unwrap();
-/// let decimal_array: DecimalArray = builder.finish();
-///
-/// assert_eq!(&DataType::Decimal(23, 6), decimal_array.data_type());
-/// assert_eq!(8_887_000_000, decimal_array.value(0));
-/// assert_eq!("8887.000000", decimal_array.value_as_string(0));
-/// assert_eq!(3, decimal_array.len());
-/// assert_eq!(1, decimal_array.null_count());
-/// assert_eq!(32, decimal_array.value_offset(2));
-/// assert_eq!(16, decimal_array.value_length());
-/// assert_eq!(23, decimal_array.precision());
-/// assert_eq!(6, decimal_array.scale());
-/// ```
-///
-pub struct DecimalArray {
- data: ArrayData,
- value_data: RawPtrBox<u8>,
- precision: usize,
- scale: usize,
- length: i32,
-}
-
-impl DecimalArray {
- /// Returns the element at index `i` as i128.
- pub fn value(&self, i: usize) -> i128 {
- assert!(i < self.data.len(), "DecimalArray out of bounds access");
- let offset = i.checked_add(self.data.offset()).unwrap();
- let raw_val = unsafe {
- let pos = self.value_offset_at(offset);
- std::slice::from_raw_parts(
- self.value_data.as_ptr().offset(pos as isize),
- (self.value_offset_at(offset + 1) - pos) as usize,
- )
- };
- let as_array = raw_val.try_into();
- match as_array {
- Ok(v) if raw_val.len() == 16 => i128::from_le_bytes(v),
- _ => panic!("DecimalArray elements are not 128bit integers."),
- }
- }
-
- /// Returns the offset for the element at index `i`.
- ///
- /// Note this doesn't do any bound checking, for performance reason.
- #[inline]
- pub fn value_offset(&self, i: usize) -> i32 {
- self.value_offset_at(self.data.offset() + i)
- }
-
- /// Returns the length for an element.
- ///
- /// All elements have the same length as the array is a fixed size.
- #[inline]
- pub fn value_length(&self) -> i32 {
- self.length
- }
-
- /// Returns a clone of the value data buffer
- pub fn value_data(&self) -> Buffer {
- self.data.buffers()[0].clone()
- }
-
- #[inline]
- fn value_offset_at(&self, i: usize) -> i32 {
- self.length * i as i32
- }
-
- #[inline]
- pub fn value_as_string(&self, row: usize) -> String {
- let decimal_string = self.value(row).to_string();
- if self.scale == 0 {
- decimal_string
- } else {
- let splits = decimal_string.split_at(decimal_string.len() - self.scale);
- format!("{}.{}", splits.0, splits.1)
- }
- }
-
- pub fn from_fixed_size_list_array(
- v: FixedSizeListArray,
- precision: usize,
- scale: usize,
- ) -> Self {
- assert_eq!(
- v.data_ref().child_data()[0].child_data().len(),
- 0,
- "DecimalArray can only be created from list array of u8 values \
- (i.e. FixedSizeList<PrimitiveArray<u8>>)."
- );
- assert_eq!(
- v.data_ref().child_data()[0].data_type(),
- &DataType::UInt8,
- "DecimalArray can only be created from FixedSizeList<u8> arrays, mismatched data types."
- );
-
- let mut builder = ArrayData::builder(DataType::Decimal(precision, scale))
- .len(v.len())
- .add_buffer(v.data_ref().child_data()[0].buffers()[0].clone());
- if let Some(bitmap) = v.data_ref().null_bitmap() {
- builder = builder.null_bit_buffer(bitmap.bits.clone())
- }
-
- let data = builder.build();
- Self::from(data)
- }
- pub fn precision(&self) -> usize {
- self.precision
- }
-
- pub fn scale(&self) -> usize {
- self.scale
- }
-}
-
-impl From<ArrayData> for DecimalArray {
- fn from(data: ArrayData) -> Self {
- assert_eq!(
- data.buffers().len(),
- 1,
- "DecimalArray data should contain 1 buffer only (values)"
- );
- let values = data.buffers()[0].as_ptr();
- let (precision, scale) = match data.data_type() {
- DataType::Decimal(precision, scale) => (*precision, *scale),
- _ => panic!("Expected data type to be Decimal"),
- };
- let length = 16;
- Self {
- data,
- value_data: unsafe { RawPtrBox::new(values) },
- precision,
- scale,
- length,
- }
- }
-}
-
-impl fmt::Debug for DecimalArray {
- fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
- write!(f, "DecimalArray<{}, {}>\n[\n", self.precision, self.scale)?;
- print_long_array(self, f, |array, index, f| {
- let formatted_decimal = array.value_as_string(index);
-
- write!(f, "{}", formatted_decimal)
- })?;
- write!(f, "]")
- }
-}
-
-impl Array for DecimalArray {
- fn as_any(&self) -> &Any {
- self
- }
-
- fn data(&self) -> &ArrayData {
- &self.data
- }
-
- /// Returns the total number of bytes of memory occupied by the buffers owned by this [DecimalArray].
- fn get_buffer_memory_size(&self) -> usize {
- self.data.get_buffer_memory_size()
- }
-
- /// Returns the total number of bytes of memory occupied physically by this [DecimalArray].
- fn get_array_memory_size(&self) -> usize {
- self.data.get_array_memory_size() + mem::size_of_val(self)
- }
-}
-
-#[cfg(test)]
-mod tests {
- use crate::{
- array::{DecimalBuilder, LargeListArray, ListArray},
- datatypes::Field,
- };
-
- use super::*;
-
- #[test]
- fn test_binary_array() {
- let values: [u8; 12] = [
- b'h', b'e', b'l', b'l', b'o', b'p', b'a', b'r', b'q', b'u', b'e', b't',
- ];
- let offsets: [i32; 4] = [0, 5, 5, 12];
-
- // Array data: ["hello", "", "parquet"]
- let array_data = ArrayData::builder(DataType::Binary)
- .len(3)
- .add_buffer(Buffer::from_slice_ref(&offsets))
- .add_buffer(Buffer::from_slice_ref(&values))
- .build();
- let binary_array = BinaryArray::from(array_data);
- assert_eq!(3, binary_array.len());
- assert_eq!(0, binary_array.null_count());
- assert_eq!([b'h', b'e', b'l', b'l', b'o'], binary_array.value(0));
- assert_eq!([b'h', b'e', b'l', b'l', b'o'], unsafe {
- binary_array.value_unchecked(0)
- });
- assert_eq!([] as [u8; 0], binary_array.value(1));
- assert_eq!([] as [u8; 0], unsafe { binary_array.value_unchecked(1) });
- assert_eq!(
- [b'p', b'a', b'r', b'q', b'u', b'e', b't'],
- binary_array.value(2)
- );
- assert_eq!([b'p', b'a', b'r', b'q', b'u', b'e', b't'], unsafe {
- binary_array.value_unchecked(2)
- });
- assert_eq!(5, binary_array.value_offsets()[2]);
- assert_eq!(7, binary_array.value_length(2));
- for i in 0..3 {
- assert!(binary_array.is_valid(i));
- assert!(!binary_array.is_null(i));
- }
-
- // Test binary array with offset
- let array_data = ArrayData::builder(DataType::Binary)
- .len(4)
- .offset(1)
- .add_buffer(Buffer::from_slice_ref(&offsets))
- .add_buffer(Buffer::from_slice_ref(&values))
- .build();
- let binary_array = BinaryArray::from(array_data);
- assert_eq!(
- [b'p', b'a', b'r', b'q', b'u', b'e', b't'],
- binary_array.value(1)
- );
- assert_eq!(5, binary_array.value_offsets()[0]);
- assert_eq!(0, binary_array.value_length(0));
- assert_eq!(5, binary_array.value_offsets()[1]);
- assert_eq!(7, binary_array.value_length(1));
- }
-
- #[test]
- fn test_large_binary_array() {
- let values: [u8; 12] = [
- b'h', b'e', b'l', b'l', b'o', b'p', b'a', b'r', b'q', b'u', b'e', b't',
- ];
- let offsets: [i64; 4] = [0, 5, 5, 12];
-
- // Array data: ["hello", "", "parquet"]
- let array_data = ArrayData::builder(DataType::LargeBinary)
- .len(3)
- .add_buffer(Buffer::from_slice_ref(&offsets))
- .add_buffer(Buffer::from_slice_ref(&values))
- .build();
- let binary_array = LargeBinaryArray::from(array_data);
- assert_eq!(3, binary_array.len());
- assert_eq!(0, binary_array.null_count());
- assert_eq!([b'h', b'e', b'l', b'l', b'o'], binary_array.value(0));
- assert_eq!([b'h', b'e', b'l', b'l', b'o'], unsafe {
- binary_array.value_unchecked(0)
- });
- assert_eq!([] as [u8; 0], binary_array.value(1));
- assert_eq!([] as [u8; 0], unsafe { binary_array.value_unchecked(1) });
- assert_eq!(
- [b'p', b'a', b'r', b'q', b'u', b'e', b't'],
- binary_array.value(2)
- );
- assert_eq!([b'p', b'a', b'r', b'q', b'u', b'e', b't'], unsafe {
- binary_array.value_unchecked(2)
- });
- assert_eq!(5, binary_array.value_offsets()[2]);
- assert_eq!(7, binary_array.value_length(2));
- for i in 0..3 {
- assert!(binary_array.is_valid(i));
- assert!(!binary_array.is_null(i));
- }
-
- // Test binary array with offset
- let array_data = ArrayData::builder(DataType::LargeBinary)
- .len(4)
- .offset(1)
- .add_buffer(Buffer::from_slice_ref(&offsets))
- .add_buffer(Buffer::from_slice_ref(&values))
- .build();
- let binary_array = LargeBinaryArray::from(array_data);
- assert_eq!(
- [b'p', b'a', b'r', b'q', b'u', b'e', b't'],
- binary_array.value(1)
- );
- assert_eq!([b'p', b'a', b'r', b'q', b'u', b'e', b't'], unsafe {
- binary_array.value_unchecked(1)
- });
- assert_eq!(5, binary_array.value_offsets()[0]);
- assert_eq!(0, binary_array.value_length(0));
- assert_eq!(5, binary_array.value_offsets()[1]);
- assert_eq!(7, binary_array.value_length(1));
- }
-
- #[test]
- fn test_binary_array_from_list_array() {
- let values: [u8; 12] = [
- b'h', b'e', b'l', b'l', b'o', b'p', b'a', b'r', b'q', b'u', b'e', b't',
- ];
- let values_data = ArrayData::builder(DataType::UInt8)
- .len(12)
- .add_buffer(Buffer::from(&values[..]))
- .build();
- let offsets: [i32; 4] = [0, 5, 5, 12];
-
- // Array data: ["hello", "", "parquet"]
- let array_data1 = ArrayData::builder(DataType::Binary)
- .len(3)
- .add_buffer(Buffer::from_slice_ref(&offsets))
- .add_buffer(Buffer::from_slice_ref(&values))
- .build();
- let binary_array1 = BinaryArray::from(array_data1);
-
- let data_type =
- DataType::List(Box::new(Field::new("item", DataType::UInt8, false)));
- let array_data2 = ArrayData::builder(data_type)
- .len(3)
- .add_buffer(Buffer::from_slice_ref(&offsets))
- .add_child_data(values_data)
- .build();
- let list_array = ListArray::from(array_data2);
- let binary_array2 = BinaryArray::from(list_array);
-
- assert_eq!(2, binary_array2.data().buffers().len());
- assert_eq!(0, binary_array2.data().child_data().len());
-
- assert_eq!(binary_array1.len(), binary_array2.len());
- assert_eq!(binary_array1.null_count(), binary_array2.null_count());
- assert_eq!(binary_array1.value_offsets(), binary_array2.value_offsets());
- for i in 0..binary_array1.len() {
- assert_eq!(binary_array1.value(i), binary_array2.value(i));
- assert_eq!(binary_array1.value(i), unsafe {
- binary_array2.value_unchecked(i)
- });
- assert_eq!(binary_array1.value_length(i), binary_array2.value_length(i));
- }
- }
-
- #[test]
- fn test_large_binary_array_from_list_array() {
- let values: [u8; 12] = [
- b'h', b'e', b'l', b'l', b'o', b'p', b'a', b'r', b'q', b'u', b'e', b't',
- ];
- let values_data = ArrayData::builder(DataType::UInt8)
- .len(12)
- .add_buffer(Buffer::from(&values[..]))
- .build();
- let offsets: [i64; 4] = [0, 5, 5, 12];
-
- // Array data: ["hello", "", "parquet"]
- let array_data1 = ArrayData::builder(DataType::LargeBinary)
- .len(3)
- .add_buffer(Buffer::from_slice_ref(&offsets))
- .add_buffer(Buffer::from_slice_ref(&values))
- .build();
- let binary_array1 = LargeBinaryArray::from(array_data1);
-
- let data_type =
- DataType::LargeList(Box::new(Field::new("item", DataType::UInt8, false)));
- let array_data2 = ArrayData::builder(data_type)
- .len(3)
- .add_buffer(Buffer::from_slice_ref(&offsets))
- .add_child_data(values_data)
- .build();
- let list_array = LargeListArray::from(array_data2);
- let binary_array2 = LargeBinaryArray::from(list_array);
-
- assert_eq!(2, binary_array2.data().buffers().len());
- assert_eq!(0, binary_array2.data().child_data().len());
-
- assert_eq!(binary_array1.len(), binary_array2.len());
- assert_eq!(binary_array1.null_count(), binary_array2.null_count());
- assert_eq!(binary_array1.value_offsets(), binary_array2.value_offsets());
- for i in 0..binary_array1.len() {
- assert_eq!(binary_array1.value(i), binary_array2.value(i));
- assert_eq!(binary_array1.value(i), unsafe {
- binary_array2.value_unchecked(i)
- });
- assert_eq!(binary_array1.value_length(i), binary_array2.value_length(i));
- }
- }
-
- fn test_generic_binary_array_from_opt_vec<T: BinaryOffsetSizeTrait>() {
- let values: Vec<Option<&[u8]>> =
- vec![Some(b"one"), Some(b"two"), None, Some(b""), Some(b"three")];
- let array = GenericBinaryArray::<T>::from_opt_vec(values);
- assert_eq!(array.len(), 5);
- assert_eq!(array.value(0), b"one");
- assert_eq!(array.value(1), b"two");
- assert_eq!(array.value(3), b"");
- assert_eq!(array.value(4), b"three");
- assert_eq!(array.is_null(0), false);
- assert_eq!(array.is_null(1), false);
- assert_eq!(array.is_null(2), true);
- assert_eq!(array.is_null(3), false);
- assert_eq!(array.is_null(4), false);
- }
-
- #[test]
- fn test_large_binary_array_from_opt_vec() {
- test_generic_binary_array_from_opt_vec::<i64>()
- }
-
- #[test]
- fn test_binary_array_from_opt_vec() {
- test_generic_binary_array_from_opt_vec::<i32>()
- }
-
- #[test]
- fn test_binary_array_from_unbound_iter() {
- // iterator that doesn't declare (upper) size bound
- let value_iter = (0..)
- .scan(0usize, |pos, i| {
- if *pos < 10 {
- *pos += 1;
- Some(Some(format!("value {}", i)))
- } else {
- // actually returns up to 10 values
- None
- }
- })
- // limited using take()
- .take(100);
-
- let (_, upper_size_bound) = value_iter.size_hint();
- // the upper bound, defined by take above, is 100
- assert_eq!(upper_size_bound, Some(100));
- let binary_array: BinaryArray = value_iter.collect();
- // but the actual number of items in the array should be 10
- assert_eq!(binary_array.len(), 10);
- }
-
- #[test]
- #[should_panic(
- expected = "assertion failed: `(left == right)`\n left: `UInt32`,\n \
- right: `UInt8`: BinaryArray can only be created from List<u8> arrays, \
- mismatched data types."
- )]
- fn test_binary_array_from_incorrect_list_array() {
- let values: [u32; 12] = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11];
- let values_data = ArrayData::builder(DataType::UInt32)
- .len(12)
- .add_buffer(Buffer::from_slice_ref(&values))
- .build();
- let offsets: [i32; 4] = [0, 5, 5, 12];
-
- let data_type =
- DataType::List(Box::new(Field::new("item", DataType::UInt32, false)));
- let array_data = ArrayData::builder(data_type)
- .len(3)
- .add_buffer(Buffer::from_slice_ref(&offsets))
- .add_child_data(values_data)
- .build();
- let list_array = ListArray::from(array_data);
- BinaryArray::from(list_array);
- }
-
- #[test]
- fn test_fixed_size_binary_array() {
- let values: [u8; 15] = *b"hellotherearrow";
-
- let array_data = ArrayData::builder(DataType::FixedSizeBinary(5))
- .len(3)
- .add_buffer(Buffer::from(&values[..]))
- .build();
- let fixed_size_binary_array = FixedSizeBinaryArray::from(array_data);
- assert_eq!(3, fixed_size_binary_array.len());
- assert_eq!(0, fixed_size_binary_array.null_count());
- assert_eq!(
- [b'h', b'e', b'l', b'l', b'o'],
- fixed_size_binary_array.value(0)
- );
- assert_eq!(
- [b't', b'h', b'e', b'r', b'e'],
- fixed_size_binary_array.value(1)
- );
- assert_eq!(
- [b'a', b'r', b'r', b'o', b'w'],
- fixed_size_binary_array.value(2)
- );
- assert_eq!(5, fixed_size_binary_array.value_length());
- assert_eq!(10, fixed_size_binary_array.value_offset(2));
- for i in 0..3 {
- assert!(fixed_size_binary_array.is_valid(i));
- assert!(!fixed_size_binary_array.is_null(i));
- }
-
- // Test binary array with offset
- let array_data = ArrayData::builder(DataType::FixedSizeBinary(5))
- .len(2)
- .offset(1)
- .add_buffer(Buffer::from(&values[..]))
- .build();
- let fixed_size_binary_array = FixedSizeBinaryArray::from(array_data);
- assert_eq!(
- [b't', b'h', b'e', b'r', b'e'],
- fixed_size_binary_array.value(0)
- );
- assert_eq!(
- [b'a', b'r', b'r', b'o', b'w'],
- fixed_size_binary_array.value(1)
- );
- assert_eq!(2, fixed_size_binary_array.len());
- assert_eq!(5, fixed_size_binary_array.value_offset(0));
- assert_eq!(5, fixed_size_binary_array.value_length());
- assert_eq!(10, fixed_size_binary_array.value_offset(1));
- }
-
- #[test]
- #[should_panic(
- expected = "FixedSizeBinaryArray can only be created from list array of u8 values \
- (i.e. FixedSizeList<PrimitiveArray<u8>>)."
- )]
- fn test_fixed_size_binary_array_from_incorrect_list_array() {
- let values: [u32; 12] = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11];
- let values_data = ArrayData::builder(DataType::UInt32)
- .len(12)
- .add_buffer(Buffer::from_slice_ref(&values))
- .add_child_data(ArrayData::builder(DataType::Boolean).build())
- .build();
-
- let array_data = ArrayData::builder(DataType::FixedSizeList(
- Box::new(Field::new("item", DataType::Binary, false)),
- 4,
- ))
- .len(3)
- .add_child_data(values_data)
- .build();
- let list_array = FixedSizeListArray::from(array_data);
- FixedSizeBinaryArray::from(list_array);
- }
-
- #[test]
- #[should_panic(expected = "BinaryArray out of bounds access")]
- fn test_binary_array_get_value_index_out_of_bound() {
- let values: [u8; 12] =
- [104, 101, 108, 108, 111, 112, 97, 114, 113, 117, 101, 116];
- let offsets: [i32; 4] = [0, 5, 5, 12];
- let array_data = ArrayData::builder(DataType::Binary)
- .len(3)
- .add_buffer(Buffer::from_slice_ref(&offsets))
- .add_buffer(Buffer::from_slice_ref(&values))
- .build();
- let binary_array = BinaryArray::from(array_data);
- binary_array.value(4);
- }
-
- #[test]
- fn test_binary_array_fmt_debug() {
- let values: [u8; 15] = *b"hellotherearrow";
-
- let array_data = ArrayData::builder(DataType::FixedSizeBinary(5))
- .len(3)
- .add_buffer(Buffer::from(&values[..]))
- .build();
- let arr = FixedSizeBinaryArray::from(array_data);
- assert_eq!(
- "FixedSizeBinaryArray<5>\n[\n [104, 101, 108, 108, 111],\n [116, 104, 101, 114, 101],\n [97, 114, 114, 111, 119],\n]",
- format!("{:?}", arr)
- );
- }
-
- #[test]
- fn test_decimal_array() {
- // let val_8887: [u8; 16] = [192, 219, 180, 17, 2, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0];
- // let val_neg_8887: [u8; 16] = [64, 36, 75, 238, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255];
- let values: [u8; 32] = [
- 192, 219, 180, 17, 2, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 64, 36, 75, 238, 253,
- 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255,
- ];
- let array_data = ArrayData::builder(DataType::Decimal(23, 6))
- .len(2)
- .add_buffer(Buffer::from(&values[..]))
- .build();
- let decimal_array = DecimalArray::from(array_data);
- assert_eq!(8_887_000_000, decimal_array.value(0));
- assert_eq!(-8_887_000_000, decimal_array.value(1));
- assert_eq!(16, decimal_array.value_length());
- }
-
- #[test]
- fn test_decimal_array_fmt_debug() {
- let values: Vec<i128> = vec![8887000000, -8887000000];
- let mut decimal_builder = DecimalBuilder::new(3, 23, 6);
-
- values.iter().for_each(|&value| {
- decimal_builder.append_value(value).unwrap();
- });
- decimal_builder.append_null().unwrap();
- let arr = decimal_builder.finish();
- assert_eq!(
- "DecimalArray<23, 6>\n[\n 8887.000000,\n -8887.000000,\n null,\n]",
- format!("{:?}", arr)
- );
- }
-
- #[test]
- fn test_fixed_size_binary_array_from_iter() {
- let input_arg = vec![vec![1, 2], vec![3, 4], vec![5, 6]];
- let arr = FixedSizeBinaryArray::try_from_iter(input_arg.into_iter()).unwrap();
-
- assert_eq!(2, arr.value_length());
- assert_eq!(3, arr.len())
- }
-
- #[test]
- fn test_fixed_size_binary_array_from_sparse_iter() {
- let input_arg = vec![
- None,
- Some(vec![7, 8]),
- Some(vec![9, 10]),
- None,
- Some(vec![13, 14]),
- ];
- let arr =
- FixedSizeBinaryArray::try_from_sparse_iter(input_arg.into_iter()).unwrap();
- assert_eq!(2, arr.value_length());
- assert_eq!(5, arr.len())
- }
-}
diff --git a/arrow/src/array/array_boolean.rs b/arrow/src/array/array_boolean.rs
deleted file mode 100644
index 8b38b53..0000000
--- a/arrow/src/array/array_boolean.rs
+++ /dev/null
@@ -1,312 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-use std::borrow::Borrow;
-use std::convert::From;
-use std::iter::{FromIterator, IntoIterator};
-use std::mem;
-use std::{any::Any, fmt};
-
-use super::*;
-use super::{array::print_long_array, raw_pointer::RawPtrBox};
-use crate::buffer::{Buffer, MutableBuffer};
-use crate::util::bit_util;
-
-/// Array of bools
-///
-/// # Example
-///
-/// ```
-/// use arrow::array::{Array, BooleanArray};
-/// let arr = BooleanArray::from(vec![Some(false), Some(true), None, Some(true)]);
-/// assert_eq!(4, arr.len());
-/// assert_eq!(1, arr.null_count());
-/// assert!(arr.is_valid(0));
-/// assert!(!arr.is_null(0));
-/// assert_eq!(false, arr.value(0));
-/// assert!(arr.is_valid(1));
-/// assert!(!arr.is_null(1));
-/// assert_eq!(true, arr.value(1));
-/// assert!(!arr.is_valid(2));
-/// assert!(arr.is_null(2));
-/// assert!(arr.is_valid(3));
-/// assert!(!arr.is_null(3));
-/// assert_eq!(true, arr.value(3));
-/// ```
-///
-pub struct BooleanArray {
- data: ArrayData,
- /// Pointer to the value array. The lifetime of this must be <= to the value buffer
- /// stored in `data`, so it's safe to store.
- raw_values: RawPtrBox<u8>,
-}
-
-impl fmt::Debug for BooleanArray {
- fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
- write!(f, "BooleanArray\n[\n")?;
- print_long_array(self, f, |array, index, f| {
- fmt::Debug::fmt(&array.value(index), f)
- })?;
- write!(f, "]")
- }
-}
-
-impl BooleanArray {
- /// Returns the length of this array.
- pub fn len(&self) -> usize {
- self.data.len()
- }
-
- /// Returns whether this array is empty.
- pub fn is_empty(&self) -> bool {
- self.data.is_empty()
- }
-
- // Returns a new boolean array builder
- pub fn builder(capacity: usize) -> BooleanBuilder {
- BooleanBuilder::new(capacity)
- }
-
- /// Returns a `Buffer` holding all the values of this array.
- ///
- /// Note this doesn't take the offset of this array into account.
- pub fn values(&self) -> &Buffer {
- &self.data.buffers()[0]
- }
-
- /// Returns the boolean value at index `i`.
- ///
- /// # Safety
- /// This doesn't check bounds, the caller must ensure that index < self.len()
- pub unsafe fn value_unchecked(&self, i: usize) -> bool {
- let offset = i + self.offset();
- bit_util::get_bit_raw(self.raw_values.as_ptr(), offset)
- }
-
- /// Returns the boolean value at index `i`.
- ///
- /// Note this doesn't do any bound checking, for performance reason.
- pub fn value(&self, i: usize) -> bool {
- debug_assert!(i < self.len());
- unsafe { self.value_unchecked(i) }
- }
-}
-
-impl Array for BooleanArray {
- fn as_any(&self) -> &Any {
- self
- }
-
- fn data(&self) -> &ArrayData {
- &self.data
- }
-
- /// Returns the total number of bytes of memory occupied by the buffers owned by this [BooleanArray].
- fn get_buffer_memory_size(&self) -> usize {
- self.data.get_buffer_memory_size()
- }
-
- /// Returns the total number of bytes of memory occupied physically by this [BooleanArray].
- fn get_array_memory_size(&self) -> usize {
- self.data.get_array_memory_size() + mem::size_of_val(self)
- }
-}
-
-impl From<Vec<bool>> for BooleanArray {
- fn from(data: Vec<bool>) -> Self {
- let mut mut_buf = MutableBuffer::new_null(data.len());
- {
- let mut_slice = mut_buf.as_slice_mut();
- for (i, b) in data.iter().enumerate() {
- if *b {
- bit_util::set_bit(mut_slice, i);
- }
- }
- }
- let array_data = ArrayData::builder(DataType::Boolean)
- .len(data.len())
- .add_buffer(mut_buf.into())
- .build();
- BooleanArray::from(array_data)
- }
-}
-
-impl From<Vec<Option<bool>>> for BooleanArray {
- fn from(data: Vec<Option<bool>>) -> Self {
- BooleanArray::from_iter(data.iter())
- }
-}
-
-impl From<ArrayData> for BooleanArray {
- fn from(data: ArrayData) -> Self {
- assert_eq!(
- data.buffers().len(),
- 1,
- "BooleanArray data should contain a single buffer only (values buffer)"
- );
- let ptr = data.buffers()[0].as_ptr();
- Self {
- data,
- raw_values: unsafe { RawPtrBox::new(ptr) },
- }
- }
-}
-
-impl<'a> IntoIterator for &'a BooleanArray {
- type Item = Option<bool>;
- type IntoIter = BooleanIter<'a>;
-
- fn into_iter(self) -> Self::IntoIter {
- BooleanIter::<'a>::new(self)
- }
-}
-
-impl<'a> BooleanArray {
- /// constructs a new iterator
- pub fn iter(&'a self) -> BooleanIter<'a> {
- BooleanIter::<'a>::new(&self)
- }
-}
-
-impl<Ptr: Borrow<Option<bool>>> FromIterator<Ptr> for BooleanArray {
- fn from_iter<I: IntoIterator<Item = Ptr>>(iter: I) -> Self {
- let iter = iter.into_iter();
- let (_, data_len) = iter.size_hint();
- let data_len = data_len.expect("Iterator must be sized"); // panic if no upper bound.
-
- let num_bytes = bit_util::ceil(data_len, 8);
- let mut null_buf = MutableBuffer::from_len_zeroed(num_bytes);
- let mut val_buf = MutableBuffer::from_len_zeroed(num_bytes);
-
- let data = val_buf.as_slice_mut();
-
- let null_slice = null_buf.as_slice_mut();
- iter.enumerate().for_each(|(i, item)| {
- if let Some(a) = item.borrow() {
- bit_util::set_bit(null_slice, i);
- if *a {
- bit_util::set_bit(data, i);
- }
- }
- });
-
- let data = ArrayData::new(
- DataType::Boolean,
- data_len,
- None,
- Some(null_buf.into()),
- 0,
- vec![val_buf.into()],
- vec![],
- );
- BooleanArray::from(data)
- }
-}
-
-#[cfg(test)]
-mod tests {
- use super::*;
-
- use crate::buffer::Buffer;
- use crate::datatypes::DataType;
-
- #[test]
- fn test_boolean_fmt_debug() {
- let arr = BooleanArray::from(vec![true, false, false]);
- assert_eq!(
- "BooleanArray\n[\n true,\n false,\n false,\n]",
- format!("{:?}", arr)
- );
- }
-
- #[test]
- fn test_boolean_with_null_fmt_debug() {
- let mut builder = BooleanArray::builder(3);
- builder.append_value(true).unwrap();
- builder.append_null().unwrap();
- builder.append_value(false).unwrap();
- let arr = builder.finish();
- assert_eq!(
- "BooleanArray\n[\n true,\n null,\n false,\n]",
- format!("{:?}", arr)
- );
- }
-
- #[test]
- fn test_boolean_array_from_vec() {
- let buf = Buffer::from([10_u8]);
- let arr = BooleanArray::from(vec![false, true, false, true]);
- assert_eq!(&buf, arr.values());
- assert_eq!(4, arr.len());
- assert_eq!(0, arr.offset());
- assert_eq!(0, arr.null_count());
- for i in 0..4 {
- assert!(!arr.is_null(i));
- assert!(arr.is_valid(i));
- assert_eq!(i == 1 || i == 3, arr.value(i), "failed at {}", i)
- }
- }
-
- #[test]
- fn test_boolean_array_from_vec_option() {
- let buf = Buffer::from([10_u8]);
- let arr = BooleanArray::from(vec![Some(false), Some(true), None, Some(true)]);
- assert_eq!(&buf, arr.values());
- assert_eq!(4, arr.len());
- assert_eq!(0, arr.offset());
- assert_eq!(1, arr.null_count());
- for i in 0..4 {
- if i == 2 {
- assert!(arr.is_null(i));
- assert!(!arr.is_valid(i));
- } else {
- assert!(!arr.is_null(i));
- assert!(arr.is_valid(i));
- assert_eq!(i == 1 || i == 3, arr.value(i), "failed at {}", i)
- }
- }
- }
-
- #[test]
- fn test_boolean_array_builder() {
- // Test building a boolean array with ArrayData builder and offset
- // 000011011
- let buf = Buffer::from([27_u8]);
- let buf2 = buf.clone();
- let data = ArrayData::builder(DataType::Boolean)
- .len(5)
- .offset(2)
- .add_buffer(buf)
- .build();
- let arr = BooleanArray::from(data);
- assert_eq!(&buf2, arr.values());
- assert_eq!(5, arr.len());
- assert_eq!(2, arr.offset());
- assert_eq!(0, arr.null_count());
- for i in 0..3 {
- assert_eq!(i != 0, arr.value(i), "failed at {}", i);
- }
- }
-
- #[test]
- #[should_panic(expected = "BooleanArray data should contain a single buffer only \
- (values buffer)")]
- fn test_boolean_array_invalid_buffer_len() {
- let data = ArrayData::builder(DataType::Boolean).len(5).build();
- BooleanArray::from(data);
- }
-}
diff --git a/arrow/src/array/array_dictionary.rs b/arrow/src/array/array_dictionary.rs
deleted file mode 100644
index 168a439..0000000
--- a/arrow/src/array/array_dictionary.rs
+++ /dev/null
@@ -1,393 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-use std::any::Any;
-use std::fmt;
-use std::iter::IntoIterator;
-use std::mem;
-use std::{convert::From, iter::FromIterator};
-
-use super::{
- make_array, Array, ArrayData, ArrayRef, PrimitiveArray, PrimitiveBuilder,
- StringArray, StringBuilder, StringDictionaryBuilder,
-};
-use crate::datatypes::ArrowNativeType;
-use crate::datatypes::{ArrowDictionaryKeyType, ArrowPrimitiveType, DataType};
-
-/// A dictionary array where each element is a single value indexed by an integer key.
-/// This is mostly used to represent strings or a limited set of primitive types as integers,
-/// for example when doing NLP analysis or representing chromosomes by name.
-///
-/// Example **with nullable** data:
-///
-/// ```
-/// use arrow::array::{DictionaryArray, Int8Array};
-/// use arrow::datatypes::Int8Type;
-/// let test = vec!["a", "a", "b", "c"];
-/// let array : DictionaryArray<Int8Type> = test.iter().map(|&x| if x == "b" {None} else {Some(x)}).collect();
-/// assert_eq!(array.keys(), &Int8Array::from(vec![Some(0), Some(0), None, Some(1)]));
-/// ```
-///
-/// Example **without nullable** data:
-///
-/// ```
-/// use arrow::array::{DictionaryArray, Int8Array};
-/// use arrow::datatypes::Int8Type;
-/// let test = vec!["a", "a", "b", "c"];
-/// let array : DictionaryArray<Int8Type> = test.into_iter().collect();
-/// assert_eq!(array.keys(), &Int8Array::from(vec![0, 0, 1, 2]));
-/// ```
-pub struct DictionaryArray<K: ArrowPrimitiveType> {
- /// Data of this dictionary. Note that this is _not_ compatible with the C Data interface,
- /// as, in the current implementation, `values` below are the first child of this struct.
- data: ArrayData,
-
- /// The keys of this dictionary. These are constructed from the buffer and null bitmap
- /// of `data`.
- /// Also, note that these do not correspond to the true values of this array. Rather, they map
- /// to the real values.
- keys: PrimitiveArray<K>,
-
- /// Array of dictionary values (can by any DataType).
- values: ArrayRef,
-
- /// Values are ordered.
- is_ordered: bool,
-}
-
-impl<'a, K: ArrowPrimitiveType> DictionaryArray<K> {
- /// Return an array view of the keys of this dictionary as a PrimitiveArray.
- pub fn keys(&self) -> &PrimitiveArray<K> {
- &self.keys
- }
-
- /// Returns the lookup key by doing reverse dictionary lookup
- pub fn lookup_key(&self, value: &str) -> Option<K::Native> {
- let rd_buf: &StringArray =
- self.values.as_any().downcast_ref::<StringArray>().unwrap();
-
- (0..rd_buf.len())
- .position(|i| rd_buf.value(i) == value)
- .map(K::Native::from_usize)
- .flatten()
- }
-
- /// Returns a reference to the dictionary values array
- pub fn values(&self) -> &ArrayRef {
- &self.values
- }
-
- /// Returns a clone of the value type of this list.
- pub fn value_type(&self) -> DataType {
- self.values.data_ref().data_type().clone()
- }
-
- /// The length of the dictionary is the length of the keys array.
- pub fn len(&self) -> usize {
- self.keys.len()
- }
-
- /// Whether this dictionary is empty
- pub fn is_empty(&self) -> bool {
- self.keys.is_empty()
- }
-
- // Currently exists for compatibility purposes with Arrow IPC.
- pub fn is_ordered(&self) -> bool {
- self.is_ordered
- }
-}
-
-/// Constructs a `DictionaryArray` from an array data reference.
-impl<T: ArrowPrimitiveType> From<ArrayData> for DictionaryArray<T> {
- fn from(data: ArrayData) -> Self {
- assert_eq!(
- data.buffers().len(),
- 1,
- "DictionaryArray data should contain a single buffer only (keys)."
- );
- assert_eq!(
- data.child_data().len(),
- 1,
- "DictionaryArray should contain a single child array (values)."
- );
-
- if let DataType::Dictionary(key_data_type, _) = data.data_type() {
- if key_data_type.as_ref() != &T::DATA_TYPE {
- panic!("DictionaryArray's data type must match.")
- };
- // create a zero-copy of the keys' data
- let keys = PrimitiveArray::<T>::from(ArrayData::new(
- T::DATA_TYPE,
- data.len(),
- Some(data.null_count()),
- data.null_buffer().cloned(),
- data.offset(),
- data.buffers().to_vec(),
- vec![],
- ));
- let values = make_array(data.child_data()[0].clone());
- Self {
- data,
- keys,
- values,
- is_ordered: false,
- }
- } else {
- panic!("DictionaryArray must have Dictionary data type.")
- }
- }
-}
-
-/// Constructs a `DictionaryArray` from an iterator of optional strings.
-impl<'a, T: ArrowPrimitiveType + ArrowDictionaryKeyType> FromIterator<Option<&'a str>>
- for DictionaryArray<T>
-{
- fn from_iter<I: IntoIterator<Item = Option<&'a str>>>(iter: I) -> Self {
- let it = iter.into_iter();
- let (lower, _) = it.size_hint();
- let key_builder = PrimitiveBuilder::<T>::new(lower);
- let value_builder = StringBuilder::new(256);
- let mut builder = StringDictionaryBuilder::new(key_builder, value_builder);
- it.for_each(|i| {
- if let Some(i) = i {
- // Note: impl ... for Result<DictionaryArray<T>> fails with
- // error[E0117]: only traits defined in the current crate can be implemented for arbitrary types
- builder
- .append(i)
- .expect("Unable to append a value to a dictionary array.");
- } else {
- builder
- .append_null()
- .expect("Unable to append a null value to a dictionary array.");
- }
- });
-
- builder.finish()
- }
-}
-
-/// Constructs a `DictionaryArray` from an iterator of strings.
-impl<'a, T: ArrowPrimitiveType + ArrowDictionaryKeyType> FromIterator<&'a str>
- for DictionaryArray<T>
-{
- fn from_iter<I: IntoIterator<Item = &'a str>>(iter: I) -> Self {
- let it = iter.into_iter();
- let (lower, _) = it.size_hint();
- let key_builder = PrimitiveBuilder::<T>::new(lower);
- let value_builder = StringBuilder::new(256);
- let mut builder = StringDictionaryBuilder::new(key_builder, value_builder);
- it.for_each(|i| {
- builder
- .append(i)
- .expect("Unable to append a value to a dictionary array.");
- });
-
- builder.finish()
- }
-}
-
-impl<T: ArrowPrimitiveType> Array for DictionaryArray<T> {
- fn as_any(&self) -> &Any {
- self
- }
-
- fn data(&self) -> &ArrayData {
- &self.data
- }
-
- fn get_buffer_memory_size(&self) -> usize {
- // Since both `keys` and `values` derive (are references from) `data`, we only need to account for `data`.
- self.data.get_buffer_memory_size()
- }
-
- fn get_array_memory_size(&self) -> usize {
- self.data.get_array_memory_size()
- + self.keys.get_array_memory_size()
- + self.values.get_array_memory_size()
- + mem::size_of_val(self)
- }
-}
-
-impl<T: ArrowPrimitiveType> fmt::Debug for DictionaryArray<T> {
- fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
- writeln!(
- f,
- "DictionaryArray {{keys: {:?} values: {:?}}}",
- self.keys, self.values
- )
- }
-}
-
-#[cfg(test)]
-mod tests {
- use super::*;
-
- use crate::{
- array::Int16Array,
- datatypes::{Int32Type, Int8Type, UInt32Type, UInt8Type},
- };
- use crate::{
- array::Int16DictionaryArray, array::PrimitiveDictionaryBuilder,
- datatypes::DataType,
- };
- use crate::{buffer::Buffer, datatypes::ToByteSlice};
-
- #[test]
- fn test_dictionary_array() {
- // Construct a value array
- let value_data = ArrayData::builder(DataType::Int8)
- .len(8)
- .add_buffer(Buffer::from(
- &[10_i8, 11, 12, 13, 14, 15, 16, 17].to_byte_slice(),
- ))
- .build();
-
- // Construct a buffer for value offsets, for the nested array:
- let keys = Buffer::from(&[2_i16, 3, 4].to_byte_slice());
-
- // Construct a dictionary array from the above two
- let key_type = DataType::Int16;
- let value_type = DataType::Int8;
- let dict_data_type =
- DataType::Dictionary(Box::new(key_type), Box::new(value_type));
- let dict_data = ArrayData::builder(dict_data_type.clone())
- .len(3)
- .add_buffer(keys.clone())
- .add_child_data(value_data.clone())
- .build();
- let dict_array = Int16DictionaryArray::from(dict_data);
-
- let values = dict_array.values();
- assert_eq!(&value_data, values.data());
- assert_eq!(DataType::Int8, dict_array.value_type());
- assert_eq!(3, dict_array.len());
-
- // Null count only makes sense in terms of the component arrays.
- assert_eq!(0, dict_array.null_count());
- assert_eq!(0, dict_array.values().null_count());
- assert_eq!(dict_array.keys(), &Int16Array::from(vec![2_i16, 3, 4]));
-
- // Now test with a non-zero offset
- let dict_data = ArrayData::builder(dict_data_type)
- .len(2)
- .offset(1)
- .add_buffer(keys)
- .add_child_data(value_data.clone())
- .build();
- let dict_array = Int16DictionaryArray::from(dict_data);
-
- let values = dict_array.values();
- assert_eq!(&value_data, values.data());
- assert_eq!(DataType::Int8, dict_array.value_type());
- assert_eq!(2, dict_array.len());
- assert_eq!(dict_array.keys(), &Int16Array::from(vec![3_i16, 4]));
- }
-
- #[test]
- fn test_dictionary_array_fmt_debug() {
- let key_builder = PrimitiveBuilder::<UInt8Type>::new(3);
- let value_builder = PrimitiveBuilder::<UInt32Type>::new(2);
- let mut builder = PrimitiveDictionaryBuilder::new(key_builder, value_builder);
- builder.append(12345678).unwrap();
- builder.append_null().unwrap();
- builder.append(22345678).unwrap();
- let array = builder.finish();
- assert_eq!(
- "DictionaryArray {keys: PrimitiveArray<UInt8>\n[\n 0,\n null,\n 1,\n] values: PrimitiveArray<UInt32>\n[\n 12345678,\n 22345678,\n]}\n",
- format!("{:?}", array)
- );
-
- let key_builder = PrimitiveBuilder::<UInt8Type>::new(20);
- let value_builder = PrimitiveBuilder::<UInt32Type>::new(2);
- let mut builder = PrimitiveDictionaryBuilder::new(key_builder, value_builder);
- for _ in 0..20 {
- builder.append(1).unwrap();
- }
- let array = builder.finish();
- assert_eq!(
- "DictionaryArray {keys: PrimitiveArray<UInt8>\n[\n 0,\n 0,\n 0,\n 0,\n 0,\n 0,\n 0,\n 0,\n 0,\n 0,\n 0,\n 0,\n 0,\n 0,\n 0,\n 0,\n 0,\n 0,\n 0,\n 0,\n] values: PrimitiveArray<UInt32>\n[\n 1,\n]}\n",
- format!("{:?}", array)
- );
- }
-
- #[test]
- fn test_dictionary_array_from_iter() {
- let test = vec!["a", "a", "b", "c"];
- let array: DictionaryArray<Int8Type> = test
- .iter()
- .map(|&x| if x == "b" { None } else { Some(x) })
- .collect();
- assert_eq!(
- "DictionaryArray {keys: PrimitiveArray<Int8>\n[\n 0,\n 0,\n null,\n 1,\n] values: StringArray\n[\n \"a\",\n \"c\",\n]}\n",
- format!("{:?}", array)
- );
-
- let array: DictionaryArray<Int8Type> = test.into_iter().collect();
- assert_eq!(
- "DictionaryArray {keys: PrimitiveArray<Int8>\n[\n 0,\n 0,\n 1,\n 2,\n] values: StringArray\n[\n \"a\",\n \"b\",\n \"c\",\n]}\n",
- format!("{:?}", array)
- );
- }
-
- #[test]
- fn test_dictionary_array_reverse_lookup_key() {
- let test = vec!["a", "a", "b", "c"];
- let array: DictionaryArray<Int8Type> = test.into_iter().collect();
-
- assert_eq!(array.lookup_key("c"), Some(2));
-
- // Direction of building a dictionary is the iterator direction
- let test = vec!["t3", "t3", "t2", "t2", "t1", "t3", "t4", "t1", "t0"];
- let array: DictionaryArray<Int8Type> = test.into_iter().collect();
-
- assert_eq!(array.lookup_key("t1"), Some(2));
- assert_eq!(array.lookup_key("non-existent"), None);
- }
-
- #[test]
- fn test_dictionary_keys_as_primitive_array() {
- let test = vec!["a", "b", "c", "a"];
- let array: DictionaryArray<Int8Type> = test.into_iter().collect();
-
- let keys = array.keys();
- assert_eq!(&DataType::Int8, keys.data_type());
- assert_eq!(0, keys.null_count());
- assert_eq!(&[0, 1, 2, 0], keys.values());
- }
-
- #[test]
- fn test_dictionary_keys_as_primitive_array_with_null() {
- let test = vec![Some("a"), None, Some("b"), None, None, Some("a")];
- let array: DictionaryArray<Int32Type> = test.into_iter().collect();
-
- let keys = array.keys();
- assert_eq!(&DataType::Int32, keys.data_type());
- assert_eq!(3, keys.null_count());
-
- assert_eq!(true, keys.is_valid(0));
- assert_eq!(false, keys.is_valid(1));
- assert_eq!(true, keys.is_valid(2));
- assert_eq!(false, keys.is_valid(3));
- assert_eq!(false, keys.is_valid(4));
- assert_eq!(true, keys.is_valid(5));
-
- assert_eq!(0, keys.value(0));
- assert_eq!(1, keys.value(2));
- assert_eq!(0, keys.value(5));
- }
-}
diff --git a/arrow/src/array/array_list.rs b/arrow/src/array/array_list.rs
deleted file mode 100644
index 0e33463..0000000
--- a/arrow/src/array/array_list.rs
+++ /dev/null
@@ -1,1056 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-use std::any::Any;
-use std::fmt;
-use std::mem;
-
-use num::Num;
-
-use super::{
- array::print_long_array, make_array, raw_pointer::RawPtrBox, Array, ArrayData,
- ArrayRef, BooleanBufferBuilder, GenericListArrayIter, PrimitiveArray,
-};
-use crate::{
- buffer::MutableBuffer,
- datatypes::{ArrowNativeType, ArrowPrimitiveType, DataType, Field},
- error::ArrowError,
-};
-
-/// trait declaring an offset size, relevant for i32 vs i64 array types.
-pub trait OffsetSizeTrait: ArrowNativeType + Num + Ord + std::ops::AddAssign {
- fn is_large() -> bool;
-}
-
-impl OffsetSizeTrait for i32 {
- #[inline]
- fn is_large() -> bool {
- false
- }
-}
-
-impl OffsetSizeTrait for i64 {
- #[inline]
- fn is_large() -> bool {
- true
- }
-}
-
-pub struct GenericListArray<OffsetSize> {
- data: ArrayData,
- values: ArrayRef,
- value_offsets: RawPtrBox<OffsetSize>,
-}
-
-impl<OffsetSize: OffsetSizeTrait> GenericListArray<OffsetSize> {
- /// Returns a reference to the values of this list.
- pub fn values(&self) -> ArrayRef {
- self.values.clone()
- }
-
- /// Returns a clone of the value type of this list.
- pub fn value_type(&self) -> DataType {
- self.values.data_ref().data_type().clone()
- }
-
- /// Returns ith value of this list array.
- /// # Safety
- /// Caller must ensure that the index is within the array bounds
- pub unsafe fn value_unchecked(&self, i: usize) -> ArrayRef {
- let end = *self.value_offsets().get_unchecked(i + 1);
- let start = *self.value_offsets().get_unchecked(i);
- self.values
- .slice(start.to_usize().unwrap(), (end - start).to_usize().unwrap())
- }
-
- /// Returns ith value of this list array.
- pub fn value(&self, i: usize) -> ArrayRef {
- let end = self.value_offsets()[i + 1];
- let start = self.value_offsets()[i];
- self.values
- .slice(start.to_usize().unwrap(), (end - start).to_usize().unwrap())
- }
-
- /// Returns the offset values in the offsets buffer
- #[inline]
- pub fn value_offsets(&self) -> &[OffsetSize] {
- // Soundness
- // pointer alignment & location is ensured by RawPtrBox
- // buffer bounds/offset is ensured by the ArrayData instance.
- unsafe {
- std::slice::from_raw_parts(
- self.value_offsets.as_ptr().add(self.data.offset()),
- self.len() + 1,
- )
- }
- }
-
- /// Returns the length for value at index `i`.
- #[inline]
- pub fn value_length(&self, i: usize) -> OffsetSize {
- let offsets = self.value_offsets();
- offsets[i + 1] - offsets[i]
- }
-
- /// constructs a new iterator
- pub fn iter<'a>(&'a self) -> GenericListArrayIter<'a, OffsetSize> {
- GenericListArrayIter::<'a, OffsetSize>::new(&self)
- }
-
- #[inline]
- fn get_type(data_type: &DataType) -> Option<&DataType> {
- if OffsetSize::is_large() {
- if let DataType::LargeList(child) = data_type {
- Some(child.data_type())
- } else {
- None
- }
- } else if let DataType::List(child) = data_type {
- Some(child.data_type())
- } else {
- None
- }
- }
-
- /// Creates a [`GenericListArray`] from an iterator of primitive values
- /// # Example
- /// ```
- /// # use arrow::array::ListArray;
- /// # use arrow::datatypes::Int32Type;
- /// let data = vec![
- /// Some(vec![Some(0), Some(1), Some(2)]),
- /// None,
- /// Some(vec![Some(3), None, Some(5)]),
- /// Some(vec![Some(6), Some(7)]),
- /// ];
- /// let list_array = ListArray::from_iter_primitive::<Int32Type, _, _>(data);
- /// println!("{:?}", list_array);
- /// ```
- pub fn from_iter_primitive<T, P, I>(iter: I) -> Self
- where
- T: ArrowPrimitiveType,
- P: AsRef<[Option<<T as ArrowPrimitiveType>::Native>]>
- + IntoIterator<Item = Option<<T as ArrowPrimitiveType>::Native>>,
- I: IntoIterator<Item = Option<P>>,
- {
- let iterator = iter.into_iter();
- let (lower, _) = iterator.size_hint();
-
- let mut offsets =
- MutableBuffer::new((lower + 1) * std::mem::size_of::<OffsetSize>());
- let mut length_so_far = OffsetSize::zero();
- offsets.push(length_so_far);
-
- let mut null_buf = BooleanBufferBuilder::new(lower);
-
- let values: PrimitiveArray<T> = iterator
- .filter_map(|maybe_slice| {
- // regardless of whether the item is Some, the offsets and null buffers must be updated.
- match &maybe_slice {
- Some(x) => {
- length_so_far +=
- OffsetSize::from_usize(x.as_ref().len()).unwrap();
- null_buf.append(true);
- }
- None => null_buf.append(false),
- };
- offsets.push(length_so_far);
- maybe_slice
- })
- .flatten()
- .collect();
-
- let field = Box::new(Field::new("item", T::DATA_TYPE, true));
- let data_type = if OffsetSize::is_large() {
- DataType::LargeList(field)
- } else {
- DataType::List(field)
- };
- let data = ArrayData::builder(data_type)
- .len(null_buf.len())
- .add_buffer(offsets.into())
- .add_child_data(values.data().clone())
- .null_bit_buffer(null_buf.into())
- .build();
- Self::from(data)
- }
-}
-
-impl<OffsetSize: OffsetSizeTrait> From<ArrayData> for GenericListArray<OffsetSize> {
- fn from(data: ArrayData) -> Self {
- Self::try_new_from_array_data(data).expect(
- "Expected infallable creation of GenericListArray from ArrayDataRef failed",
- )
- }
-}
-
-impl<OffsetSize: OffsetSizeTrait> GenericListArray<OffsetSize> {
- fn try_new_from_array_data(data: ArrayData) -> Result<Self, ArrowError> {
- if data.buffers().len() != 1 {
- return Err(ArrowError::InvalidArgumentError(
- format!("ListArray data should contain a single buffer only (value offsets), had {}",
- data.len())));
- }
-
- if data.child_data().len() != 1 {
- return Err(ArrowError::InvalidArgumentError(format!(
- "ListArray should contain a single child array (values array), had {}",
- data.child_data().len()
- )));
- }
-
- let values = data.child_data()[0].clone();
-
- if let Some(child_data_type) = Self::get_type(data.data_type()) {
- if values.data_type() != child_data_type {
- return Err(ArrowError::InvalidArgumentError(format!(
- "[Large]ListArray's child datatype {:?} does not \
- correspond to the List's datatype {:?}",
- values.data_type(),
- child_data_type
- )));
- }
- } else {
- return Err(ArrowError::InvalidArgumentError(format!(
- "[Large]ListArray's datatype must be [Large]ListArray(). It is {:?}",
- data.data_type()
- )));
- }
-
- let values = make_array(values);
- let value_offsets = data.buffers()[0].as_ptr();
-
- let value_offsets = unsafe { RawPtrBox::<OffsetSize>::new(value_offsets) };
- unsafe {
- if !(*value_offsets.as_ptr().offset(0)).is_zero() {
- return Err(ArrowError::InvalidArgumentError(String::from(
- "offsets do not start at zero",
- )));
- }
- }
- Ok(Self {
- data,
- values,
- value_offsets,
- })
- }
-}
-
-impl<OffsetSize: 'static + OffsetSizeTrait> Array for GenericListArray<OffsetSize> {
- fn as_any(&self) -> &Any {
- self
- }
-
- fn data(&self) -> &ArrayData {
- &self.data
- }
-
- /// Returns the total number of bytes of memory occupied by the buffers owned by this [ListArray].
- fn get_buffer_memory_size(&self) -> usize {
- self.data.get_buffer_memory_size()
- }
-
- /// Returns the total number of bytes of memory occupied physically by this [ListArray].
- fn get_array_memory_size(&self) -> usize {
- self.data.get_array_memory_size() + mem::size_of_val(self)
- }
-}
-
-impl<OffsetSize: OffsetSizeTrait> fmt::Debug for GenericListArray<OffsetSize> {
- fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
- let prefix = if OffsetSize::is_large() { "Large" } else { "" };
-
- write!(f, "{}ListArray\n[\n", prefix)?;
- print_long_array(self, f, |array, index, f| {
- fmt::Debug::fmt(&array.value(index), f)
- })?;
- write!(f, "]")
- }
-}
-
-/// A list array where each element is a variable-sized sequence of values with the same
-/// type whose memory offsets between elements are represented by a i32.
-pub type ListArray = GenericListArray<i32>;
-
-/// A list array where each element is a variable-sized sequence of values with the same
-/// type whose memory offsets between elements are represented by a i64.
-pub type LargeListArray = GenericListArray<i64>;
-
-/// A list array where each element is a fixed-size sequence of values with the same
-/// type whose maximum length is represented by a i32.
-pub struct FixedSizeListArray {
- data: ArrayData,
- values: ArrayRef,
- length: i32,
-}
-
-impl FixedSizeListArray {
- /// Returns a reference to the values of this list.
- pub fn values(&self) -> ArrayRef {
- self.values.clone()
- }
-
- /// Returns a clone of the value type of this list.
- pub fn value_type(&self) -> DataType {
- self.values.data_ref().data_type().clone()
- }
-
- /// Returns ith value of this list array.
- pub fn value(&self, i: usize) -> ArrayRef {
- self.values
- .slice(self.value_offset(i) as usize, self.value_length() as usize)
- }
-
- /// Returns the offset for value at index `i`.
- ///
- /// Note this doesn't do any bound checking, for performance reason.
- #[inline]
- pub fn value_offset(&self, i: usize) -> i32 {
- self.value_offset_at(self.data.offset() + i)
- }
-
- /// Returns the length for value at index `i`.
- ///
- /// Note this doesn't do any bound checking, for performance reason.
- #[inline]
- pub const fn value_length(&self) -> i32 {
- self.length
- }
-
- #[inline]
- const fn value_offset_at(&self, i: usize) -> i32 {
- i as i32 * self.length
- }
-}
-
-impl From<ArrayData> for FixedSizeListArray {
- fn from(data: ArrayData) -> Self {
- assert_eq!(
- data.buffers().len(),
- 0,
- "FixedSizeListArray data should not contain a buffer for value offsets"
- );
- assert_eq!(
- data.child_data().len(),
- 1,
- "FixedSizeListArray should contain a single child array (values array)"
- );
- let values = make_array(data.child_data()[0].clone());
- let length = match data.data_type() {
- DataType::FixedSizeList(_, len) => {
- if *len > 0 {
- // check that child data is multiple of length
- assert_eq!(
- values.len() % *len as usize,
- 0,
- "FixedSizeListArray child array length should be a multiple of {}",
- len
- );
- }
-
- *len
- }
- _ => {
- panic!("FixedSizeListArray data should contain a FixedSizeList data type")
- }
- };
- Self {
- data,
- values,
- length,
- }
- }
-}
-
-impl Array for FixedSizeListArray {
- fn as_any(&self) -> &Any {
- self
- }
-
- fn data(&self) -> &ArrayData {
- &self.data
- }
-
- /// Returns the total number of bytes of memory occupied by the buffers owned by this [FixedSizeListArray].
- fn get_buffer_memory_size(&self) -> usize {
- self.data.get_buffer_memory_size() + self.values().get_buffer_memory_size()
- }
-
- /// Returns the total number of bytes of memory occupied physically by this [FixedSizeListArray].
- fn get_array_memory_size(&self) -> usize {
- self.data.get_array_memory_size()
- + self.values().get_array_memory_size()
- + mem::size_of_val(self)
- }
-}
-
-impl fmt::Debug for FixedSizeListArray {
- fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
- write!(f, "FixedSizeListArray<{}>\n[\n", self.value_length())?;
- print_long_array(self, f, |array, index, f| {
- fmt::Debug::fmt(&array.value(index), f)
- })?;
- write!(f, "]")
- }
-}
-
-#[cfg(test)]
-mod tests {
- use crate::{
- alloc,
- array::ArrayData,
- array::Int32Array,
- buffer::Buffer,
- datatypes::Field,
- datatypes::{Int32Type, ToByteSlice},
- util::bit_util,
- };
-
- use super::*;
-
- fn create_from_buffers() -> ListArray {
- // Construct a value array
- let value_data = ArrayData::builder(DataType::Int32)
- .len(8)
- .add_buffer(Buffer::from(&[0, 1, 2, 3, 4, 5, 6, 7].to_byte_slice()))
- .build();
-
- // Construct a buffer for value offsets, for the nested array:
- // [[0, 1, 2], [3, 4, 5], [6, 7]]
- let value_offsets = Buffer::from(&[0, 3, 6, 8].to_byte_slice());
-
- // Construct a list array from the above two
- let list_data_type =
- DataType::List(Box::new(Field::new("item", DataType::Int32, true)));
- let list_data = ArrayData::builder(list_data_type)
- .len(3)
- .add_buffer(value_offsets)
- .add_child_data(value_data)
- .build();
- ListArray::from(list_data)
- }
-
- #[test]
- fn test_from_iter_primitive() {
- let data = vec![
- Some(vec![Some(0), Some(1), Some(2)]),
- Some(vec![Some(3), Some(4), Some(5)]),
- Some(vec![Some(6), Some(7)]),
- ];
- let list_array = ListArray::from_iter_primitive::<Int32Type, _, _>(data);
-
- let another = create_from_buffers();
- assert_eq!(list_array, another)
- }
-
- #[test]
- fn test_list_array() {
- // Construct a value array
- let value_data = ArrayData::builder(DataType::Int32)
- .len(8)
- .add_buffer(Buffer::from_slice_ref(&[0, 1, 2, 3, 4, 5, 6, 7]))
- .build();
-
- // Construct a buffer for value offsets, for the nested array:
- // [[0, 1, 2], [3, 4, 5], [6, 7]]
- let value_offsets = Buffer::from_slice_ref(&[0, 3, 6, 8]);
-
- // Construct a list array from the above two
- let list_data_type =
- DataType::List(Box::new(Field::new("item", DataType::Int32, false)));
- let list_data = ArrayData::builder(list_data_type.clone())
- .len(3)
- .add_buffer(value_offsets.clone())
- .add_child_data(value_data.clone())
- .build();
- let list_array = ListArray::from(list_data);
-
- let values = list_array.values();
- assert_eq!(&value_data, values.data());
- assert_eq!(DataType::Int32, list_array.value_type());
- assert_eq!(3, list_array.len());
- assert_eq!(0, list_array.null_count());
- assert_eq!(6, list_array.value_offsets()[2]);
- assert_eq!(2, list_array.value_length(2));
- assert_eq!(
- 0,
- list_array
- .value(0)
- .as_any()
- .downcast_ref::<Int32Array>()
- .unwrap()
- .value(0)
- );
- assert_eq!(
- 0,
- unsafe { list_array.value_unchecked(0) }
- .as_any()
- .downcast_ref::<Int32Array>()
- .unwrap()
- .value(0)
- );
- for i in 0..3 {
- assert!(list_array.is_valid(i));
- assert!(!list_array.is_null(i));
- }
-
- // Now test with a non-zero offset
- let list_data = ArrayData::builder(list_data_type)
- .len(3)
- .offset(1)
- .add_buffer(value_offsets)
- .add_child_data(value_data.clone())
- .build();
- let list_array = ListArray::from(list_data);
-
- let values = list_array.values();
- assert_eq!(&value_data, values.data());
- assert_eq!(DataType::Int32, list_array.value_type());
- assert_eq!(3, list_array.len());
- assert_eq!(0, list_array.null_count());
- assert_eq!(6, list_array.value_offsets()[1]);
- assert_eq!(2, list_array.value_length(1));
- assert_eq!(
- 3,
- list_array
- .value(0)
- .as_any()
- .downcast_ref::<Int32Array>()
- .unwrap()
- .value(0)
- );
- assert_eq!(
- 3,
- unsafe { list_array.value_unchecked(0) }
- .as_any()
- .downcast_ref::<Int32Array>()
- .unwrap()
- .value(0)
- );
- }
-
- #[test]
- fn test_large_list_array() {
- // Construct a value array
- let value_data = ArrayData::builder(DataType::Int32)
- .len(8)
- .add_buffer(Buffer::from_slice_ref(&[0, 1, 2, 3, 4, 5, 6, 7]))
- .build();
-
- // Construct a buffer for value offsets, for the nested array:
- // [[0, 1, 2], [3, 4, 5], [6, 7]]
- let value_offsets = Buffer::from_slice_ref(&[0i64, 3, 6, 8]);
-
- // Construct a list array from the above two
- let list_data_type =
- DataType::LargeList(Box::new(Field::new("item", DataType::Int32, false)));
- let list_data = ArrayData::builder(list_data_type.clone())
- .len(3)
- .add_buffer(value_offsets.clone())
- .add_child_data(value_data.clone())
- .build();
- let list_array = LargeListArray::from(list_data);
-
- let values = list_array.values();
- assert_eq!(&value_data, values.data());
- assert_eq!(DataType::Int32, list_array.value_type());
- assert_eq!(3, list_array.len());
- assert_eq!(0, list_array.null_count());
- assert_eq!(6, list_array.value_offsets()[2]);
- assert_eq!(2, list_array.value_length(2));
- assert_eq!(
- 0,
- list_array
- .value(0)
- .as_any()
- .downcast_ref::<Int32Array>()
- .unwrap()
- .value(0)
- );
- assert_eq!(
- 0,
- unsafe { list_array.value_unchecked(0) }
- .as_any()
- .downcast_ref::<Int32Array>()
- .unwrap()
- .value(0)
- );
- for i in 0..3 {
- assert!(list_array.is_valid(i));
- assert!(!list_array.is_null(i));
- }
-
- // Now test with a non-zero offset
- let list_data = ArrayData::builder(list_data_type)
- .len(3)
- .offset(1)
- .add_buffer(value_offsets)
- .add_child_data(value_data.clone())
- .build();
- let list_array = LargeListArray::from(list_data);
-
- let values = list_array.values();
- assert_eq!(&value_data, values.data());
- assert_eq!(DataType::Int32, list_array.value_type());
- assert_eq!(3, list_array.len());
- assert_eq!(0, list_array.null_count());
- assert_eq!(6, list_array.value_offsets()[1]);
- assert_eq!(2, list_array.value_length(1));
- assert_eq!(
- 3,
- list_array
- .value(0)
- .as_any()
- .downcast_ref::<Int32Array>()
- .unwrap()
- .value(0)
- );
- assert_eq!(
- 3,
- unsafe { list_array.value_unchecked(0) }
- .as_any()
- .downcast_ref::<Int32Array>()
- .unwrap()
- .value(0)
- );
- }
-
- #[test]
- fn test_fixed_size_list_array() {
- // Construct a value array
- let value_data = ArrayData::builder(DataType::Int32)
- .len(9)
- .add_buffer(Buffer::from_slice_ref(&[0, 1, 2, 3, 4, 5, 6, 7, 8]))
- .build();
-
- // Construct a list array from the above two
- let list_data_type = DataType::FixedSizeList(
- Box::new(Field::new("item", DataType::Int32, false)),
- 3,
- );
- let list_data = ArrayData::builder(list_data_type.clone())
- .len(3)
- .add_child_data(value_data.clone())
- .build();
- let list_array = FixedSizeListArray::from(list_data);
-
- let values = list_array.values();
- assert_eq!(&value_data, values.data());
- assert_eq!(DataType::Int32, list_array.value_type());
- assert_eq!(3, list_array.len());
- assert_eq!(0, list_array.null_count());
- assert_eq!(6, list_array.value_offset(2));
- assert_eq!(3, list_array.value_length());
- assert_eq!(
- 0,
- list_array
- .value(0)
- .as_any()
- .downcast_ref::<Int32Array>()
- .unwrap()
- .value(0)
- );
- for i in 0..3 {
- assert!(list_array.is_valid(i));
- assert!(!list_array.is_null(i));
- }
-
- // Now test with a non-zero offset
- let list_data = ArrayData::builder(list_data_type)
- .len(3)
- .offset(1)
- .add_child_data(value_data.clone())
- .build();
- let list_array = FixedSizeListArray::from(list_data);
-
- let values = list_array.values();
- assert_eq!(&value_data, values.data());
- assert_eq!(DataType::Int32, list_array.value_type());
- assert_eq!(3, list_array.len());
- assert_eq!(0, list_array.null_count());
- assert_eq!(
- 3,
- list_array
- .value(0)
- .as_any()
- .downcast_ref::<Int32Array>()
- .unwrap()
- .value(0)
- );
- assert_eq!(6, list_array.value_offset(1));
- assert_eq!(3, list_array.value_length());
- }
-
- #[test]
- #[should_panic(
- expected = "FixedSizeListArray child array length should be a multiple of 3"
- )]
- fn test_fixed_size_list_array_unequal_children() {
- // Construct a value array
- let value_data = ArrayData::builder(DataType::Int32)
- .len(8)
- .add_buffer(Buffer::from_slice_ref(&[0, 1, 2, 3, 4, 5, 6, 7]))
- .build();
-
- // Construct a list array from the above two
- let list_data_type = DataType::FixedSizeList(
- Box::new(Field::new("item", DataType::Int32, false)),
- 3,
- );
- let list_data = ArrayData::builder(list_data_type)
- .len(3)
- .add_child_data(value_data)
- .build();
- FixedSizeListArray::from(list_data);
- }
-
- #[test]
- fn test_list_array_slice() {
- // Construct a value array
- let value_data = ArrayData::builder(DataType::Int32)
- .len(10)
- .add_buffer(Buffer::from_slice_ref(&[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]))
- .build();
-
- // Construct a buffer for value offsets, for the nested array:
- // [[0, 1], null, null, [2, 3], [4, 5], null, [6, 7, 8], null, [9]]
- let value_offsets = Buffer::from_slice_ref(&[0, 2, 2, 2, 4, 6, 6, 9, 9, 10]);
- // 01011001 00000001
- let mut null_bits: [u8; 2] = [0; 2];
- bit_util::set_bit(&mut null_bits, 0);
- bit_util::set_bit(&mut null_bits, 3);
- bit_util::set_bit(&mut null_bits, 4);
- bit_util::set_bit(&mut null_bits, 6);
- bit_util::set_bit(&mut null_bits, 8);
-
- // Construct a list array from the above two
- let list_data_type =
- DataType::List(Box::new(Field::new("item", DataType::Int32, false)));
- let list_data = ArrayData::builder(list_data_type)
- .len(9)
- .add_buffer(value_offsets)
- .add_child_data(value_data.clone())
- .null_bit_buffer(Buffer::from(null_bits))
- .build();
- let list_array = ListArray::from(list_data);
-
- let values = list_array.values();
- assert_eq!(&value_data, values.data());
- assert_eq!(DataType::Int32, list_array.value_type());
- assert_eq!(9, list_array.len());
- assert_eq!(4, list_array.null_count());
- assert_eq!(2, list_array.value_offsets()[3]);
- assert_eq!(2, list_array.value_length(3));
-
- let sliced_array = list_array.slice(1, 6);
- assert_eq!(6, sliced_array.len());
- assert_eq!(1, sliced_array.offset());
- assert_eq!(3, sliced_array.null_count());
-
- for i in 0..sliced_array.len() {
- if bit_util::get_bit(&null_bits, sliced_array.offset() + i) {
- assert!(sliced_array.is_valid(i));
- } else {
- assert!(sliced_array.is_null(i));
- }
- }
-
- // Check offset and length for each non-null value.
- let sliced_list_array =
- sliced_array.as_any().downcast_ref::<ListArray>().unwrap();
- assert_eq!(2, sliced_list_array.value_offsets()[2]);
- assert_eq!(2, sliced_list_array.value_length(2));
- assert_eq!(4, sliced_list_array.value_offsets()[3]);
- assert_eq!(2, sliced_list_array.value_length(3));
- assert_eq!(6, sliced_list_array.value_offsets()[5]);
- assert_eq!(3, sliced_list_array.value_length(5));
- }
-
- #[test]
- fn test_large_list_array_slice() {
- // Construct a value array
- let value_data = ArrayData::builder(DataType::Int32)
- .len(10)
- .add_buffer(Buffer::from_slice_ref(&[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]))
- .build();
-
- // Construct a buffer for value offsets, for the nested array:
- // [[0, 1], null, null, [2, 3], [4, 5], null, [6, 7, 8], null, [9]]
- let value_offsets = Buffer::from_slice_ref(&[0i64, 2, 2, 2, 4, 6, 6, 9, 9, 10]);
- // 01011001 00000001
- let mut null_bits: [u8; 2] = [0; 2];
- bit_util::set_bit(&mut null_bits, 0);
- bit_util::set_bit(&mut null_bits, 3);
- bit_util::set_bit(&mut null_bits, 4);
- bit_util::set_bit(&mut null_bits, 6);
- bit_util::set_bit(&mut null_bits, 8);
-
- // Construct a list array from the above two
- let list_data_type =
- DataType::LargeList(Box::new(Field::new("item", DataType::Int32, false)));
- let list_data = ArrayData::builder(list_data_type)
- .len(9)
- .add_buffer(value_offsets)
- .add_child_data(value_data.clone())
- .null_bit_buffer(Buffer::from(null_bits))
- .build();
- let list_array = LargeListArray::from(list_data);
-
- let values = list_array.values();
- assert_eq!(&value_data, values.data());
- assert_eq!(DataType::Int32, list_array.value_type());
- assert_eq!(9, list_array.len());
- assert_eq!(4, list_array.null_count());
- assert_eq!(2, list_array.value_offsets()[3]);
- assert_eq!(2, list_array.value_length(3));
-
- let sliced_array = list_array.slice(1, 6);
- assert_eq!(6, sliced_array.len());
- assert_eq!(1, sliced_array.offset());
- assert_eq!(3, sliced_array.null_count());
-
- for i in 0..sliced_array.len() {
- if bit_util::get_bit(&null_bits, sliced_array.offset() + i) {
- assert!(sliced_array.is_valid(i));
- } else {
- assert!(sliced_array.is_null(i));
- }
- }
-
- // Check offset and length for each non-null value.
- let sliced_list_array = sliced_array
- .as_any()
- .downcast_ref::<LargeListArray>()
- .unwrap();
- assert_eq!(2, sliced_list_array.value_offsets()[2]);
- assert_eq!(2, sliced_list_array.value_length(2));
- assert_eq!(4, sliced_list_array.value_offsets()[3]);
- assert_eq!(2, sliced_list_array.value_length(3));
- assert_eq!(6, sliced_list_array.value_offsets()[5]);
- assert_eq!(3, sliced_list_array.value_length(5));
- }
-
- #[test]
- #[should_panic(expected = "index out of bounds: the len is 10 but the index is 11")]
- fn test_list_array_index_out_of_bound() {
- // Construct a value array
- let value_data = ArrayData::builder(DataType::Int32)
- .len(10)
- .add_buffer(Buffer::from_slice_ref(&[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]))
- .build();
-
- // Construct a buffer for value offsets, for the nested array:
- // [[0, 1], null, null, [2, 3], [4, 5], null, [6, 7, 8], null, [9]]
- let value_offsets = Buffer::from_slice_ref(&[0i64, 2, 2, 2, 4, 6, 6, 9, 9, 10]);
- // 01011001 00000001
- let mut null_bits: [u8; 2] = [0; 2];
- bit_util::set_bit(&mut null_bits, 0);
- bit_util::set_bit(&mut null_bits, 3);
- bit_util::set_bit(&mut null_bits, 4);
- bit_util::set_bit(&mut null_bits, 6);
- bit_util::set_bit(&mut null_bits, 8);
-
- // Construct a list array from the above two
- let list_data_type =
- DataType::LargeList(Box::new(Field::new("item", DataType::Int32, false)));
- let list_data = ArrayData::builder(list_data_type)
- .len(9)
- .add_buffer(value_offsets)
- .add_child_data(value_data)
- .null_bit_buffer(Buffer::from(null_bits))
- .build();
- let list_array = LargeListArray::from(list_data);
- assert_eq!(9, list_array.len());
-
- list_array.value(10);
- }
-
- #[test]
- fn test_fixed_size_list_array_slice() {
- // Construct a value array
- let value_data = ArrayData::builder(DataType::Int32)
- .len(10)
- .add_buffer(Buffer::from_slice_ref(&[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]))
- .build();
-
- // Set null buts for the nested array:
- // [[0, 1], null, null, [6, 7], [8, 9]]
- // 01011001 00000001
- let mut null_bits: [u8; 1] = [0; 1];
- bit_util::set_bit(&mut null_bits, 0);
- bit_util::set_bit(&mut null_bits, 3);
- bit_util::set_bit(&mut null_bits, 4);
-
- // Construct a fixed size list array from the above two
- let list_data_type = DataType::FixedSizeList(
- Box::new(Field::new("item", DataType::Int32, false)),
- 2,
- );
- let list_data = ArrayData::builder(list_data_type)
- .len(5)
- .add_child_data(value_data.clone())
- .null_bit_buffer(Buffer::from(null_bits))
- .build();
- let list_array = FixedSizeListArray::from(list_data);
-
- let values = list_array.values();
- assert_eq!(&value_data, values.data());
- assert_eq!(DataType::Int32, list_array.value_type());
- assert_eq!(5, list_array.len());
- assert_eq!(2, list_array.null_count());
- assert_eq!(6, list_array.value_offset(3));
- assert_eq!(2, list_array.value_length());
-
- let sliced_array = list_array.slice(1, 4);
- assert_eq!(4, sliced_array.len());
- assert_eq!(1, sliced_array.offset());
- assert_eq!(2, sliced_array.null_count());
-
- for i in 0..sliced_array.len() {
- if bit_util::get_bit(&null_bits, sliced_array.offset() + i) {
- assert!(sliced_array.is_valid(i));
- } else {
- assert!(sliced_array.is_null(i));
- }
- }
-
- // Check offset and length for each non-null value.
- let sliced_list_array = sliced_array
- .as_any()
- .downcast_ref::<FixedSizeListArray>()
- .unwrap();
- assert_eq!(2, sliced_list_array.value_length());
- assert_eq!(6, sliced_list_array.value_offset(2));
- assert_eq!(8, sliced_list_array.value_offset(3));
- }
-
- #[test]
- #[should_panic(expected = "assertion failed: (offset + length) <= self.len()")]
- fn test_fixed_size_list_array_index_out_of_bound() {
- // Construct a value array
- let value_data = ArrayData::builder(DataType::Int32)
- .len(10)
- .add_buffer(Buffer::from_slice_ref(&[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]))
- .build();
-
- // Set null buts for the nested array:
- // [[0, 1], null, null, [6, 7], [8, 9]]
- // 01011001 00000001
- let mut null_bits: [u8; 1] = [0; 1];
- bit_util::set_bit(&mut null_bits, 0);
- bit_util::set_bit(&mut null_bits, 3);
- bit_util::set_bit(&mut null_bits, 4);
-
- // Construct a fixed size list array from the above two
- let list_data_type = DataType::FixedSizeList(
- Box::new(Field::new("item", DataType::Int32, false)),
- 2,
- );
- let list_data = ArrayData::builder(list_data_type)
- .len(5)
- .add_child_data(value_data)
- .null_bit_buffer(Buffer::from(null_bits))
- .build();
- let list_array = FixedSizeListArray::from(list_data);
-
- list_array.value(10);
- }
-
- #[test]
- #[should_panic(
- expected = "ListArray data should contain a single buffer only (value offsets)"
- )]
- fn test_list_array_invalid_buffer_len() {
- let value_data = ArrayData::builder(DataType::Int32)
- .len(8)
- .add_buffer(Buffer::from_slice_ref(&[0, 1, 2, 3, 4, 5, 6, 7]))
- .build();
- let list_data_type =
- DataType::List(Box::new(Field::new("item", DataType::Int32, false)));
- let list_data = ArrayData::builder(list_data_type)
- .len(3)
- .add_child_data(value_data)
- .build();
- ListArray::from(list_data);
- }
-
- #[test]
- #[should_panic(
- expected = "ListArray should contain a single child array (values array)"
- )]
- fn test_list_array_invalid_child_array_len() {
- let value_offsets = Buffer::from_slice_ref(&[0, 2, 5, 7]);
- let list_data_type =
- DataType::List(Box::new(Field::new("item", DataType::Int32, false)));
- let list_data = ArrayData::builder(list_data_type)
- .len(3)
- .add_buffer(value_offsets)
- .build();
- ListArray::from(list_data);
- }
-
- #[test]
- #[should_panic(expected = "offsets do not start at zero")]
- fn test_list_array_invalid_value_offset_start() {
- let value_data = ArrayData::builder(DataType::Int32)
- .len(8)
- .add_buffer(Buffer::from_slice_ref(&[0, 1, 2, 3, 4, 5, 6, 7]))
- .build();
-
- let value_offsets = Buffer::from_slice_ref(&[2, 2, 5, 7]);
-
- let list_data_type =
- DataType::List(Box::new(Field::new("item", DataType::Int32, false)));
- let list_data = ArrayData::builder(list_data_type)
- .len(3)
- .add_buffer(value_offsets)
- .add_child_data(value_data)
- .build();
- ListArray::from(list_data);
- }
-
- #[test]
- #[should_panic(expected = "memory is not aligned")]
- fn test_primitive_array_alignment() {
- let ptr = alloc::allocate_aligned::<u8>(8);
- let buf = unsafe { Buffer::from_raw_parts(ptr, 8, 8) };
- let buf2 = buf.slice(1);
- let array_data = ArrayData::builder(DataType::Int32).add_buffer(buf2).build();
- Int32Array::from(array_data);
- }
-
- #[test]
- #[should_panic(expected = "memory is not aligned")]
- fn test_list_array_alignment() {
- let ptr = alloc::allocate_aligned::<u8>(8);
- let buf = unsafe { Buffer::from_raw_parts(ptr, 8, 8) };
- let buf2 = buf.slice(1);
-
- let values: [i32; 8] = [0; 8];
- let value_data = ArrayData::builder(DataType::Int32)
- .add_buffer(Buffer::from_slice_ref(&values))
- .build();
-
- let list_data_type =
- DataType::List(Box::new(Field::new("item", DataType::Int32, false)));
- let list_data = ArrayData::builder(list_data_type)
- .add_buffer(buf2)
- .add_child_data(value_data)
- .build();
- ListArray::from(list_data);
- }
-}
diff --git a/arrow/src/array/array_primitive.rs b/arrow/src/array/array_primitive.rs
deleted file mode 100644
index 9101865..0000000
--- a/arrow/src/array/array_primitive.rs
+++ /dev/null
@@ -1,944 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-use std::any::Any;
-use std::borrow::Borrow;
-use std::convert::From;
-use std::fmt;
-use std::iter::{FromIterator, IntoIterator};
-use std::mem;
-
-use chrono::{prelude::*, Duration};
-
-use super::array::print_long_array;
-use super::raw_pointer::RawPtrBox;
-use super::*;
-use crate::temporal_conversions;
-use crate::util::bit_util;
-use crate::{
- buffer::{Buffer, MutableBuffer},
- util::trusted_len_unzip,
-};
-
-/// Number of seconds in a day
-const SECONDS_IN_DAY: i64 = 86_400;
-/// Number of milliseconds in a second
-const MILLISECONDS: i64 = 1_000;
-/// Number of microseconds in a second
-const MICROSECONDS: i64 = 1_000_000;
-/// Number of nanoseconds in a second
-const NANOSECONDS: i64 = 1_000_000_000;
-
-/// Array whose elements are of primitive types.
-pub struct PrimitiveArray<T: ArrowPrimitiveType> {
- /// Underlying ArrayData
- /// # Safety
- /// must have exactly one buffer, aligned to type T
- data: ArrayData,
- /// Pointer to the value array. The lifetime of this must be <= to the value buffer
- /// stored in `data`, so it's safe to store.
- /// # Safety
- /// raw_values must have a value equivalent to `data.buffers()[0].raw_data()`
- /// raw_values must have alignment for type T::NativeType
- raw_values: RawPtrBox<T::Native>,
-}
-
-impl<T: ArrowPrimitiveType> PrimitiveArray<T> {
- /// Returns the length of this array.
- #[inline]
- pub fn len(&self) -> usize {
- self.data.len()
- }
-
- /// Returns whether this array is empty.
- pub fn is_empty(&self) -> bool {
- self.data.is_empty()
- }
-
- /// Returns a slice of the values of this array
- #[inline]
- pub fn values(&self) -> &[T::Native] {
- // Soundness
- // raw_values alignment & location is ensured by fn from(ArrayDataRef)
- // buffer bounds/offset is ensured by the ArrayData instance.
- unsafe {
- std::slice::from_raw_parts(
- self.raw_values.as_ptr().add(self.data.offset()),
- self.len(),
- )
- }
- }
-
- // Returns a new primitive array builder
- pub fn builder(capacity: usize) -> PrimitiveBuilder<T> {
- PrimitiveBuilder::<T>::new(capacity)
- }
-
- /// Returns the primitive value at index `i`.
- ///
- /// # Safety
- ///
- /// caller must ensure that the passed in offset is less than the array len()
- #[inline]
- pub unsafe fn value_unchecked(&self, i: usize) -> T::Native {
- let offset = i + self.offset();
- *self.raw_values.as_ptr().add(offset)
- }
-
- /// Returns the primitive value at index `i`.
- ///
- /// Note this doesn't do any bound checking, for performance reason.
- /// # Safety
- /// caller must ensure that the passed in offset is less than the array len()
- #[inline]
- pub fn value(&self, i: usize) -> T::Native {
- debug_assert!(i < self.len());
- unsafe { self.value_unchecked(i) }
- }
-
- /// Creates a PrimitiveArray based on an iterator of values without nulls
- pub fn from_iter_values<I: IntoIterator<Item = T::Native>>(iter: I) -> Self {
- let val_buf: Buffer = iter.into_iter().collect();
- let data = ArrayData::new(
- T::DATA_TYPE,
- val_buf.len() / mem::size_of::<<T as ArrowPrimitiveType>::Native>(),
- None,
- None,
- 0,
- vec![val_buf],
- vec![],
- );
- PrimitiveArray::from(data)
- }
-
- /// Creates a PrimitiveArray based on a constant value with `count` elements
- pub fn from_value(value: T::Native, count: usize) -> Self {
- // # Safety: length is known
- let val_buf = unsafe { Buffer::from_trusted_len_iter((0..count).map(|_| value)) };
- let data = ArrayData::new(
- T::DATA_TYPE,
- val_buf.len() / mem::size_of::<<T as ArrowPrimitiveType>::Native>(),
- None,
- None,
- 0,
- vec![val_buf],
- vec![],
- );
- PrimitiveArray::from(data)
- }
-}
-
-impl<T: ArrowPrimitiveType> Array for PrimitiveArray<T> {
- fn as_any(&self) -> &Any {
- self
- }
-
- fn data(&self) -> &ArrayData {
- &self.data
- }
-
- /// Returns the total number of bytes of memory occupied by the buffers owned by this [PrimitiveArray].
- fn get_buffer_memory_size(&self) -> usize {
- self.data.get_buffer_memory_size()
- }
-
- /// Returns the total number of bytes of memory occupied physically by this [PrimitiveArray].
- fn get_array_memory_size(&self) -> usize {
- self.data.get_array_memory_size() + mem::size_of::<RawPtrBox<T::Native>>()
- }
-}
-
-fn as_datetime<T: ArrowPrimitiveType>(v: i64) -> Option<NaiveDateTime> {
- match T::DATA_TYPE {
- DataType::Date32 => Some(temporal_conversions::date32_to_datetime(v as i32)),
- DataType::Date64 => Some(temporal_conversions::date64_to_datetime(v)),
- DataType::Time32(_) | DataType::Time64(_) => None,
- DataType::Timestamp(unit, _) => match unit {
- TimeUnit::Second => Some(temporal_conversions::timestamp_s_to_datetime(v)),
- TimeUnit::Millisecond => {
- Some(temporal_conversions::timestamp_ms_to_datetime(v))
- }
- TimeUnit::Microsecond => {
- Some(temporal_conversions::timestamp_us_to_datetime(v))
- }
- TimeUnit::Nanosecond => {
- Some(temporal_conversions::timestamp_ns_to_datetime(v))
- }
- },
- // interval is not yet fully documented [ARROW-3097]
- DataType::Interval(_) => None,
- _ => None,
- }
-}
-
-fn as_date<T: ArrowPrimitiveType>(v: i64) -> Option<NaiveDate> {
- as_datetime::<T>(v).map(|datetime| datetime.date())
-}
-
-fn as_time<T: ArrowPrimitiveType>(v: i64) -> Option<NaiveTime> {
- match T::DATA_TYPE {
- DataType::Time32(unit) => {
- // safe to immediately cast to u32 as `self.value(i)` is positive i32
- let v = v as u32;
- match unit {
- TimeUnit::Second => Some(temporal_conversions::time32s_to_time(v as i32)),
- TimeUnit::Millisecond => {
- Some(temporal_conversions::time32ms_to_time(v as i32))
- }
- _ => None,
- }
- }
- DataType::Time64(unit) => match unit {
- TimeUnit::Microsecond => Some(temporal_conversions::time64us_to_time(v)),
- TimeUnit::Nanosecond => Some(temporal_conversions::time64ns_to_time(v)),
- _ => None,
- },
- DataType::Timestamp(_, _) => as_datetime::<T>(v).map(|datetime| datetime.time()),
- DataType::Date32 | DataType::Date64 => Some(NaiveTime::from_hms(0, 0, 0)),
- DataType::Interval(_) => None,
- _ => None,
- }
-}
-
-fn as_duration<T: ArrowPrimitiveType>(v: i64) -> Option<Duration> {
- match T::DATA_TYPE {
- DataType::Duration(unit) => match unit {
- TimeUnit::Second => Some(temporal_conversions::duration_s_to_duration(v)),
- TimeUnit::Millisecond => {
- Some(temporal_conversions::duration_ms_to_duration(v))
- }
- TimeUnit::Microsecond => {
- Some(temporal_conversions::duration_us_to_duration(v))
- }
- TimeUnit::Nanosecond => {
- Some(temporal_conversions::duration_ns_to_duration(v))
- }
- },
- _ => None,
- }
-}
-
-impl<T: ArrowTemporalType + ArrowNumericType> PrimitiveArray<T>
-where
- i64: std::convert::From<T::Native>,
-{
- /// Returns value as a chrono `NaiveDateTime`, handling time resolution
- ///
- /// If a data type cannot be converted to `NaiveDateTime`, a `None` is returned.
- /// A valid value is expected, thus the user should first check for validity.
- pub fn value_as_datetime(&self, i: usize) -> Option<NaiveDateTime> {
- as_datetime::<T>(i64::from(self.value(i)))
- }
-
- /// Returns value as a chrono `NaiveDate` by using `Self::datetime()`
- ///
- /// If a data type cannot be converted to `NaiveDate`, a `None` is returned
- pub fn value_as_date(&self, i: usize) -> Option<NaiveDate> {
- self.value_as_datetime(i).map(|datetime| datetime.date())
- }
-
- /// Returns a value as a chrono `NaiveTime`
- ///
- /// `Date32` and `Date64` return UTC midnight as they do not have time resolution
- pub fn value_as_time(&self, i: usize) -> Option<NaiveTime> {
- as_time::<T>(i64::from(self.value(i)))
- }
-
- /// Returns a value as a chrono `Duration`
- ///
- /// If a data type cannot be converted to `Duration`, a `None` is returned
- pub fn value_as_duration(&self, i: usize) -> Option<Duration> {
- as_duration::<T>(i64::from(self.value(i)))
- }
-}
-
-impl<T: ArrowPrimitiveType> fmt::Debug for PrimitiveArray<T> {
- fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
- write!(f, "PrimitiveArray<{:?}>\n[\n", T::DATA_TYPE)?;
- print_long_array(self, f, |array, index, f| match T::DATA_TYPE {
- DataType::Date32 | DataType::Date64 => {
- let v = self.value(index).to_isize().unwrap() as i64;
- match as_date::<T>(v) {
- Some(date) => write!(f, "{:?}", date),
- None => write!(f, "null"),
- }
- }
- DataType::Time32(_) | DataType::Time64(_) => {
- let v = self.value(index).to_isize().unwrap() as i64;
- match as_time::<T>(v) {
- Some(time) => write!(f, "{:?}", time),
- None => write!(f, "null"),
- }
- }
- DataType::Timestamp(_, _) => {
- let v = self.value(index).to_isize().unwrap() as i64;
- match as_datetime::<T>(v) {
- Some(datetime) => write!(f, "{:?}", datetime),
- None => write!(f, "null"),
- }
- }
- _ => fmt::Debug::fmt(&array.value(index), f),
- })?;
- write!(f, "]")
- }
-}
-
-impl<'a, T: ArrowPrimitiveType> IntoIterator for &'a PrimitiveArray<T> {
- type Item = Option<<T as ArrowPrimitiveType>::Native>;
- type IntoIter = PrimitiveIter<'a, T>;
-
- fn into_iter(self) -> Self::IntoIter {
- PrimitiveIter::<'a, T>::new(self)
- }
-}
-
-impl<'a, T: ArrowPrimitiveType> PrimitiveArray<T> {
- /// constructs a new iterator
- pub fn iter(&'a self) -> PrimitiveIter<'a, T> {
- PrimitiveIter::<'a, T>::new(&self)
- }
-}
-
-impl<T: ArrowPrimitiveType, Ptr: Borrow<Option<<T as ArrowPrimitiveType>::Native>>>
- FromIterator<Ptr> for PrimitiveArray<T>
-{
- fn from_iter<I: IntoIterator<Item = Ptr>>(iter: I) -> Self {
- let iter = iter.into_iter();
- let (lower, _) = iter.size_hint();
-
- let mut null_buf = BooleanBufferBuilder::new(lower);
-
- let buffer: Buffer = iter
- .map(|item| {
- if let Some(a) = item.borrow() {
- null_buf.append(true);
- *a
- } else {
- null_buf.append(false);
- // this ensures that null items on the buffer are not arbitrary.
- // This is important because falible operations can use null values (e.g. a vectorized "add")
- // which may panic (e.g. overflow if the number on the slots happen to be very large).
- T::Native::default()
- }
- })
- .collect();
-
- let data = ArrayData::new(
- T::DATA_TYPE,
- null_buf.len(),
- None,
- Some(null_buf.into()),
- 0,
- vec![buffer],
- vec![],
- );
- PrimitiveArray::from(data)
- }
-}
-
-impl<T: ArrowPrimitiveType> PrimitiveArray<T> {
- /// Creates a [`PrimitiveArray`] from an iterator of trusted length.
- /// # Safety
- /// The iterator must be [`TrustedLen`](https://doc.rust-lang.org/std/iter/trait.TrustedLen.html).
- /// I.e. that `size_hint().1` correctly reports its length.
- #[inline]
- pub unsafe fn from_trusted_len_iter<I, P>(iter: I) -> Self
- where
- P: std::borrow::Borrow<Option<<T as ArrowPrimitiveType>::Native>>,
- I: IntoIterator<Item = P>,
- {
- let iterator = iter.into_iter();
- let (_, upper) = iterator.size_hint();
- let len = upper.expect("trusted_len_unzip requires an upper limit");
-
- let (null, buffer) = trusted_len_unzip(iterator);
-
- let data =
- ArrayData::new(T::DATA_TYPE, len, None, Some(null), 0, vec![buffer], vec![]);
- PrimitiveArray::from(data)
- }
-}
-
-// TODO: the macro is needed here because we'd get "conflicting implementations" error
-// otherwise with both `From<Vec<T::Native>>` and `From<Vec<Option<T::Native>>>`.
-// We should revisit this in future.
-macro_rules! def_numeric_from_vec {
- ( $ty:ident ) => {
- impl From<Vec<<$ty as ArrowPrimitiveType>::Native>> for PrimitiveArray<$ty> {
- fn from(data: Vec<<$ty as ArrowPrimitiveType>::Native>) -> Self {
- let array_data = ArrayData::builder($ty::DATA_TYPE)
- .len(data.len())
- .add_buffer(Buffer::from_slice_ref(&data))
- .build();
- PrimitiveArray::from(array_data)
- }
- }
-
- // Constructs a primitive array from a vector. Should only be used for testing.
- impl From<Vec<Option<<$ty as ArrowPrimitiveType>::Native>>>
- for PrimitiveArray<$ty>
- {
- fn from(data: Vec<Option<<$ty as ArrowPrimitiveType>::Native>>) -> Self {
- PrimitiveArray::from_iter(data.iter())
- }
- }
- };
-}
-
-def_numeric_from_vec!(Int8Type);
-def_numeric_from_vec!(Int16Type);
-def_numeric_from_vec!(Int32Type);
-def_numeric_from_vec!(Int64Type);
-def_numeric_from_vec!(UInt8Type);
-def_numeric_from_vec!(UInt16Type);
-def_numeric_from_vec!(UInt32Type);
-def_numeric_from_vec!(UInt64Type);
-def_numeric_from_vec!(Float32Type);
-def_numeric_from_vec!(Float64Type);
-
-def_numeric_from_vec!(Date32Type);
-def_numeric_from_vec!(Date64Type);
-def_numeric_from_vec!(Time32SecondType);
-def_numeric_from_vec!(Time32MillisecondType);
-def_numeric_from_vec!(Time64MicrosecondType);
-def_numeric_from_vec!(Time64NanosecondType);
-def_numeric_from_vec!(IntervalYearMonthType);
-def_numeric_from_vec!(IntervalDayTimeType);
-def_numeric_from_vec!(DurationSecondType);
-def_numeric_from_vec!(DurationMillisecondType);
-def_numeric_from_vec!(DurationMicrosecondType);
-def_numeric_from_vec!(DurationNanosecondType);
-def_numeric_from_vec!(TimestampSecondType);
-def_numeric_from_vec!(TimestampMillisecondType);
-def_numeric_from_vec!(TimestampMicrosecondType);
-def_numeric_from_vec!(TimestampNanosecondType);
-
-impl<T: ArrowTimestampType> PrimitiveArray<T> {
- /// Construct a timestamp array from a vec of i64 values and an optional timezone
- pub fn from_vec(data: Vec<i64>, timezone: Option<String>) -> Self {
- let array_data =
- ArrayData::builder(DataType::Timestamp(T::get_time_unit(), timezone))
- .len(data.len())
- .add_buffer(Buffer::from_slice_ref(&data))
- .build();
- PrimitiveArray::from(array_data)
- }
-}
-
-impl<T: ArrowTimestampType> PrimitiveArray<T> {
- /// Construct a timestamp array from a vec of Option<i64> values and an optional timezone
- pub fn from_opt_vec(data: Vec<Option<i64>>, timezone: Option<String>) -> Self {
- // TODO: duplicated from def_numeric_from_vec! macro, it looks possible to convert to generic
- let data_len = data.len();
- let mut null_buf = MutableBuffer::new_null(data_len);
- let mut val_buf = MutableBuffer::new(data_len * mem::size_of::<i64>());
-
- {
- let null_slice = null_buf.as_slice_mut();
- for (i, v) in data.iter().enumerate() {
- if let Some(n) = v {
- bit_util::set_bit(null_slice, i);
- val_buf.push(*n);
- } else {
- val_buf.push(0i64);
- }
- }
- }
-
- let array_data =
- ArrayData::builder(DataType::Timestamp(T::get_time_unit(), timezone))
- .len(data_len)
- .add_buffer(val_buf.into())
- .null_bit_buffer(null_buf.into())
- .build();
- PrimitiveArray::from(array_data)
- }
-}
-
-/// Constructs a `PrimitiveArray` from an array data reference.
-impl<T: ArrowPrimitiveType> From<ArrayData> for PrimitiveArray<T> {
- fn from(data: ArrayData) -> Self {
- assert_eq!(
- data.buffers().len(),
- 1,
- "PrimitiveArray data should contain a single buffer only (values buffer)"
- );
-
- let ptr = data.buffers()[0].as_ptr();
- Self {
- data,
- raw_values: unsafe { RawPtrBox::new(ptr) },
- }
- }
-}
-
-#[cfg(test)]
-mod tests {
- use super::*;
-
- use std::thread;
-
- use crate::buffer::Buffer;
- use crate::datatypes::DataType;
-
- #[test]
- fn test_primitive_array_from_vec() {
- let buf = Buffer::from_slice_ref(&[0, 1, 2, 3, 4]);
- let arr = Int32Array::from(vec![0, 1, 2, 3, 4]);
- assert_eq!(buf, arr.data.buffers()[0]);
- assert_eq!(5, arr.len());
- assert_eq!(0, arr.offset());
- assert_eq!(0, arr.null_count());
- for i in 0..5 {
- assert!(!arr.is_null(i));
- assert!(arr.is_valid(i));
- assert_eq!(i as i32, arr.value(i));
- }
-
- assert_eq!(64, arr.get_buffer_memory_size());
- assert_eq!(136, arr.get_array_memory_size());
- }
-
- #[test]
- fn test_primitive_array_from_vec_option() {
- // Test building a primitive array with null values
- let arr = Int32Array::from(vec![Some(0), None, Some(2), None, Some(4)]);
- assert_eq!(5, arr.len());
- assert_eq!(0, arr.offset());
- assert_eq!(2, arr.null_count());
- for i in 0..5 {
- if i % 2 == 0 {
- assert!(!arr.is_null(i));
- assert!(arr.is_valid(i));
- assert_eq!(i as i32, arr.value(i));
- } else {
- assert!(arr.is_null(i));
- assert!(!arr.is_valid(i));
- }
- }
-
- assert_eq!(128, arr.get_buffer_memory_size());
- assert_eq!(216, arr.get_array_memory_size());
- }
-
- #[test]
- fn test_date64_array_from_vec_option() {
- // Test building a primitive array with null values
- // we use Int32 and Int64 as a backing array, so all Int32 and Int64 conventions
- // work
- let arr: PrimitiveArray<Date64Type> =
- vec![Some(1550902545147), None, Some(1550902545147)].into();
- assert_eq!(3, arr.len());
- assert_eq!(0, arr.offset());
- assert_eq!(1, arr.null_count());
- for i in 0..3 {
- if i % 2 == 0 {
- assert!(!arr.is_null(i));
- assert!(arr.is_valid(i));
- assert_eq!(1550902545147, arr.value(i));
- // roundtrip to and from datetime
- assert_eq!(
- 1550902545147,
- arr.value_as_datetime(i).unwrap().timestamp_millis()
- );
- } else {
- assert!(arr.is_null(i));
- assert!(!arr.is_valid(i));
- }
- }
- }
-
- #[test]
- fn test_time32_millisecond_array_from_vec() {
- // 1: 00:00:00.001
- // 37800005: 10:30:00.005
- // 86399210: 23:59:59.210
- let arr: PrimitiveArray<Time32MillisecondType> =
- vec![1, 37_800_005, 86_399_210].into();
- assert_eq!(3, arr.len());
- assert_eq!(0, arr.offset());
- assert_eq!(0, arr.null_count());
- let formatted = vec!["00:00:00.001", "10:30:00.005", "23:59:59.210"];
- for (i, formatted) in formatted.iter().enumerate().take(3) {
- // check that we can't create dates or datetimes from time instances
- assert_eq!(None, arr.value_as_datetime(i));
- assert_eq!(None, arr.value_as_date(i));
- let time = arr.value_as_time(i).unwrap();
- assert_eq!(*formatted, time.format("%H:%M:%S%.3f").to_string());
- }
- }
-
- #[test]
- fn test_time64_nanosecond_array_from_vec() {
- // Test building a primitive array with null values
- // we use Int32 and Int64 as a backing array, so all Int32 and Int64 conventions
- // work
-
- // 1e6: 00:00:00.001
- // 37800005e6: 10:30:00.005
- // 86399210e6: 23:59:59.210
- let arr: PrimitiveArray<Time64NanosecondType> =
- vec![1_000_000, 37_800_005_000_000, 86_399_210_000_000].into();
- assert_eq!(3, arr.len());
- assert_eq!(0, arr.offset());
- assert_eq!(0, arr.null_count());
- let formatted = vec!["00:00:00.001", "10:30:00.005", "23:59:59.210"];
- for (i, item) in formatted.iter().enumerate().take(3) {
- // check that we can't create dates or datetimes from time instances
- assert_eq!(None, arr.value_as_datetime(i));
- assert_eq!(None, arr.value_as_date(i));
- let time = arr.value_as_time(i).unwrap();
- assert_eq!(*item, time.format("%H:%M:%S%.3f").to_string());
- }
- }
-
- #[test]
- fn test_interval_array_from_vec() {
- // intervals are currently not treated specially, but are Int32 and Int64 arrays
- let arr = IntervalYearMonthArray::from(vec![Some(1), None, Some(-5)]);
- assert_eq!(3, arr.len());
- assert_eq!(0, arr.offset());
- assert_eq!(1, arr.null_count());
- assert_eq!(1, arr.value(0));
- assert_eq!(1, arr.values()[0]);
- assert!(arr.is_null(1));
- assert_eq!(-5, arr.value(2));
- assert_eq!(-5, arr.values()[2]);
-
- // a day_time interval contains days and milliseconds, but we do not yet have accessors for the values
- let arr = IntervalDayTimeArray::from(vec![Some(1), None, Some(-5)]);
- assert_eq!(3, arr.len());
- assert_eq!(0, arr.offset());
- assert_eq!(1, arr.null_count());
- assert_eq!(1, arr.value(0));
- assert_eq!(1, arr.values()[0]);
- assert!(arr.is_null(1));
- assert_eq!(-5, arr.value(2));
- assert_eq!(-5, arr.values()[2]);
- }
-
- #[test]
- fn test_duration_array_from_vec() {
- let arr = DurationSecondArray::from(vec![Some(1), None, Some(-5)]);
- assert_eq!(3, arr.len());
- assert_eq!(0, arr.offset());
- assert_eq!(1, arr.null_count());
- assert_eq!(1, arr.value(0));
- assert_eq!(1, arr.values()[0]);
- assert!(arr.is_null(1));
- assert_eq!(-5, arr.value(2));
- assert_eq!(-5, arr.values()[2]);
-
- let arr = DurationMillisecondArray::from(vec![Some(1), None, Some(-5)]);
- assert_eq!(3, arr.len());
- assert_eq!(0, arr.offset());
- assert_eq!(1, arr.null_count());
- assert_eq!(1, arr.value(0));
- assert_eq!(1, arr.values()[0]);
- assert!(arr.is_null(1));
- assert_eq!(-5, arr.value(2));
- assert_eq!(-5, arr.values()[2]);
-
- let arr = DurationMicrosecondArray::from(vec![Some(1), None, Some(-5)]);
- assert_eq!(3, arr.len());
- assert_eq!(0, arr.offset());
- assert_eq!(1, arr.null_count());
- assert_eq!(1, arr.value(0));
- assert_eq!(1, arr.values()[0]);
- assert!(arr.is_null(1));
- assert_eq!(-5, arr.value(2));
- assert_eq!(-5, arr.values()[2]);
-
- let arr = DurationNanosecondArray::from(vec![Some(1), None, Some(-5)]);
- assert_eq!(3, arr.len());
- assert_eq!(0, arr.offset());
- assert_eq!(1, arr.null_count());
- assert_eq!(1, arr.value(0));
- assert_eq!(1, arr.values()[0]);
- assert!(arr.is_null(1));
- assert_eq!(-5, arr.value(2));
- assert_eq!(-5, arr.values()[2]);
- }
-
- #[test]
- fn test_timestamp_array_from_vec() {
- let arr = TimestampSecondArray::from_vec(vec![1, -5], None);
- assert_eq!(2, arr.len());
- assert_eq!(0, arr.offset());
- assert_eq!(0, arr.null_count());
- assert_eq!(1, arr.value(0));
- assert_eq!(-5, arr.value(1));
- assert_eq!(&[1, -5], arr.values());
-
- let arr = TimestampMillisecondArray::from_vec(vec![1, -5], None);
- assert_eq!(2, arr.len());
- assert_eq!(0, arr.offset());
- assert_eq!(0, arr.null_count());
- assert_eq!(1, arr.value(0));
- assert_eq!(-5, arr.value(1));
- assert_eq!(&[1, -5], arr.values());
-
- let arr = TimestampMicrosecondArray::from_vec(vec![1, -5], None);
- assert_eq!(2, arr.len());
- assert_eq!(0, arr.offset());
- assert_eq!(0, arr.null_count());
- assert_eq!(1, arr.value(0));
- assert_eq!(-5, arr.value(1));
- assert_eq!(&[1, -5], arr.values());
-
- let arr = TimestampNanosecondArray::from_vec(vec![1, -5], None);
- assert_eq!(2, arr.len());
- assert_eq!(0, arr.offset());
- assert_eq!(0, arr.null_count());
- assert_eq!(1, arr.value(0));
- assert_eq!(-5, arr.value(1));
- assert_eq!(&[1, -5], arr.values());
- }
-
- #[test]
- fn test_primitive_array_slice() {
- let arr = Int32Array::from(vec![
- Some(0),
- None,
- Some(2),
- None,
- Some(4),
- Some(5),
- Some(6),
- None,
- None,
- ]);
- assert_eq!(9, arr.len());
- assert_eq!(0, arr.offset());
- assert_eq!(4, arr.null_count());
-
- let arr2 = arr.slice(2, 5);
- assert_eq!(5, arr2.len());
- assert_eq!(2, arr2.offset());
- assert_eq!(1, arr2.null_count());
-
- for i in 0..arr2.len() {
- assert_eq!(i == 1, arr2.is_null(i));
- assert_eq!(i != 1, arr2.is_valid(i));
- }
- let int_arr2 = arr2.as_any().downcast_ref::<Int32Array>().unwrap();
- assert_eq!(2, int_arr2.values()[0]);
- assert_eq!(&[4, 5, 6], &int_arr2.values()[2..5]);
-
- let arr3 = arr2.slice(2, 3);
- assert_eq!(3, arr3.len());
- assert_eq!(4, arr3.offset());
- assert_eq!(0, arr3.null_count());
-
- let int_arr3 = arr3.as_any().downcast_ref::<Int32Array>().unwrap();
- assert_eq!(&[4, 5, 6], int_arr3.values());
- assert_eq!(4, int_arr3.value(0));
- assert_eq!(5, int_arr3.value(1));
- assert_eq!(6, int_arr3.value(2));
- }
-
- #[test]
- fn test_boolean_array_slice() {
- let arr = BooleanArray::from(vec![
- Some(true),
- None,
- Some(false),
- None,
- Some(true),
- Some(false),
- Some(true),
- Some(false),
- None,
- Some(true),
- ]);
-
- assert_eq!(10, arr.len());
- assert_eq!(0, arr.offset());
- assert_eq!(3, arr.null_count());
-
- let arr2 = arr.slice(3, 5);
- assert_eq!(5, arr2.len());
- assert_eq!(3, arr2.offset());
- assert_eq!(1, arr2.null_count());
-
- let bool_arr = arr2.as_any().downcast_ref::<BooleanArray>().unwrap();
-
- assert_eq!(false, bool_arr.is_valid(0));
-
- assert_eq!(true, bool_arr.is_valid(1));
- assert_eq!(true, bool_arr.value(1));
-
- assert_eq!(true, bool_arr.is_valid(2));
- assert_eq!(false, bool_arr.value(2));
-
- assert_eq!(true, bool_arr.is_valid(3));
- assert_eq!(true, bool_arr.value(3));
-
- assert_eq!(true, bool_arr.is_valid(4));
- assert_eq!(false, bool_arr.value(4));
- }
-
- #[test]
- fn test_int32_fmt_debug() {
- let arr = Int32Array::from(vec![0, 1, 2, 3, 4]);
- assert_eq!(
- "PrimitiveArray<Int32>\n[\n 0,\n 1,\n 2,\n 3,\n 4,\n]",
- format!("{:?}", arr)
- );
- }
-
- #[test]
- fn test_fmt_debug_up_to_20_elements() {
- (1..=20).for_each(|i| {
- let values = (0..i).collect::<Vec<i16>>();
- let array_expected = format!(
- "PrimitiveArray<Int16>\n[\n{}\n]",
- values
- .iter()
- .map(|v| { format!(" {},", v) })
- .collect::<Vec<String>>()
- .join("\n")
- );
- let array = Int16Array::from(values);
-
- assert_eq!(array_expected, format!("{:?}", array));
- })
- }
-
- #[test]
- fn test_int32_with_null_fmt_debug() {
- let mut builder = Int32Array::builder(3);
- builder.append_slice(&[0, 1]).unwrap();
- builder.append_null().unwrap();
- builder.append_slice(&[3, 4]).unwrap();
- let arr = builder.finish();
- assert_eq!(
- "PrimitiveArray<Int32>\n[\n 0,\n 1,\n null,\n 3,\n 4,\n]",
- format!("{:?}", arr)
- );
- }
-
- #[test]
- fn test_timestamp_fmt_debug() {
- let arr: PrimitiveArray<TimestampMillisecondType> =
- TimestampMillisecondArray::from_vec(
- vec![1546214400000, 1546214400000, -1546214400000],
- None,
- );
- assert_eq!(
- "PrimitiveArray<Timestamp(Millisecond, None)>\n[\n 2018-12-31T00:00:00,\n 2018-12-31T00:00:00,\n 1921-01-02T00:00:00,\n]",
- format!("{:?}", arr)
- );
- }
-
- #[test]
- fn test_date32_fmt_debug() {
- let arr: PrimitiveArray<Date32Type> = vec![12356, 13548, -365].into();
- assert_eq!(
- "PrimitiveArray<Date32>\n[\n 2003-10-31,\n 2007-02-04,\n 1969-01-01,\n]",
- format!("{:?}", arr)
- );
- }
-
- #[test]
- fn test_time32second_fmt_debug() {
- let arr: PrimitiveArray<Time32SecondType> = vec![7201, 60054].into();
- assert_eq!(
- "PrimitiveArray<Time32(Second)>\n[\n 02:00:01,\n 16:40:54,\n]",
- format!("{:?}", arr)
- );
- }
-
- #[test]
- #[should_panic(expected = "invalid time")]
- fn test_time32second_invalid_neg() {
- // The panic should come from chrono, not from arrow
- let arr: PrimitiveArray<Time32SecondType> = vec![-7201, -60054].into();
- println!("{:?}", arr);
- }
-
- #[test]
- fn test_primitive_array_builder() {
- // Test building a primitive array with ArrayData builder and offset
- let buf = Buffer::from_slice_ref(&[0, 1, 2, 3, 4]);
- let buf2 = buf.clone();
- let data = ArrayData::builder(DataType::Int32)
- .len(5)
- .offset(2)
- .add_buffer(buf)
- .build();
- let arr = Int32Array::from(data);
- assert_eq!(buf2, arr.data.buffers()[0]);
- assert_eq!(5, arr.len());
- assert_eq!(0, arr.null_count());
- for i in 0..3 {
- assert_eq!((i + 2) as i32, arr.value(i));
- }
- }
-
- #[test]
- fn test_primitive_from_iter_values() {
- // Test building a primitive array with from_iter_values
- let arr: PrimitiveArray<Int32Type> = PrimitiveArray::from_iter_values(0..10);
- assert_eq!(10, arr.len());
- assert_eq!(0, arr.null_count());
- for i in 0..10i32 {
- assert_eq!(i, arr.value(i as usize));
- }
- }
-
- #[test]
- fn test_primitive_array_from_unbound_iter() {
- // iterator that doesn't declare (upper) size bound
- let value_iter = (0..)
- .scan(0usize, |pos, i| {
- if *pos < 10 {
- *pos += 1;
- Some(Some(i))
- } else {
- // actually returns up to 10 values
- None
- }
- })
- // limited using take()
- .take(100);
-
- let (_, upper_size_bound) = value_iter.size_hint();
- // the upper bound, defined by take above, is 100
- assert_eq!(upper_size_bound, Some(100));
- let primitive_array: PrimitiveArray<Int32Type> = value_iter.collect();
- // but the actual number of items in the array should be 10
- assert_eq!(primitive_array.len(), 10);
- }
-
- #[test]
- #[should_panic(expected = "PrimitiveArray data should contain a single buffer only \
- (values buffer)")]
- fn test_primitive_array_invalid_buffer_len() {
- let data = ArrayData::builder(DataType::Int32).len(5).build();
- Int32Array::from(data);
- }
-
- #[test]
- fn test_access_array_concurrently() {
- let a = Int32Array::from(vec![5, 6, 7, 8, 9]);
- let ret = thread::spawn(move || a.value(3)).join();
-
- assert!(ret.is_ok());
- assert_eq!(8, ret.ok().unwrap());
- }
-}
diff --git a/arrow/src/array/array_string.rs b/arrow/src/array/array_string.rs
deleted file mode 100644
index 8b0ad10..0000000
--- a/arrow/src/array/array_string.rs
+++ /dev/null
@@ -1,544 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-use std::convert::From;
-use std::fmt;
-use std::mem;
-use std::{any::Any, iter::FromIterator};
-
-use super::{
- array::print_long_array, raw_pointer::RawPtrBox, Array, ArrayData, GenericListArray,
- GenericStringIter, OffsetSizeTrait,
-};
-use crate::buffer::Buffer;
-use crate::util::bit_util;
-use crate::{buffer::MutableBuffer, datatypes::DataType};
-
-/// Like OffsetSizeTrait, but specialized for Strings
-// This allow us to expose a constant datatype for the GenericStringArray
-pub trait StringOffsetSizeTrait: OffsetSizeTrait {
- const DATA_TYPE: DataType;
-}
-
-impl StringOffsetSizeTrait for i32 {
- const DATA_TYPE: DataType = DataType::Utf8;
-}
-
-impl StringOffsetSizeTrait for i64 {
- const DATA_TYPE: DataType = DataType::LargeUtf8;
-}
-
-/// Generic struct for \[Large\]StringArray
-pub struct GenericStringArray<OffsetSize: StringOffsetSizeTrait> {
- data: ArrayData,
- value_offsets: RawPtrBox<OffsetSize>,
- value_data: RawPtrBox<u8>,
-}
-
-impl<OffsetSize: StringOffsetSizeTrait> GenericStringArray<OffsetSize> {
- /// Returns the length for the element at index `i`.
- #[inline]
- pub fn value_length(&self, i: usize) -> OffsetSize {
- let offsets = self.value_offsets();
- offsets[i + 1] - offsets[i]
- }
-
- /// Returns the offset values in the offsets buffer
- #[inline]
- pub fn value_offsets(&self) -> &[OffsetSize] {
- // Soundness
- // pointer alignment & location is ensured by RawPtrBox
- // buffer bounds/offset is ensured by the ArrayData instance.
- unsafe {
- std::slice::from_raw_parts(
- self.value_offsets.as_ptr().add(self.data.offset()),
- self.len() + 1,
- )
- }
- }
-
- /// Returns a clone of the value data buffer
- pub fn value_data(&self) -> Buffer {
- self.data.buffers()[1].clone()
- }
-
- /// Returns the element at index
- /// # Safety
- /// caller is responsible for ensuring that index is within the array bounds
- pub unsafe fn value_unchecked(&self, i: usize) -> &str {
- let end = self.value_offsets().get_unchecked(i + 1);
- let start = self.value_offsets().get_unchecked(i);
-
- // Soundness
- // pointer alignment & location is ensured by RawPtrBox
- // buffer bounds/offset is ensured by the value_offset invariants
- // ISSUE: utf-8 well formedness is not checked
-
- // Safety of `to_isize().unwrap()`
- // `start` and `end` are &OffsetSize, which is a generic type that implements the
- // OffsetSizeTrait. Currently, only i32 and i64 implement OffsetSizeTrait,
- // both of which should cleanly cast to isize on an architecture that supports
- // 32/64-bit offsets
- let slice = std::slice::from_raw_parts(
- self.value_data.as_ptr().offset(start.to_isize().unwrap()),
- (*end - *start).to_usize().unwrap(),
- );
- std::str::from_utf8_unchecked(slice)
- }
-
- /// Returns the element at index `i` as &str
- pub fn value(&self, i: usize) -> &str {
- assert!(i < self.data.len(), "StringArray out of bounds access");
- //Soundness: length checked above, offset buffer length is 1 larger than logical array length
- let end = unsafe { self.value_offsets().get_unchecked(i + 1) };
- let start = unsafe { self.value_offsets().get_unchecked(i) };
-
- // Soundness
- // pointer alignment & location is ensured by RawPtrBox
- // buffer bounds/offset is ensured by the value_offset invariants
- // ISSUE: utf-8 well formedness is not checked
- unsafe {
- // Safety of `to_isize().unwrap()`
- // `start` and `end` are &OffsetSize, which is a generic type that implements the
- // OffsetSizeTrait. Currently, only i32 and i64 implement OffsetSizeTrait,
- // both of which should cleanly cast to isize on an architecture that supports
- // 32/64-bit offsets
- let slice = std::slice::from_raw_parts(
- self.value_data.as_ptr().offset(start.to_isize().unwrap()),
- (*end - *start).to_usize().unwrap(),
- );
- std::str::from_utf8_unchecked(slice)
- }
- }
-
- fn from_list(v: GenericListArray<OffsetSize>) -> Self {
- assert_eq!(
- v.data().child_data()[0].child_data().len(),
- 0,
- "StringArray can only be created from list array of u8 values \
- (i.e. List<PrimitiveArray<u8>>)."
- );
- assert_eq!(
- v.data().child_data()[0].data_type(),
- &DataType::UInt8,
- "StringArray can only be created from List<u8> arrays, mismatched data types."
- );
-
- let mut builder = ArrayData::builder(OffsetSize::DATA_TYPE)
- .len(v.len())
- .add_buffer(v.data().buffers()[0].clone())
- .add_buffer(v.data().child_data()[0].buffers()[0].clone());
- if let Some(bitmap) = v.data().null_bitmap() {
- builder = builder.null_bit_buffer(bitmap.bits.clone())
- }
-
- let data = builder.build();
- Self::from(data)
- }
-
- pub(crate) fn from_vec(v: Vec<&str>) -> Self {
- let mut offsets =
- MutableBuffer::new((v.len() + 1) * std::mem::size_of::<OffsetSize>());
- let mut values = MutableBuffer::new(0);
-
- let mut length_so_far = OffsetSize::zero();
- offsets.push(length_so_far);
-
- for s in &v {
- length_so_far += OffsetSize::from_usize(s.len()).unwrap();
- offsets.push(length_so_far);
- values.extend_from_slice(s.as_bytes());
- }
- let array_data = ArrayData::builder(OffsetSize::DATA_TYPE)
- .len(v.len())
- .add_buffer(offsets.into())
- .add_buffer(values.into())
- .build();
- Self::from(array_data)
- }
-
- pub(crate) fn from_opt_vec(v: Vec<Option<&str>>) -> Self {
- v.into_iter().collect()
- }
-
- /// Creates a `GenericStringArray` based on an iterator of values without nulls
- pub fn from_iter_values<Ptr, I: IntoIterator<Item = Ptr>>(iter: I) -> Self
- where
- Ptr: AsRef<str>,
- {
- let iter = iter.into_iter();
- let (_, data_len) = iter.size_hint();
- let data_len = data_len.expect("Iterator must be sized"); // panic if no upper bound.
-
- let mut offsets =
- MutableBuffer::new((data_len + 1) * std::mem::size_of::<OffsetSize>());
- let mut values = MutableBuffer::new(0);
-
- let mut length_so_far = OffsetSize::zero();
- offsets.push(length_so_far);
-
- for i in iter {
- let s = i.as_ref();
- length_so_far += OffsetSize::from_usize(s.len()).unwrap();
- offsets.push(length_so_far);
- values.extend_from_slice(s.as_bytes());
- }
- let array_data = ArrayData::builder(OffsetSize::DATA_TYPE)
- .len(data_len)
- .add_buffer(offsets.into())
- .add_buffer(values.into())
- .build();
- Self::from(array_data)
- }
-}
-
-impl<'a, Ptr, OffsetSize: StringOffsetSizeTrait> FromIterator<Option<Ptr>>
- for GenericStringArray<OffsetSize>
-where
- Ptr: AsRef<str>,
-{
- fn from_iter<I: IntoIterator<Item = Option<Ptr>>>(iter: I) -> Self {
- let iter = iter.into_iter();
- let (_, data_len) = iter.size_hint();
- let data_len = data_len.expect("Iterator must be sized"); // panic if no upper bound.
-
- let offset_size = std::mem::size_of::<OffsetSize>();
- let mut offsets = MutableBuffer::new((data_len + 1) * offset_size);
- let mut values = MutableBuffer::new(0);
- let mut null_buf = MutableBuffer::new_null(data_len);
- let null_slice = null_buf.as_slice_mut();
- let mut length_so_far = OffsetSize::zero();
- offsets.push(length_so_far);
-
- for (i, s) in iter.enumerate() {
- let value_bytes = if let Some(ref s) = s {
- // set null bit
- bit_util::set_bit(null_slice, i);
- let s_bytes = s.as_ref().as_bytes();
- length_so_far += OffsetSize::from_usize(s_bytes.len()).unwrap();
- s_bytes
- } else {
- b""
- };
- values.extend_from_slice(value_bytes);
- offsets.push(length_so_far);
- }
-
- // calculate actual data_len, which may be different from the iterator's upper bound
- let data_len = (offsets.len() / offset_size) - 1;
- let array_data = ArrayData::builder(OffsetSize::DATA_TYPE)
- .len(data_len)
- .add_buffer(offsets.into())
- .add_buffer(values.into())
- .null_bit_buffer(null_buf.into())
- .build();
- Self::from(array_data)
- }
-}
-
-impl<'a, T: StringOffsetSizeTrait> IntoIterator for &'a GenericStringArray<T> {
- type Item = Option<&'a str>;
- type IntoIter = GenericStringIter<'a, T>;
-
- fn into_iter(self) -> Self::IntoIter {
- GenericStringIter::<'a, T>::new(self)
- }
-}
-
-impl<'a, T: StringOffsetSizeTrait> GenericStringArray<T> {
- /// constructs a new iterator
- pub fn iter(&'a self) -> GenericStringIter<'a, T> {
- GenericStringIter::<'a, T>::new(&self)
- }
-}
-
-impl<OffsetSize: StringOffsetSizeTrait> fmt::Debug for GenericStringArray<OffsetSize> {
- fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
- let prefix = if OffsetSize::is_large() { "Large" } else { "" };
-
- write!(f, "{}StringArray\n[\n", prefix)?;
- print_long_array(self, f, |array, index, f| {
- fmt::Debug::fmt(&array.value(index), f)
- })?;
- write!(f, "]")
- }
-}
-
-impl<OffsetSize: StringOffsetSizeTrait> Array for GenericStringArray<OffsetSize> {
- fn as_any(&self) -> &Any {
- self
- }
-
- fn data(&self) -> &ArrayData {
- &self.data
- }
-
- /// Returns the total number of bytes of memory occupied by the buffers owned by this [$name].
- fn get_buffer_memory_size(&self) -> usize {
- self.data.get_buffer_memory_size()
- }
-
- /// Returns the total number of bytes of memory occupied physically by this [$name].
- fn get_array_memory_size(&self) -> usize {
- self.data.get_array_memory_size() + mem::size_of_val(self)
- }
-}
-
-impl<OffsetSize: StringOffsetSizeTrait> From<ArrayData>
- for GenericStringArray<OffsetSize>
-{
- fn from(data: ArrayData) -> Self {
- assert_eq!(
- data.data_type(),
- &<OffsetSize as StringOffsetSizeTrait>::DATA_TYPE,
- "[Large]StringArray expects Datatype::[Large]Utf8"
- );
- assert_eq!(
- data.buffers().len(),
- 2,
- "StringArray data should contain 2 buffers only (offsets and values)"
- );
- let offsets = data.buffers()[0].as_ptr();
- let values = data.buffers()[1].as_ptr();
- Self {
- data,
- value_offsets: unsafe { RawPtrBox::new(offsets) },
- value_data: unsafe { RawPtrBox::new(values) },
- }
- }
-}
-
-impl<OffsetSize: StringOffsetSizeTrait> From<Vec<Option<&str>>>
- for GenericStringArray<OffsetSize>
-{
- fn from(v: Vec<Option<&str>>) -> Self {
- GenericStringArray::<OffsetSize>::from_opt_vec(v)
- }
-}
-
-impl<OffsetSize: StringOffsetSizeTrait> From<Vec<&str>>
- for GenericStringArray<OffsetSize>
-{
- fn from(v: Vec<&str>) -> Self {
- GenericStringArray::<OffsetSize>::from_vec(v)
- }
-}
-
-/// An array where each element is a variable-sized sequence of bytes representing a string
-/// whose maximum length (in bytes) is represented by a i32.
-///
-/// Example
-///
-/// ```
-/// use arrow::array::StringArray;
-/// let array = StringArray::from(vec![Some("foo"), None, Some("bar")]);
-/// assert_eq!(array.value(0), "foo");
-/// ```
-pub type StringArray = GenericStringArray<i32>;
-
-/// An array where each element is a variable-sized sequence of bytes representing a string
-/// whose maximum length (in bytes) is represented by a i64.
-///
-/// Example
-///
-/// ```
-/// use arrow::array::LargeStringArray;
-/// let array = LargeStringArray::from(vec![Some("foo"), None, Some("bar")]);
-/// assert_eq!(array.value(2), "bar");
-/// ```
-pub type LargeStringArray = GenericStringArray<i64>;
-
-impl<T: StringOffsetSizeTrait> From<GenericListArray<T>> for GenericStringArray<T> {
- fn from(v: GenericListArray<T>) -> Self {
- GenericStringArray::<T>::from_list(v)
- }
-}
-
-#[cfg(test)]
-mod tests {
- use crate::array::{ListBuilder, StringBuilder};
-
- use super::*;
-
- #[test]
- fn test_string_array_from_u8_slice() {
- let values: Vec<&str> = vec!["hello", "", "parquet"];
-
- // Array data: ["hello", "", "parquet"]
- let string_array = StringArray::from(values);
-
- assert_eq!(3, string_array.len());
- assert_eq!(0, string_array.null_count());
- assert_eq!("hello", string_array.value(0));
- assert_eq!("hello", unsafe { string_array.value_unchecked(0) });
- assert_eq!("", string_array.value(1));
- assert_eq!("", unsafe { string_array.value_unchecked(1) });
- assert_eq!("parquet", string_array.value(2));
- assert_eq!("parquet", unsafe { string_array.value_unchecked(2) });
- assert_eq!(5, string_array.value_offsets()[2]);
- assert_eq!(7, string_array.value_length(2));
- for i in 0..3 {
- assert!(string_array.is_valid(i));
- assert!(!string_array.is_null(i));
- }
- }
-
- #[test]
- #[should_panic(expected = "[Large]StringArray expects Datatype::[Large]Utf8")]
- fn test_string_array_from_int() {
- let array = LargeStringArray::from(vec!["a", "b"]);
- StringArray::from(array.data().clone());
- }
-
- #[test]
- fn test_large_string_array_from_u8_slice() {
- let values: Vec<&str> = vec!["hello", "", "parquet"];
-
- // Array data: ["hello", "", "parquet"]
- let string_array = LargeStringArray::from(values);
-
- assert_eq!(3, string_array.len());
- assert_eq!(0, string_array.null_count());
- assert_eq!("hello", string_array.value(0));
- assert_eq!("hello", unsafe { string_array.value_unchecked(0) });
- assert_eq!("", string_array.value(1));
- assert_eq!("", unsafe { string_array.value_unchecked(1) });
- assert_eq!("parquet", string_array.value(2));
- assert_eq!("parquet", unsafe { string_array.value_unchecked(2) });
- assert_eq!(5, string_array.value_offsets()[2]);
- assert_eq!(7, string_array.value_length(2));
- for i in 0..3 {
- assert!(string_array.is_valid(i));
- assert!(!string_array.is_null(i));
- }
- }
-
- #[test]
- fn test_nested_string_array() {
- let string_builder = StringBuilder::new(3);
- let mut list_of_string_builder = ListBuilder::new(string_builder);
-
- list_of_string_builder.values().append_value("foo").unwrap();
- list_of_string_builder.values().append_value("bar").unwrap();
- list_of_string_builder.append(true).unwrap();
-
- list_of_string_builder
- .values()
- .append_value("foobar")
- .unwrap();
- list_of_string_builder.append(true).unwrap();
- let list_of_strings = list_of_string_builder.finish();
-
- assert_eq!(list_of_strings.len(), 2);
-
- let first_slot = list_of_strings.value(0);
- let first_list = first_slot.as_any().downcast_ref::<StringArray>().unwrap();
- assert_eq!(first_list.len(), 2);
- assert_eq!(first_list.value(0), "foo");
- assert_eq!(unsafe { first_list.value_unchecked(0) }, "foo");
- assert_eq!(first_list.value(1), "bar");
- assert_eq!(unsafe { first_list.value_unchecked(1) }, "bar");
-
- let second_slot = list_of_strings.value(1);
- let second_list = second_slot.as_any().downcast_ref::<StringArray>().unwrap();
- assert_eq!(second_list.len(), 1);
- assert_eq!(second_list.value(0), "foobar");
- assert_eq!(unsafe { second_list.value_unchecked(0) }, "foobar");
- }
-
- #[test]
- #[should_panic(expected = "StringArray out of bounds access")]
- fn test_string_array_get_value_index_out_of_bound() {
- let values: [u8; 12] = [
- b'h', b'e', b'l', b'l', b'o', b'p', b'a', b'r', b'q', b'u', b'e', b't',
- ];
- let offsets: [i32; 4] = [0, 5, 5, 12];
- let array_data = ArrayData::builder(DataType::Utf8)
- .len(3)
- .add_buffer(Buffer::from_slice_ref(&offsets))
- .add_buffer(Buffer::from_slice_ref(&values))
- .build();
- let string_array = StringArray::from(array_data);
- string_array.value(4);
- }
-
- #[test]
- fn test_string_array_fmt_debug() {
- let arr: StringArray = vec!["hello", "arrow"].into();
- assert_eq!(
- "StringArray\n[\n \"hello\",\n \"arrow\",\n]",
- format!("{:?}", arr)
- );
- }
-
- #[test]
- fn test_large_string_array_fmt_debug() {
- let arr: LargeStringArray = vec!["hello", "arrow"].into();
- assert_eq!(
- "LargeStringArray\n[\n \"hello\",\n \"arrow\",\n]",
- format!("{:?}", arr)
- );
- }
-
- #[test]
- fn test_string_array_from_iter() {
- let data = vec![Some("hello"), None, Some("arrow")];
- // from Vec<Option<&str>>
- let array1 = StringArray::from(data.clone());
- // from Iterator<Option<&str>>
- let array2: StringArray = data.clone().into_iter().collect();
- // from Iterator<Option<String>>
- let array3: StringArray =
- data.into_iter().map(|x| x.map(|s| s.to_string())).collect();
-
- assert_eq!(array1, array2);
- assert_eq!(array2, array3);
- }
-
- #[test]
- fn test_string_array_from_iter_values() {
- let data = vec!["hello", "hello2"];
- let array1 = StringArray::from_iter_values(data.iter());
-
- assert_eq!(array1.value(0), "hello");
- assert_eq!(array1.value(1), "hello2");
- }
-
- #[test]
- fn test_string_array_from_unbound_iter() {
- // iterator that doesn't declare (upper) size bound
- let string_iter = (0..)
- .scan(0usize, |pos, i| {
- if *pos < 10 {
- *pos += 1;
- Some(Some(format!("value {}", i)))
- } else {
- // actually returns up to 10 values
- None
- }
- })
- // limited using take()
- .take(100);
-
- let (_, upper_size_bound) = string_iter.size_hint();
- // the upper bound, defined by take above, is 100
- assert_eq!(upper_size_bound, Some(100));
- let string_array: StringArray = string_iter.collect();
- // but the actual number of items in the array should be 10
- assert_eq!(string_array.len(), 10);
- }
-}
diff --git a/arrow/src/array/array_struct.rs b/arrow/src/array/array_struct.rs
deleted file mode 100644
index 59ee527..0000000
--- a/arrow/src/array/array_struct.rs
+++ /dev/null
@@ -1,531 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-use std::any::Any;
-use std::convert::{From, TryFrom};
-use std::fmt;
-use std::iter::IntoIterator;
-use std::mem;
-
-use super::{make_array, Array, ArrayData, ArrayRef};
-use crate::datatypes::DataType;
-use crate::error::{ArrowError, Result};
-use crate::{
- buffer::{buffer_bin_or, Buffer},
- datatypes::Field,
-};
-
-/// A nested array type where each child (called *field*) is represented by a separate
-/// array.
-pub struct StructArray {
- data: ArrayData,
- pub(crate) boxed_fields: Vec<ArrayRef>,
-}
-
-impl StructArray {
- /// Returns the field at `pos`.
- pub fn column(&self, pos: usize) -> &ArrayRef {
- &self.boxed_fields[pos]
- }
-
- /// Return the number of fields in this struct array
- pub fn num_columns(&self) -> usize {
- self.boxed_fields.len()
- }
-
- /// Returns the fields of the struct array
- pub fn columns(&self) -> Vec<&ArrayRef> {
- self.boxed_fields.iter().collect()
- }
-
- /// Returns child array refs of the struct array
- pub fn columns_ref(&self) -> Vec<ArrayRef> {
- self.boxed_fields.clone()
- }
-
- /// Return field names in this struct array
- pub fn column_names(&self) -> Vec<&str> {
- match self.data.data_type() {
- DataType::Struct(fields) => fields
- .iter()
- .map(|f| f.name().as_str())
- .collect::<Vec<&str>>(),
- _ => unreachable!("Struct array's data type is not struct!"),
- }
- }
-
- /// Return child array whose field name equals to column_name
- ///
- /// Note: A schema can currently have duplicate field names, in which case
- /// the first field will always be selected.
- /// This issue will be addressed in [ARROW-11178](https://issues.apache.org/jira/browse/ARROW-11178)
- pub fn column_by_name(&self, column_name: &str) -> Option<&ArrayRef> {
- self.column_names()
- .iter()
- .position(|c| c == &column_name)
- .map(|pos| self.column(pos))
- }
-}
-
-impl From<ArrayData> for StructArray {
- fn from(data: ArrayData) -> Self {
- let mut boxed_fields = vec![];
- for cd in data.child_data() {
- let child_data = if data.offset() != 0 || data.len() != cd.len() {
- cd.slice(data.offset(), data.len())
- } else {
- cd.clone()
- };
- boxed_fields.push(make_array(child_data));
- }
- Self { data, boxed_fields }
- }
-}
-
-impl TryFrom<Vec<(&str, ArrayRef)>> for StructArray {
- type Error = ArrowError;
-
- /// builds a StructArray from a vector of names and arrays.
- /// This errors if the values have a different length.
- /// An entry is set to Null when all values are null.
- fn try_from(values: Vec<(&str, ArrayRef)>) -> Result<Self> {
- let values_len = values.len();
-
- // these will be populated
- let mut fields = Vec::with_capacity(values_len);
- let mut child_data = Vec::with_capacity(values_len);
-
- // len: the size of the arrays.
- let mut len: Option<usize> = None;
- // null: the null mask of the arrays.
- let mut null: Option<Buffer> = None;
- for (field_name, array) in values {
- let child_datum = array.data();
- let child_datum_len = child_datum.len();
- if let Some(len) = len {
- if len != child_datum_len {
- return Err(ArrowError::InvalidArgumentError(
- format!("Array of field \"{}\" has length {}, but previous elements have length {}.
- All arrays in every entry in a struct array must have the same length.", field_name, child_datum_len, len)
- ));
- }
- } else {
- len = Some(child_datum_len)
- }
- child_data.push(child_datum.clone());
- fields.push(Field::new(
- field_name,
- array.data_type().clone(),
- child_datum.null_buffer().is_some(),
- ));
-
- if let Some(child_null_buffer) = child_datum.null_buffer() {
- let child_datum_offset = child_datum.offset();
-
- null = Some(if let Some(null_buffer) = &null {
- buffer_bin_or(
- null_buffer,
- 0,
- child_null_buffer,
- child_datum_offset,
- child_datum_len,
- )
- } else {
- child_null_buffer.bit_slice(child_datum_offset, child_datum_len)
- });
- } else if null.is_some() {
- // when one of the fields has no nulls, them there is no null in the array
- null = None;
- }
- }
- let len = len.unwrap();
-
- let mut builder = ArrayData::builder(DataType::Struct(fields))
- .len(len)
- .child_data(child_data);
- if let Some(null_buffer) = null {
- builder = builder.null_bit_buffer(null_buffer);
- }
-
- Ok(StructArray::from(builder.build()))
- }
-}
-
-impl Array for StructArray {
- fn as_any(&self) -> &Any {
- self
- }
-
- fn data(&self) -> &ArrayData {
- &self.data
- }
-
- /// Returns the length (i.e., number of elements) of this array
- fn len(&self) -> usize {
- self.data_ref().len()
- }
-
- /// Returns the total number of bytes of memory occupied by the buffers owned by this [StructArray].
- fn get_buffer_memory_size(&self) -> usize {
- self.data.get_buffer_memory_size()
- }
-
- /// Returns the total number of bytes of memory occupied physically by this [StructArray].
- fn get_array_memory_size(&self) -> usize {
- self.data.get_array_memory_size() + mem::size_of_val(self)
- }
-}
-
-impl From<Vec<(Field, ArrayRef)>> for StructArray {
- fn from(v: Vec<(Field, ArrayRef)>) -> Self {
- let (field_types, field_values): (Vec<_>, Vec<_>) = v.into_iter().unzip();
-
- // Check the length of the child arrays
- let length = field_values[0].len();
- for i in 1..field_values.len() {
- assert_eq!(
- length,
- field_values[i].len(),
- "all child arrays of a StructArray must have the same length"
- );
- assert_eq!(
- field_types[i].data_type(),
- field_values[i].data().data_type(),
- "the field data types must match the array data in a StructArray"
- )
- }
-
- let data = ArrayData::builder(DataType::Struct(field_types))
- .child_data(field_values.into_iter().map(|a| a.data().clone()).collect())
- .len(length)
- .build();
- Self::from(data)
- }
-}
-
-impl fmt::Debug for StructArray {
- fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
- write!(f, "StructArray\n[\n")?;
- for (child_index, name) in self.column_names().iter().enumerate() {
- let column = self.column(child_index);
- writeln!(
- f,
- "-- child {}: \"{}\" ({:?})",
- child_index,
- name,
- column.data_type()
- )?;
- fmt::Debug::fmt(column, f)?;
- writeln!(f)?;
- }
- write!(f, "]")
- }
-}
-
-impl From<(Vec<(Field, ArrayRef)>, Buffer)> for StructArray {
- fn from(pair: (Vec<(Field, ArrayRef)>, Buffer)) -> Self {
- let (field_types, field_values): (Vec<_>, Vec<_>) = pair.0.into_iter().unzip();
-
- // Check the length of the child arrays
- let length = field_values[0].len();
- for i in 1..field_values.len() {
- assert_eq!(
- length,
- field_values[i].len(),
- "all child arrays of a StructArray must have the same length"
- );
- assert_eq!(
- field_types[i].data_type(),
- field_values[i].data().data_type(),
- "the field data types must match the array data in a StructArray"
- )
- }
-
- let data = ArrayData::builder(DataType::Struct(field_types))
- .null_bit_buffer(pair.1)
- .child_data(field_values.into_iter().map(|a| a.data().clone()).collect())
- .len(length)
- .build();
- Self::from(data)
- }
-}
-
-#[cfg(test)]
-mod tests {
- use super::*;
-
- use std::sync::Arc;
-
- use crate::{
- array::BooleanArray, array::Float32Array, array::Float64Array, array::Int32Array,
- array::StringArray, bitmap::Bitmap,
- };
- use crate::{
- array::Int64Array,
- datatypes::{DataType, Field},
- };
- use crate::{buffer::Buffer, datatypes::ToByteSlice};
-
- #[test]
- fn test_struct_array_builder() {
- let array = BooleanArray::from(vec![false, false, true, true]);
- let boolean_data = array.data();
- let array = Int64Array::from(vec![42, 28, 19, 31]);
- let int_data = array.data();
-
- let fields = vec![
- Field::new("a", DataType::Boolean, false),
- Field::new("b", DataType::Int64, false),
- ];
- let struct_array_data = ArrayData::builder(DataType::Struct(fields))
- .len(4)
- .add_child_data(boolean_data.clone())
- .add_child_data(int_data.clone())
- .build();
- let struct_array = StructArray::from(struct_array_data);
-
- assert_eq!(boolean_data, struct_array.column(0).data());
- assert_eq!(int_data, struct_array.column(1).data());
- }
-
- #[test]
- fn test_struct_array_from() {
- let boolean = Arc::new(BooleanArray::from(vec![false, false, true, true]));
- let int = Arc::new(Int32Array::from(vec![42, 28, 19, 31]));
-
- let struct_array = StructArray::from(vec![
- (
- Field::new("b", DataType::Boolean, false),
- boolean.clone() as ArrayRef,
- ),
- (
- Field::new("c", DataType::Int32, false),
- int.clone() as ArrayRef,
- ),
- ]);
- assert_eq!(struct_array.column(0).as_ref(), boolean.as_ref());
- assert_eq!(struct_array.column(1).as_ref(), int.as_ref());
- assert_eq!(4, struct_array.len());
- assert_eq!(0, struct_array.null_count());
- assert_eq!(0, struct_array.offset());
- }
-
- /// validates that the in-memory representation follows [the spec](https://arrow.apache.org/docs/format/Columnar.html#struct-layout)
- #[test]
- fn test_struct_array_from_vec() {
- let strings: ArrayRef = Arc::new(StringArray::from(vec![
- Some("joe"),
- None,
- None,
- Some("mark"),
- ]));
- let ints: ArrayRef =
- Arc::new(Int32Array::from(vec![Some(1), Some(2), None, Some(4)]));
-
- let arr =
- StructArray::try_from(vec![("f1", strings.clone()), ("f2", ints.clone())])
- .unwrap();
-
- let struct_data = arr.data();
- assert_eq!(4, struct_data.len());
- assert_eq!(1, struct_data.null_count());
- assert_eq!(
- // 00001011
- &Some(Bitmap::from(Buffer::from(&[11_u8]))),
- struct_data.null_bitmap()
- );
-
- let expected_string_data = ArrayData::builder(DataType::Utf8)
- .len(4)
- .null_bit_buffer(Buffer::from(&[9_u8]))
- .add_buffer(Buffer::from(&[0, 3, 3, 3, 7].to_byte_slice()))
- .add_buffer(Buffer::from(b"joemark"))
- .build();
-
- let expected_int_data = ArrayData::builder(DataType::Int32)
- .len(4)
- .null_bit_buffer(Buffer::from(&[11_u8]))
- .add_buffer(Buffer::from(&[1, 2, 0, 4].to_byte_slice()))
- .build();
-
- assert_eq!(&expected_string_data, arr.column(0).data());
-
- // TODO: implement equality for ArrayData
- assert_eq!(expected_int_data.len(), arr.column(1).data().len());
- assert_eq!(
- expected_int_data.null_count(),
- arr.column(1).data().null_count()
- );
- assert_eq!(
- expected_int_data.null_bitmap(),
- arr.column(1).data().null_bitmap()
- );
- let expected_value_buf = expected_int_data.buffers()[0].clone();
- let actual_value_buf = arr.column(1).data().buffers()[0].clone();
- for i in 0..expected_int_data.len() {
- if !expected_int_data.is_null(i) {
- assert_eq!(
- expected_value_buf.as_slice()[i * 4..(i + 1) * 4],
- actual_value_buf.as_slice()[i * 4..(i + 1) * 4]
- );
- }
- }
- }
-
- #[test]
- fn test_struct_array_from_vec_error() {
- let strings: ArrayRef = Arc::new(StringArray::from(vec![
- Some("joe"),
- None,
- None,
- // 3 elements, not 4
- ]));
- let ints: ArrayRef =
- Arc::new(Int32Array::from(vec![Some(1), Some(2), None, Some(4)]));
-
- let arr =
- StructArray::try_from(vec![("f1", strings.clone()), ("f2", ints.clone())]);
-
- match arr {
- Err(ArrowError::InvalidArgumentError(e)) => {
- assert!(e.starts_with("Array of field \"f2\" has length 4, but previous elements have length 3."));
- }
- _ => panic!("This test got an unexpected error type"),
- };
- }
-
- #[test]
- #[should_panic(
- expected = "the field data types must match the array data in a StructArray"
- )]
- fn test_struct_array_from_mismatched_types() {
- StructArray::from(vec![
- (
- Field::new("b", DataType::Int16, false),
- Arc::new(BooleanArray::from(vec![false, false, true, true]))
- as Arc<Array>,
- ),
- (
- Field::new("c", DataType::Utf8, false),
- Arc::new(Int32Array::from(vec![42, 28, 19, 31])),
- ),
- ]);
- }
-
- #[test]
- fn test_struct_array_slice() {
- let boolean_data = ArrayData::builder(DataType::Boolean)
- .len(5)
- .add_buffer(Buffer::from([0b00010000]))
- .null_bit_buffer(Buffer::from([0b00010001]))
- .build();
- let int_data = ArrayData::builder(DataType::Int32)
- .len(5)
- .add_buffer(Buffer::from([0, 28, 42, 0, 0].to_byte_slice()))
- .null_bit_buffer(Buffer::from([0b00000110]))
- .build();
-
- let mut field_types = vec![];
- field_types.push(Field::new("a", DataType::Boolean, false));
- field_types.push(Field::new("b", DataType::Int32, false));
- let struct_array_data = ArrayData::builder(DataType::Struct(field_types))
- .len(5)
- .add_child_data(boolean_data.clone())
- .add_child_data(int_data.clone())
- .null_bit_buffer(Buffer::from([0b00010111]))
- .build();
- let struct_array = StructArray::from(struct_array_data);
-
- assert_eq!(5, struct_array.len());
- assert_eq!(1, struct_array.null_count());
- assert!(struct_array.is_valid(0));
- assert!(struct_array.is_valid(1));
- assert!(struct_array.is_valid(2));
- assert!(struct_array.is_null(3));
- assert!(struct_array.is_valid(4));
- assert_eq!(&boolean_data, struct_array.column(0).data());
- assert_eq!(&int_data, struct_array.column(1).data());
-
- let c0 = struct_array.column(0);
- let c0 = c0.as_any().downcast_ref::<BooleanArray>().unwrap();
- assert_eq!(5, c0.len());
- assert_eq!(3, c0.null_count());
- assert!(c0.is_valid(0));
- assert_eq!(false, c0.value(0));
- assert!(c0.is_null(1));
- assert!(c0.is_null(2));
- assert!(c0.is_null(3));
- assert!(c0.is_valid(4));
- assert_eq!(true, c0.value(4));
-
- let c1 = struct_array.column(1);
- let c1 = c1.as_any().downcast_ref::<Int32Array>().unwrap();
- assert_eq!(5, c1.len());
- assert_eq!(3, c1.null_count());
- assert!(c1.is_null(0));
- assert!(c1.is_valid(1));
- assert_eq!(28, c1.value(1));
- assert!(c1.is_valid(2));
- assert_eq!(42, c1.value(2));
- assert!(c1.is_null(3));
- assert!(c1.is_null(4));
-
- let sliced_array = struct_array.slice(2, 3);
- let sliced_array = sliced_array.as_any().downcast_ref::<StructArray>().unwrap();
- assert_eq!(3, sliced_array.len());
- assert_eq!(2, sliced_array.offset());
- assert_eq!(1, sliced_array.null_count());
- assert!(sliced_array.is_valid(0));
- assert!(sliced_array.is_null(1));
- assert!(sliced_array.is_valid(2));
-
- let sliced_c0 = sliced_array.column(0);
- let sliced_c0 = sliced_c0.as_any().downcast_ref::<BooleanArray>().unwrap();
- assert_eq!(3, sliced_c0.len());
- assert_eq!(2, sliced_c0.offset());
- assert!(sliced_c0.is_null(0));
- assert!(sliced_c0.is_null(1));
- assert!(sliced_c0.is_valid(2));
- assert_eq!(true, sliced_c0.value(2));
-
- let sliced_c1 = sliced_array.column(1);
- let sliced_c1 = sliced_c1.as_any().downcast_ref::<Int32Array>().unwrap();
- assert_eq!(3, sliced_c1.len());
- assert_eq!(2, sliced_c1.offset());
- assert!(sliced_c1.is_valid(0));
- assert_eq!(42, sliced_c1.value(0));
- assert!(sliced_c1.is_null(1));
- assert!(sliced_c1.is_null(2));
- }
-
- #[test]
- #[should_panic(
- expected = "all child arrays of a StructArray must have the same length"
- )]
- fn test_invalid_struct_child_array_lengths() {
- StructArray::from(vec![
- (
- Field::new("b", DataType::Float32, false),
- Arc::new(Float32Array::from(vec![1.1])) as Arc<Array>,
- ),
- (
- Field::new("c", DataType::Float64, false),
- Arc::new(Float64Array::from(vec![2.2, 3.3])),
- ),
- ]);
- }
-}
diff --git a/arrow/src/array/array_union.rs b/arrow/src/array/array_union.rs
deleted file mode 100644
index 083d5bb..0000000
--- a/arrow/src/array/array_union.rs
+++ /dev/null
@@ -1,831 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-//! Contains the `UnionArray` type.
-//!
-//! Each slot in a `UnionArray` can have a value chosen from a number of types. Each of the
-//! possible types are named like the fields of a [`StructArray`](crate::array::StructArray).
-//! A `UnionArray` can have two possible memory layouts, "dense" or "sparse". For more information
-//! on please see the [specification](https://arrow.apache.org/docs/format/Columnar.html#union-layout).
-//!
-//! Builders are provided for `UnionArray`'s involving primitive types. `UnionArray`'s of nested
-//! types are also supported but not via `UnionBuilder`, see the tests for examples.
-//!
-//! # Example: Dense Memory Layout
-//!
-//! ```
-//! use arrow::array::UnionBuilder;
-//! use arrow::datatypes::{Float64Type, Int32Type};
-//!
-//! # fn main() -> arrow::error::Result<()> {
-//! let mut builder = UnionBuilder::new_dense(3);
-//! builder.append::<Int32Type>("a", 1).unwrap();
-//! builder.append::<Float64Type>("b", 3.0).unwrap();
-//! builder.append::<Int32Type>("a", 4).unwrap();
-//! let union = builder.build().unwrap();
-//!
-//! assert_eq!(union.type_id(0), 0_i8);
-//! assert_eq!(union.type_id(1), 1_i8);
-//! assert_eq!(union.type_id(2), 0_i8);
-//!
-//! assert_eq!(union.value_offset(0), 0_i32);
-//! assert_eq!(union.value_offset(1), 0_i32);
-//! assert_eq!(union.value_offset(2), 1_i32);
-//!
-//! # Ok(())
-//! # }
-//! ```
-//!
-//! # Example: Sparse Memory Layout
-//! ```
-//! use arrow::array::UnionBuilder;
-//! use arrow::datatypes::{Float64Type, Int32Type};
-//!
-//! # fn main() -> arrow::error::Result<()> {
-//! let mut builder = UnionBuilder::new_sparse(3);
-//! builder.append::<Int32Type>("a", 1).unwrap();
-//! builder.append::<Float64Type>("b", 3.0).unwrap();
-//! builder.append::<Int32Type>("a", 4).unwrap();
-//! let union = builder.build().unwrap();
-//!
-//! assert_eq!(union.type_id(0), 0_i8);
-//! assert_eq!(union.type_id(1), 1_i8);
-//! assert_eq!(union.type_id(2), 0_i8);
-//!
-//! assert_eq!(union.value_offset(0), 0_i32);
-//! assert_eq!(union.value_offset(1), 1_i32);
-//! assert_eq!(union.value_offset(2), 2_i32);
-//!
-//! # Ok(())
-//! # }
-//! ```
-use crate::array::{data::count_nulls, make_array, Array, ArrayData, ArrayRef};
-use crate::buffer::Buffer;
-use crate::datatypes::*;
-use crate::error::{ArrowError, Result};
-
-use core::fmt;
-use std::any::Any;
-use std::mem;
-use std::mem::size_of;
-
-/// An Array that can represent slots of varying types.
-pub struct UnionArray {
- data: ArrayData,
- boxed_fields: Vec<ArrayRef>,
-}
-
-impl UnionArray {
- /// Creates a new `UnionArray`.
- ///
- /// Accepts type ids, child arrays and optionally offsets (for dense unions) to create
- /// a new `UnionArray`. This method makes no attempt to validate the data provided by the
- /// caller and assumes that each of the components are correct and consistent with each other.
- /// See `try_new` for an alternative that validates the data provided.
- ///
- /// # Data Consistency
- ///
- /// The `type_ids` `Buffer` should contain `i8` values. These values should be greater than
- /// zero and must be less than the number of children provided in `child_arrays`. These values
- /// are used to index into the `child_arrays`.
- ///
- /// The `value_offsets` `Buffer` is only provided in the case of a dense union, sparse unions
- /// should use `None`. If provided the `value_offsets` `Buffer` should contain `i32` values.
- /// These values should be greater than zero and must be less than the length of the overall
- /// array.
- ///
- /// In both cases above we use signed integer types to maintain compatibility with other
- /// Arrow implementations.
- ///
- /// In both of the cases above we are accepting `Buffer`'s which are assumed to be representing
- /// `i8` and `i32` values respectively. `Buffer` objects are untyped and no attempt is made
- /// to ensure that the data provided is valid.
- pub fn new(
- type_ids: Buffer,
- value_offsets: Option<Buffer>,
- child_arrays: Vec<(Field, ArrayRef)>,
- bitmap_data: Option<Buffer>,
- ) -> Self {
- let (field_types, field_values): (Vec<_>, Vec<_>) =
- child_arrays.into_iter().unzip();
- let len = type_ids.len();
- let mut builder = ArrayData::builder(DataType::Union(field_types))
- .add_buffer(type_ids)
- .child_data(field_values.into_iter().map(|a| a.data().clone()).collect())
- .len(len);
- if let Some(bitmap) = bitmap_data {
- builder = builder.null_bit_buffer(bitmap)
- }
- let data = match value_offsets {
- Some(b) => builder.add_buffer(b).build(),
- None => builder.build(),
- };
- Self::from(data)
- }
- /// Attempts to create a new `UnionArray` and validates the inputs provided.
- pub fn try_new(
- type_ids: Buffer,
- value_offsets: Option<Buffer>,
- child_arrays: Vec<(Field, ArrayRef)>,
- bitmap: Option<Buffer>,
- ) -> Result<Self> {
- if let Some(b) = &value_offsets {
- let nulls = count_nulls(bitmap.as_ref(), 0, type_ids.len());
- if ((type_ids.len() - nulls) * 4) != b.len() {
- return Err(ArrowError::InvalidArgumentError(
- "Type Ids and Offsets represent a different number of array slots."
- .to_string(),
- ));
- }
- }
-
- // Check the type_ids
- let type_id_slice: &[i8] = unsafe { type_ids.typed_data() };
- let invalid_type_ids = type_id_slice
- .iter()
- .filter(|i| *i < &0)
- .collect::<Vec<&i8>>();
- if !invalid_type_ids.is_empty() {
- return Err(ArrowError::InvalidArgumentError(format!(
- "Type Ids must be positive and cannot be greater than the number of \
- child arrays, found:\n{:?}",
- invalid_type_ids
- )));
- }
-
- // Check the value offsets if provided
- if let Some(offset_buffer) = &value_offsets {
- let max_len = type_ids.len() as i32;
- let offsets_slice: &[i32] = unsafe { offset_buffer.typed_data() };
- let invalid_offsets = offsets_slice
- .iter()
- .filter(|i| *i < &0 || *i > &max_len)
- .collect::<Vec<&i32>>();
- if !invalid_offsets.is_empty() {
- return Err(ArrowError::InvalidArgumentError(format!(
- "Offsets must be positive and within the length of the Array, \
- found:\n{:?}",
- invalid_offsets
- )));
- }
- }
-
- Ok(Self::new(type_ids, value_offsets, child_arrays, bitmap))
- }
-
- /// Accesses the child array for `type_id`.
- ///
- /// # Panics
- ///
- /// Panics if the `type_id` provided is less than zero or greater than the number of types
- /// in the `Union`.
- pub fn child(&self, type_id: i8) -> ArrayRef {
- assert!(0 <= type_id);
- assert!((type_id as usize) < self.boxed_fields.len());
- self.boxed_fields[type_id as usize].clone()
- }
-
- /// Returns the `type_id` for the array slot at `index`.
- ///
- /// # Panics
- ///
- /// Panics if `index` is greater than the length of the array.
- pub fn type_id(&self, index: usize) -> i8 {
- assert!(index - self.offset() < self.len());
- self.data().buffers()[0].as_slice()[index] as i8
- }
-
- /// Returns the offset into the underlying values array for the array slot at `index`.
- ///
- /// # Panics
- ///
- /// Panics if `index` is greater than the length of the array.
- pub fn value_offset(&self, index: usize) -> i32 {
- assert!(index - self.offset() < self.len());
- if self.is_dense() {
- // In format v4 unions had their own validity bitmap and offsets are compressed by omitting null values
- // Starting with v5 unions don't have a validity bitmap and it's possible to directly index into the offsets buffer
- let valid_slots = match self.data.null_buffer() {
- Some(b) => b.count_set_bits_offset(0, index),
- None => index,
- };
- self.data().buffers()[1].as_slice()[valid_slots * size_of::<i32>()] as i32
- } else {
- index as i32
- }
- }
-
- /// Returns the array's value at `index`.
- ///
- /// # Panics
- ///
- /// Panics if `index` is greater than the length of the array.
- pub fn value(&self, index: usize) -> ArrayRef {
- let type_id = self.type_id(self.offset() + index);
- let value_offset = self.value_offset(self.offset() + index) as usize;
- let child_data = self.boxed_fields[type_id as usize].clone();
- child_data.slice(value_offset, 1)
- }
-
- /// Returns the names of the types in the union.
- pub fn type_names(&self) -> Vec<&str> {
- match self.data.data_type() {
- DataType::Union(fields) => fields
- .iter()
- .map(|f| f.name().as_str())
- .collect::<Vec<&str>>(),
- _ => unreachable!("Union array's data type is not a union!"),
- }
- }
-
- /// Returns whether the `UnionArray` is dense (or sparse if `false`).
- fn is_dense(&self) -> bool {
- self.data().buffers().len() == 2
- }
-}
-
-impl From<ArrayData> for UnionArray {
- fn from(data: ArrayData) -> Self {
- let mut boxed_fields = vec![];
- for cd in data.child_data() {
- boxed_fields.push(make_array(cd.clone()));
- }
- Self { data, boxed_fields }
- }
-}
-
-impl Array for UnionArray {
- fn as_any(&self) -> &Any {
- self
- }
-
- fn data(&self) -> &ArrayData {
- &self.data
- }
-
- /// Returns the total number of bytes of memory occupied by the buffers owned by this [UnionArray].
- fn get_buffer_memory_size(&self) -> usize {
- let mut size = self.data.get_buffer_memory_size();
- for field in &self.boxed_fields {
- size += field.get_buffer_memory_size();
- }
- size
- }
-
- /// Returns the total number of bytes of memory occupied physically by this [UnionArray].
- fn get_array_memory_size(&self) -> usize {
- let mut size = self.data.get_array_memory_size();
- size += mem::size_of_val(self) - mem::size_of_val(&self.boxed_fields);
- for field in &self.boxed_fields {
- size += field.get_array_memory_size();
- }
- size
- }
-}
-
-impl fmt::Debug for UnionArray {
- fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
- let header = if self.is_dense() {
- "UnionArray(Dense)\n["
- } else {
- "UnionArray(Sparse)\n["
- };
- writeln!(f, "{}", header)?;
-
- writeln!(f, "-- type id buffer:")?;
- writeln!(f, "{:?}", self.data().buffers()[0])?;
-
- if self.is_dense() {
- writeln!(f, "-- offsets buffer:")?;
- writeln!(f, "{:?}", self.data().buffers()[1])?;
- }
-
- for (child_index, name) in self.type_names().iter().enumerate() {
- let column = &self.boxed_fields[child_index];
- writeln!(
- f,
- "-- child {}: \"{}\" ({:?})",
- child_index,
- *name,
- column.data_type()
- )?;
- fmt::Debug::fmt(column, f)?;
- writeln!(f)?;
- }
- writeln!(f, "]")
- }
-}
-
-#[cfg(test)]
-mod tests {
- use super::*;
-
- use std::sync::Arc;
-
- use crate::array::*;
- use crate::buffer::Buffer;
- use crate::datatypes::{DataType, Field};
-
- #[test]
- fn test_dense_i32() {
- let mut builder = UnionBuilder::new_dense(7);
- builder.append::<Int32Type>("a", 1).unwrap();
- builder.append::<Int32Type>("b", 2).unwrap();
- builder.append::<Int32Type>("c", 3).unwrap();
- builder.append::<Int32Type>("a", 4).unwrap();
- builder.append::<Int32Type>("c", 5).unwrap();
- builder.append::<Int32Type>("a", 6).unwrap();
- builder.append::<Int32Type>("b", 7).unwrap();
- let union = builder.build().unwrap();
-
- let expected_type_ids = vec![0_i8, 1, 2, 0, 2, 0, 1];
- let expected_value_offsets = vec![0_i32, 0, 0, 1, 1, 2, 1];
- let expected_array_values = [1_i32, 2, 3, 4, 5, 6, 7];
-
- // Check type ids
- assert_eq!(
- union.data().buffers()[0],
- Buffer::from_slice_ref(&expected_type_ids)
- );
- for (i, id) in expected_type_ids.iter().enumerate() {
- assert_eq!(id, &union.type_id(i));
- }
-
- // Check offsets
- assert_eq!(
- union.data().buffers()[1],
- Buffer::from_slice_ref(&expected_value_offsets)
- );
- for (i, id) in expected_value_offsets.iter().enumerate() {
- assert_eq!(&union.value_offset(i), id);
- }
-
- // Check data
- assert_eq!(
- union.data().child_data()[0].buffers()[0],
- Buffer::from_slice_ref(&[1_i32, 4, 6])
- );
- assert_eq!(
- union.data().child_data()[1].buffers()[0],
- Buffer::from_slice_ref(&[2_i32, 7])
- );
- assert_eq!(
- union.data().child_data()[2].buffers()[0],
- Buffer::from_slice_ref(&[3_i32, 5]),
- );
-
- assert_eq!(expected_array_values.len(), union.len());
- for (i, expected_value) in expected_array_values.iter().enumerate() {
- assert_eq!(false, union.is_null(i));
- let slot = union.value(i);
- let slot = slot.as_any().downcast_ref::<Int32Array>().unwrap();
- assert_eq!(slot.len(), 1);
- let value = slot.value(0);
- assert_eq!(expected_value, &value);
- }
- }
-
- #[test]
- fn test_dense_mixed() {
- let mut builder = UnionBuilder::new_dense(7);
- builder.append::<Int32Type>("a", 1).unwrap();
- builder.append::<Int64Type>("c", 3).unwrap();
- builder.append::<Int32Type>("a", 4).unwrap();
- builder.append::<Int64Type>("c", 5).unwrap();
- builder.append::<Int32Type>("a", 6).unwrap();
- let union = builder.build().unwrap();
-
- assert_eq!(5, union.len());
- for i in 0..union.len() {
- let slot = union.value(i);
- assert_eq!(false, union.is_null(i));
- match i {
- 0 => {
- let slot = slot.as_any().downcast_ref::<Int32Array>().unwrap();
- assert_eq!(slot.len(), 1);
- let value = slot.value(0);
- assert_eq!(1_i32, value);
- }
- 1 => {
- let slot = slot.as_any().downcast_ref::<Int64Array>().unwrap();
- assert_eq!(slot.len(), 1);
- let value = slot.value(0);
- assert_eq!(3_i64, value);
- }
- 2 => {
- let slot = slot.as_any().downcast_ref::<Int32Array>().unwrap();
- assert_eq!(slot.len(), 1);
- let value = slot.value(0);
- assert_eq!(4_i32, value);
- }
- 3 => {
- let slot = slot.as_any().downcast_ref::<Int64Array>().unwrap();
- assert_eq!(slot.len(), 1);
- let value = slot.value(0);
- assert_eq!(5_i64, value);
- }
- 4 => {
- let slot = slot.as_any().downcast_ref::<Int32Array>().unwrap();
- assert_eq!(slot.len(), 1);
- let value = slot.value(0);
- assert_eq!(6_i32, value);
- }
- _ => unreachable!(),
- }
- }
- }
-
- #[test]
- fn test_dense_mixed_with_nulls() {
- let mut builder = UnionBuilder::new_dense(7);
- builder.append::<Int32Type>("a", 1).unwrap();
- builder.append::<Int64Type>("c", 3).unwrap();
- builder.append::<Int32Type>("a", 10).unwrap();
- builder.append_null().unwrap();
- builder.append::<Int32Type>("a", 6).unwrap();
- let union = builder.build().unwrap();
-
- assert_eq!(5, union.len());
- for i in 0..union.len() {
- let slot = union.value(i);
- match i {
- 0 => {
- let slot = slot.as_any().downcast_ref::<Int32Array>().unwrap();
- assert_eq!(false, union.is_null(i));
- assert_eq!(slot.len(), 1);
- let value = slot.value(0);
- assert_eq!(1_i32, value);
- }
- 1 => {
- let slot = slot.as_any().downcast_ref::<Int64Array>().unwrap();
- assert_eq!(false, union.is_null(i));
- assert_eq!(slot.len(), 1);
- let value = slot.value(0);
- assert_eq!(3_i64, value);
- }
- 2 => {
- let slot = slot.as_any().downcast_ref::<Int32Array>().unwrap();
- assert_eq!(false, union.is_null(i));
- assert_eq!(slot.len(), 1);
- let value = slot.value(0);
- assert_eq!(10_i32, value);
- }
- 3 => assert!(union.is_null(i)),
- 4 => {
- let slot = slot.as_any().downcast_ref::<Int32Array>().unwrap();
- assert_eq!(false, union.is_null(i));
- assert_eq!(slot.len(), 1);
- let value = slot.value(0);
- assert_eq!(6_i32, value);
- }
- _ => unreachable!(),
- }
- }
- }
-
- #[test]
- fn test_dense_mixed_with_nulls_and_offset() {
- let mut builder = UnionBuilder::new_dense(7);
- builder.append::<Int32Type>("a", 1).unwrap();
- builder.append::<Int64Type>("c", 3).unwrap();
- builder.append::<Int32Type>("a", 10).unwrap();
- builder.append_null().unwrap();
- builder.append::<Int32Type>("a", 6).unwrap();
- let union = builder.build().unwrap();
-
- let slice = union.slice(2, 3);
- let new_union = slice.as_any().downcast_ref::<UnionArray>().unwrap();
-
- assert_eq!(3, new_union.len());
- for i in 0..new_union.len() {
- let slot = new_union.value(i);
- match i {
- 0 => {
- let slot = slot.as_any().downcast_ref::<Int32Array>().unwrap();
- assert_eq!(false, union.is_null(i));
- assert_eq!(slot.len(), 1);
- let value = slot.value(0);
- assert_eq!(10_i32, value);
- }
- 1 => assert!(new_union.is_null(i)),
- 2 => {
- let slot = slot.as_any().downcast_ref::<Int32Array>().unwrap();
- assert_eq!(false, union.is_null(i));
- assert_eq!(slot.len(), 1);
- let value = slot.value(0);
- assert_eq!(6_i32, value);
- }
- _ => unreachable!(),
- }
- }
- }
-
- #[test]
- fn test_dense_mixed_with_str() {
- let string_array = StringArray::from(vec!["foo", "bar", "baz"]);
- let int_array = Int32Array::from(vec![5, 6]);
- let float_array = Float64Array::from(vec![10.0]);
-
- let type_ids = [1_i8, 0, 0, 2, 0, 1];
- let value_offsets = [0_i32, 0, 1, 0, 2, 1];
-
- let type_id_buffer = Buffer::from_slice_ref(&type_ids);
- let value_offsets_buffer = Buffer::from_slice_ref(&value_offsets);
-
- let mut children: Vec<(Field, Arc<Array>)> = Vec::new();
- children.push((
- Field::new("A", DataType::Utf8, false),
- Arc::new(string_array),
- ));
- children.push((Field::new("B", DataType::Int32, false), Arc::new(int_array)));
- children.push((
- Field::new("C", DataType::Float64, false),
- Arc::new(float_array),
- ));
- let array = UnionArray::try_new(
- type_id_buffer,
- Some(value_offsets_buffer),
- children,
- None,
- )
- .unwrap();
-
- // Check type ids
- assert_eq!(Buffer::from_slice_ref(&type_ids), array.data().buffers()[0]);
- for (i, id) in type_ids.iter().enumerate() {
- assert_eq!(id, &array.type_id(i));
- }
-
- // Check offsets
- assert_eq!(
- Buffer::from_slice_ref(&value_offsets),
- array.data().buffers()[1]
- );
- for (i, id) in value_offsets.iter().enumerate() {
- assert_eq!(id, &array.value_offset(i));
- }
-
- // Check values
- assert_eq!(6, array.len());
-
- let slot = array.value(0);
- let value = slot.as_any().downcast_ref::<Int32Array>().unwrap().value(0);
- assert_eq!(5, value);
-
- let slot = array.value(1);
- let value = slot
- .as_any()
- .downcast_ref::<StringArray>()
- .unwrap()
- .value(0);
- assert_eq!("foo", value);
-
- let slot = array.value(2);
- let value = slot
- .as_any()
- .downcast_ref::<StringArray>()
- .unwrap()
- .value(0);
- assert_eq!("bar", value);
-
- let slot = array.value(3);
- let value = slot
- .as_any()
- .downcast_ref::<Float64Array>()
- .unwrap()
- .value(0);
- assert!(10.0 - value < f64::EPSILON);
-
- let slot = array.value(4);
- let value = slot
- .as_any()
- .downcast_ref::<StringArray>()
- .unwrap()
- .value(0);
- assert_eq!("baz", value);
-
- let slot = array.value(5);
- let value = slot.as_any().downcast_ref::<Int32Array>().unwrap().value(0);
- assert_eq!(6, value);
- }
-
- #[test]
- fn test_sparse_i32() {
- let mut builder = UnionBuilder::new_sparse(7);
- builder.append::<Int32Type>("a", 1).unwrap();
- builder.append::<Int32Type>("b", 2).unwrap();
- builder.append::<Int32Type>("c", 3).unwrap();
- builder.append::<Int32Type>("a", 4).unwrap();
- builder.append::<Int32Type>("c", 5).unwrap();
- builder.append::<Int32Type>("a", 6).unwrap();
- builder.append::<Int32Type>("b", 7).unwrap();
- let union = builder.build().unwrap();
-
- let expected_type_ids = vec![0_i8, 1, 2, 0, 2, 0, 1];
- let expected_array_values = [1_i32, 2, 3, 4, 5, 6, 7];
-
- // Check type ids
- assert_eq!(
- Buffer::from_slice_ref(&expected_type_ids),
- union.data().buffers()[0]
- );
- for (i, id) in expected_type_ids.iter().enumerate() {
- assert_eq!(id, &union.type_id(i));
- }
-
- // Check offsets, sparse union should only have a single buffer
- assert_eq!(union.data().buffers().len(), 1);
-
- // Check data
- assert_eq!(
- union.data().child_data()[0].buffers()[0],
- Buffer::from_slice_ref(&[1_i32, 0, 0, 4, 0, 6, 0]),
- );
- assert_eq!(
- Buffer::from_slice_ref(&[0_i32, 2_i32, 0, 0, 0, 0, 7]),
- union.data().child_data()[1].buffers()[0]
- );
- assert_eq!(
- Buffer::from_slice_ref(&[0_i32, 0, 3_i32, 0, 5, 0, 0]),
- union.data().child_data()[2].buffers()[0]
- );
-
- assert_eq!(expected_array_values.len(), union.len());
- for (i, expected_value) in expected_array_values.iter().enumerate() {
- assert_eq!(false, union.is_null(i));
- let slot = union.value(i);
- let slot = slot.as_any().downcast_ref::<Int32Array>().unwrap();
- assert_eq!(slot.len(), 1);
- let value = slot.value(0);
- assert_eq!(expected_value, &value);
- }
- }
-
- #[test]
- fn test_sparse_mixed() {
- let mut builder = UnionBuilder::new_sparse(5);
- builder.append::<Int32Type>("a", 1).unwrap();
- builder.append::<Float64Type>("c", 3.0).unwrap();
- builder.append::<Int32Type>("a", 4).unwrap();
- builder.append::<Float64Type>("c", 5.0).unwrap();
- builder.append::<Int32Type>("a", 6).unwrap();
- let union = builder.build().unwrap();
-
- let expected_type_ids = vec![0_i8, 1, 0, 1, 0];
-
- // Check type ids
- assert_eq!(
- Buffer::from_slice_ref(&expected_type_ids),
- union.data().buffers()[0]
- );
- for (i, id) in expected_type_ids.iter().enumerate() {
- assert_eq!(id, &union.type_id(i));
- }
-
- // Check offsets, sparse union should only have a single buffer, i.e. no offsets
- assert_eq!(union.data().buffers().len(), 1);
-
- for i in 0..union.len() {
- let slot = union.value(i);
- assert_eq!(false, union.is_null(i));
- match i {
- 0 => {
- let slot = slot.as_any().downcast_ref::<Int32Array>().unwrap();
- assert_eq!(slot.len(), 1);
- let value = slot.value(0);
- assert_eq!(1_i32, value);
- }
- 1 => {
- let slot = slot.as_any().downcast_ref::<Float64Array>().unwrap();
- assert_eq!(slot.len(), 1);
- let value = slot.value(0);
- assert!(value - 3_f64 < f64::EPSILON);
- }
- 2 => {
- let slot = slot.as_any().downcast_ref::<Int32Array>().unwrap();
- assert_eq!(slot.len(), 1);
- let value = slot.value(0);
- assert_eq!(4_i32, value);
- }
- 3 => {
- let slot = slot.as_any().downcast_ref::<Float64Array>().unwrap();
- assert_eq!(slot.len(), 1);
- let value = slot.value(0);
- assert!(5_f64 - value < f64::EPSILON);
- }
- 4 => {
- let slot = slot.as_any().downcast_ref::<Int32Array>().unwrap();
- assert_eq!(slot.len(), 1);
- let value = slot.value(0);
- assert_eq!(6_i32, value);
- }
- _ => unreachable!(),
- }
- }
- }
-
- #[test]
- fn test_sparse_mixed_with_nulls() {
- let mut builder = UnionBuilder::new_sparse(5);
- builder.append::<Int32Type>("a", 1).unwrap();
- builder.append_null().unwrap();
- builder.append::<Float64Type>("c", 3.0).unwrap();
- builder.append::<Int32Type>("a", 4).unwrap();
- let union = builder.build().unwrap();
-
- let expected_type_ids = vec![0_i8, 0, 1, 0];
-
- // Check type ids
- assert_eq!(
- Buffer::from_slice_ref(&expected_type_ids),
- union.data().buffers()[0]
- );
- for (i, id) in expected_type_ids.iter().enumerate() {
- assert_eq!(id, &union.type_id(i));
- }
-
- // Check offsets, sparse union should only have a single buffer, i.e. no offsets
- assert_eq!(union.data().buffers().len(), 1);
-
- for i in 0..union.len() {
- let slot = union.value(i);
- match i {
- 0 => {
- let slot = slot.as_any().downcast_ref::<Int32Array>().unwrap();
- assert_eq!(false, union.is_null(i));
- assert_eq!(slot.len(), 1);
- let value = slot.value(0);
- assert_eq!(1_i32, value);
- }
- 1 => assert!(union.is_null(i)),
- 2 => {
- let slot = slot.as_any().downcast_ref::<Float64Array>().unwrap();
- assert_eq!(false, union.is_null(i));
- assert_eq!(slot.len(), 1);
- let value = slot.value(0);
- assert!(value - 3_f64 < f64::EPSILON);
- }
- 3 => {
- let slot = slot.as_any().downcast_ref::<Int32Array>().unwrap();
- assert_eq!(false, union.is_null(i));
- assert_eq!(slot.len(), 1);
- let value = slot.value(0);
- assert_eq!(4_i32, value);
- }
- _ => unreachable!(),
- }
- }
- }
-
- #[test]
- fn test_sparse_mixed_with_nulls_and_offset() {
- let mut builder = UnionBuilder::new_sparse(5);
- builder.append::<Int32Type>("a", 1).unwrap();
- builder.append_null().unwrap();
- builder.append::<Float64Type>("c", 3.0).unwrap();
- builder.append_null().unwrap();
- builder.append::<Int32Type>("a", 4).unwrap();
- let union = builder.build().unwrap();
-
- let slice = union.slice(1, 4);
- let new_union = slice.as_any().downcast_ref::<UnionArray>().unwrap();
-
- assert_eq!(4, new_union.len());
- for i in 0..new_union.len() {
- let slot = new_union.value(i);
- match i {
- 0 => assert!(new_union.is_null(i)),
- 1 => {
- let slot = slot.as_any().downcast_ref::<Float64Array>().unwrap();
- assert_eq!(false, new_union.is_null(i));
- assert_eq!(slot.len(), 1);
- let value = slot.value(0);
- assert!(value - 3_f64 < f64::EPSILON);
- }
- 2 => assert!(new_union.is_null(i)),
- 3 => {
- let slot = slot.as_any().downcast_ref::<Int32Array>().unwrap();
- assert_eq!(false, new_union.is_null(i));
- assert_eq!(slot.len(), 1);
- let value = slot.value(0);
- assert_eq!(4_i32, value);
- }
- _ => unreachable!(),
- }
- }
- }
-}
diff --git a/arrow/src/array/builder.rs b/arrow/src/array/builder.rs
deleted file mode 100644
index 3bbeaf7..0000000
--- a/arrow/src/array/builder.rs
+++ /dev/null
@@ -1,3308 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-//! Defines a [`BufferBuilder`](crate::array::BufferBuilder) capable
-//! of creating a [`Buffer`](crate::buffer::Buffer) which can be used
-//! as an internal buffer in an [`ArrayData`](crate::array::ArrayData)
-//! object.
-
-use std::any::Any;
-use std::collections::HashMap;
-use std::fmt;
-use std::marker::PhantomData;
-use std::mem;
-use std::sync::Arc;
-
-use crate::array::*;
-use crate::buffer::{Buffer, MutableBuffer};
-use crate::datatypes::*;
-use crate::error::{ArrowError, Result};
-use crate::util::bit_util;
-
-/// Converts a `MutableBuffer` to a `BufferBuilder<T>`.
-///
-/// `slots` is the number of array slots currently represented in the `MutableBuffer`.
-pub(crate) fn mutable_buffer_to_builder<T: ArrowNativeType>(
- mutable_buffer: MutableBuffer,
- slots: usize,
-) -> BufferBuilder<T> {
- BufferBuilder::<T> {
- buffer: mutable_buffer,
- len: slots,
- _marker: PhantomData,
- }
-}
-
-/// Converts a `BufferBuilder<T>` into its underlying `MutableBuffer`.
-///
-/// `From` is not implemented because associated type bounds are unstable.
-pub(crate) fn builder_to_mutable_buffer<T: ArrowNativeType>(
- builder: BufferBuilder<T>,
-) -> MutableBuffer {
- builder.buffer
-}
-
-/// Builder for creating a [`Buffer`](crate::buffer::Buffer) object.
-///
-/// A [`Buffer`](crate::buffer::Buffer) is the underlying data
-/// structure of Arrow's [`Arrays`](crate::array::Array).
-///
-/// For all supported types, there are type definitions for the
-/// generic version of `BufferBuilder<T>`, e.g. `UInt8BufferBuilder`.
-///
-/// # Example:
-///
-/// ```
-/// use arrow::array::UInt8BufferBuilder;
-///
-/// # fn main() -> arrow::error::Result<()> {
-/// let mut builder = UInt8BufferBuilder::new(100);
-/// builder.append_slice(&[42, 43, 44]);
-/// builder.append(45);
-/// let buffer = builder.finish();
-///
-/// assert_eq!(unsafe { buffer.typed_data::<u8>() }, &[42, 43, 44, 45]);
-/// # Ok(())
-/// # }
-/// ```
-#[derive(Debug)]
-pub struct BufferBuilder<T: ArrowNativeType> {
- buffer: MutableBuffer,
- len: usize,
- _marker: PhantomData<T>,
-}
-
-impl<T: ArrowNativeType> BufferBuilder<T> {
- /// Creates a new builder with initial capacity for _at least_ `capacity`
- /// elements of type `T`.
- ///
- /// The capacity can later be manually adjusted with the
- /// [`reserve()`](BufferBuilder::reserve) method.
- /// Also the
- /// [`append()`](BufferBuilder::append),
- /// [`append_slice()`](BufferBuilder::append_slice) and
- /// [`advance()`](BufferBuilder::advance)
- /// methods automatically increase the capacity if needed.
- ///
- /// # Example:
- ///
- /// ```
- /// use arrow::array::UInt8BufferBuilder;
- ///
- /// let mut builder = UInt8BufferBuilder::new(10);
- ///
- /// assert!(builder.capacity() >= 10);
- /// ```
- #[inline]
- pub fn new(capacity: usize) -> Self {
- let buffer = MutableBuffer::new(capacity * mem::size_of::<T>());
-
- Self {
- buffer,
- len: 0,
- _marker: PhantomData,
- }
- }
-
- /// Returns the current number of array elements in the internal buffer.
- ///
- /// # Example:
- ///
- /// ```
- /// use arrow::array::UInt8BufferBuilder;
- ///
- /// let mut builder = UInt8BufferBuilder::new(10);
- /// builder.append(42);
- ///
- /// assert_eq!(builder.len(), 1);
- /// ```
- pub fn len(&self) -> usize {
- self.len
- }
-
- /// Returns whether the internal buffer is empty.
- ///
- /// # Example:
- ///
- /// ```
- /// use arrow::array::UInt8BufferBuilder;
- ///
- /// let mut builder = UInt8BufferBuilder::new(10);
- /// builder.append(42);
- ///
- /// assert_eq!(builder.is_empty(), false);
- /// ```
- pub fn is_empty(&self) -> bool {
- self.len == 0
- }
-
- /// Returns the actual capacity (number of elements) of the internal buffer.
- ///
- /// Note: the internal capacity returned by this method might be larger than
- /// what you'd expect after setting the capacity in the `new()` or `reserve()`
- /// functions.
- pub fn capacity(&self) -> usize {
- let byte_capacity = self.buffer.capacity();
- byte_capacity / std::mem::size_of::<T>()
- }
-
- /// Increases the number of elements in the internal buffer by `n`
- /// and resizes the buffer as needed.
- ///
- /// The values of the newly added elements are 0.
- /// This method is usually used when appending `NULL` values to the buffer
- /// as they still require physical memory space.
- ///
- /// # Example:
- ///
- /// ```
- /// use arrow::array::UInt8BufferBuilder;
- ///
- /// let mut builder = UInt8BufferBuilder::new(10);
- /// builder.advance(2);
- ///
- /// assert_eq!(builder.len(), 2);
- /// ```
- #[inline]
- pub fn advance(&mut self, i: usize) {
- let new_buffer_len = (self.len + i) * mem::size_of::<T>();
- self.buffer.resize(new_buffer_len, 0);
- self.len += i;
- }
-
- /// Reserves memory for _at least_ `n` more elements of type `T`.
- ///
- /// # Example:
- ///
- /// ```
- /// use arrow::array::UInt8BufferBuilder;
- ///
- /// let mut builder = UInt8BufferBuilder::new(10);
- /// builder.reserve(10);
- ///
- /// assert!(builder.capacity() >= 20);
- /// ```
- #[inline]
- pub fn reserve(&mut self, n: usize) {
- self.buffer.reserve(n * mem::size_of::<T>());
- }
-
- /// Appends a value of type `T` into the builder,
- /// growing the internal buffer as needed.
- ///
- /// # Example:
- ///
- /// ```
- /// use arrow::array::UInt8BufferBuilder;
- ///
- /// let mut builder = UInt8BufferBuilder::new(10);
- /// builder.append(42);
- ///
- /// assert_eq!(builder.len(), 1);
- /// ```
- #[inline]
- pub fn append(&mut self, v: T) {
- self.reserve(1);
- self.buffer.push(v);
- self.len += 1;
- }
-
- /// Appends a value of type `T` into the builder N times,
- /// growing the internal buffer as needed.
- ///
- /// # Example:
- ///
- /// ```
- /// use arrow::array::UInt8BufferBuilder;
- ///
- /// let mut builder = UInt8BufferBuilder::new(10);
- /// builder.append_n(10, 42);
- ///
- /// assert_eq!(builder.len(), 10);
- /// ```
- #[inline]
- pub fn append_n(&mut self, n: usize, v: T) {
- self.reserve(n);
- for _ in 0..n {
- self.buffer.push(v);
- }
- self.len += n;
- }
-
- /// Appends a slice of type `T`, growing the internal buffer as needed.
- ///
- /// # Example:
- ///
- /// ```
- /// use arrow::array::UInt8BufferBuilder;
- ///
- /// let mut builder = UInt8BufferBuilder::new(10);
- /// builder.append_slice(&[42, 44, 46]);
- ///
- /// assert_eq!(builder.len(), 3);
- /// ```
- #[inline]
- pub fn append_slice(&mut self, slice: &[T]) {
- self.buffer.extend_from_slice(slice);
- self.len += slice.len();
- }
-
- /// Resets this builder and returns an immutable [`Buffer`](crate::buffer::Buffer).
- ///
- /// # Example:
- ///
- /// ```
- /// use arrow::array::UInt8BufferBuilder;
- ///
- /// let mut builder = UInt8BufferBuilder::new(10);
- /// builder.append_slice(&[42, 44, 46]);
- ///
- /// let buffer = builder.finish();
- ///
- /// assert_eq!(unsafe { buffer.typed_data::<u8>() }, &[42, 44, 46]);
- /// ```
- #[inline]
- pub fn finish(&mut self) -> Buffer {
- let buf = std::mem::replace(&mut self.buffer, MutableBuffer::new(0));
- self.len = 0;
- buf.into()
- }
-}
-
-#[derive(Debug)]
-pub struct BooleanBufferBuilder {
- buffer: MutableBuffer,
- len: usize,
-}
-
-impl BooleanBufferBuilder {
- #[inline]
- pub fn new(capacity: usize) -> Self {
- let byte_capacity = bit_util::ceil(capacity, 8);
- let buffer = MutableBuffer::from_len_zeroed(byte_capacity);
- Self { buffer, len: 0 }
- }
-
- #[inline]
- pub fn len(&self) -> usize {
- self.len
- }
-
- #[inline]
- pub fn set_bit(&mut self, index: usize, v: bool) {
- if v {
- bit_util::set_bit(self.buffer.as_mut(), index);
- } else {
- bit_util::unset_bit(self.buffer.as_mut(), index);
- }
- }
-
- #[inline]
- pub fn is_empty(&self) -> bool {
- self.len == 0
- }
-
- #[inline]
- pub fn capacity(&self) -> usize {
- self.buffer.capacity() * 8
- }
-
- #[inline]
- pub fn advance(&mut self, additional: usize) {
- let new_len = self.len + additional;
- let new_len_bytes = bit_util::ceil(new_len, 8);
- if new_len_bytes > self.buffer.len() {
- self.buffer.resize(new_len_bytes, 0);
- }
- self.len = new_len;
- }
-
- /// Reserve space to at least `additional` new bits.
- /// Capacity will be `>= self.len() + additional`.
- /// New bytes are uninitialized and reading them is undefined behavior.
- #[inline]
- pub fn reserve(&mut self, additional: usize) {
- let capacity = self.len + additional;
- if capacity > self.capacity() {
- // convert differential to bytes
- let additional = bit_util::ceil(capacity, 8) - self.buffer.len();
- self.buffer.reserve(additional);
- }
- }
-
- #[inline]
- pub fn append(&mut self, v: bool) {
- self.advance(1);
- if v {
- unsafe { bit_util::set_bit_raw(self.buffer.as_mut_ptr(), self.len - 1) };
- }
- }
-
- #[inline]
- pub fn append_n(&mut self, additional: usize, v: bool) {
- self.advance(additional);
- if additional > 0 && v {
- let offset = self.len() - additional;
- (0..additional).for_each(|i| unsafe {
- bit_util::set_bit_raw(self.buffer.as_mut_ptr(), offset + i)
- })
- }
- }
-
- #[inline]
- pub fn append_slice(&mut self, slice: &[bool]) {
- let additional = slice.len();
- self.advance(additional);
-
- let offset = self.len() - additional;
- for (i, v) in slice.iter().enumerate() {
- if *v {
- unsafe { bit_util::set_bit_raw(self.buffer.as_mut_ptr(), offset + i) }
- }
- }
- }
-
- #[inline]
- pub fn finish(&mut self) -> Buffer {
- let buf = std::mem::replace(&mut self.buffer, MutableBuffer::new(0));
- self.len = 0;
- buf.into()
- }
-}
-
-impl From<BooleanBufferBuilder> for Buffer {
- #[inline]
- fn from(builder: BooleanBufferBuilder) -> Self {
- builder.buffer.into()
- }
-}
-
-/// Trait for dealing with different array builders at runtime
-///
-/// # Example
-///
-/// ```
-/// # use arrow::{
-/// # array::{ArrayBuilder, ArrayRef, Float64Builder, Int64Builder, StringArray, StringBuilder},
-/// # error::ArrowError,
-/// # };
-/// # fn main() -> std::result::Result<(), ArrowError> {
-/// // Create
-/// let mut data_builders: Vec<Box<dyn ArrayBuilder>> = vec![
-/// Box::new(Float64Builder::new(1024)),
-/// Box::new(Int64Builder::new(1024)),
-/// Box::new(StringBuilder::new(1024)),
-/// ];
-///
-/// // Fill
-/// data_builders[0]
-/// .as_any_mut()
-/// .downcast_mut::<Float64Builder>()
-/// .unwrap()
-/// .append_value(3.14)?;
-/// data_builders[1]
-/// .as_any_mut()
-/// .downcast_mut::<Int64Builder>()
-/// .unwrap()
-/// .append_value(-1)?;
-/// data_builders[2]
-/// .as_any_mut()
-/// .downcast_mut::<StringBuilder>()
-/// .unwrap()
-/// .append_value("🍎")?;
-///
-/// // Finish
-/// let array_refs: Vec<ArrayRef> = data_builders
-/// .iter_mut()
-/// .map(|builder| builder.finish())
-/// .collect();
-/// assert_eq!(array_refs[0].len(), 1);
-/// assert_eq!(array_refs[1].is_null(0), false);
-/// assert_eq!(
-/// array_refs[2]
-/// .as_any()
-/// .downcast_ref::<StringArray>()
-/// .unwrap()
-/// .value(0),
-/// "🍎"
-/// );
-/// # Ok(())
-/// # }
-/// ```
-pub trait ArrayBuilder: Any + Send {
- /// Returns the number of array slots in the builder
- fn len(&self) -> usize;
-
- /// Returns whether number of array slots is zero
- fn is_empty(&self) -> bool;
-
- /// Builds the array
- fn finish(&mut self) -> ArrayRef;
-
- /// Returns the builder as a non-mutable `Any` reference.
- ///
- /// This is most useful when one wants to call non-mutable APIs on a specific builder
- /// type. In this case, one can first cast this into a `Any`, and then use
- /// `downcast_ref` to get a reference on the specific builder.
- fn as_any(&self) -> &Any;
-
- /// Returns the builder as a mutable `Any` reference.
- ///
- /// This is most useful when one wants to call mutable APIs on a specific builder
- /// type. In this case, one can first cast this into a `Any`, and then use
- /// `downcast_mut` to get a reference on the specific builder.
- fn as_any_mut(&mut self) -> &mut Any;
-
- /// Returns the boxed builder as a box of `Any`.
- fn into_box_any(self: Box<Self>) -> Box<Any>;
-}
-
-/// Array builder for fixed-width primitive types
-#[derive(Debug)]
-pub struct BooleanBuilder {
- values_builder: BooleanBufferBuilder,
- bitmap_builder: BooleanBufferBuilder,
-}
-
-impl BooleanBuilder {
- /// Creates a new primitive array builder
- pub fn new(capacity: usize) -> Self {
- Self {
- values_builder: BooleanBufferBuilder::new(capacity),
- bitmap_builder: BooleanBufferBuilder::new(capacity),
- }
- }
-
- /// Returns the capacity of this builder measured in slots of type `T`
- pub fn capacity(&self) -> usize {
- self.values_builder.capacity()
- }
-
- /// Appends a value of type `T` into the builder
- #[inline]
- pub fn append_value(&mut self, v: bool) -> Result<()> {
- self.bitmap_builder.append(true);
- self.values_builder.append(v);
- Ok(())
- }
-
- /// Appends a null slot into the builder
- #[inline]
- pub fn append_null(&mut self) -> Result<()> {
- self.bitmap_builder.append(false);
- self.values_builder.advance(1);
- Ok(())
- }
-
- /// Appends an `Option<T>` into the builder
- #[inline]
- pub fn append_option(&mut self, v: Option<bool>) -> Result<()> {
- match v {
- None => self.append_null()?,
- Some(v) => self.append_value(v)?,
- };
- Ok(())
- }
-
- /// Appends a slice of type `T` into the builder
- #[inline]
- pub fn append_slice(&mut self, v: &[bool]) -> Result<()> {
- self.bitmap_builder.append_n(v.len(), true);
- self.values_builder.append_slice(v);
- Ok(())
- }
-
- /// Appends values from a slice of type `T` and a validity boolean slice
- #[inline]
- pub fn append_values(&mut self, values: &[bool], is_valid: &[bool]) -> Result<()> {
- if values.len() != is_valid.len() {
- return Err(ArrowError::InvalidArgumentError(
- "Value and validity lengths must be equal".to_string(),
- ));
- }
- self.bitmap_builder.append_slice(is_valid);
- self.values_builder.append_slice(values);
- Ok(())
- }
-
- /// Builds the [BooleanArray] and reset this builder.
- pub fn finish(&mut self) -> BooleanArray {
- let len = self.len();
- let null_bit_buffer = self.bitmap_builder.finish();
- let null_count = len - null_bit_buffer.count_set_bits();
- let mut builder = ArrayData::builder(DataType::Boolean)
- .len(len)
- .add_buffer(self.values_builder.finish());
- if null_count > 0 {
- builder = builder.null_bit_buffer(null_bit_buffer);
- }
- let data = builder.build();
- BooleanArray::from(data)
- }
-}
-
-impl ArrayBuilder for BooleanBuilder {
- /// Returns the builder as a non-mutable `Any` reference.
- fn as_any(&self) -> &Any {
- self
- }
-
- /// Returns the builder as a mutable `Any` reference.
- fn as_any_mut(&mut self) -> &mut Any {
- self
- }
-
- /// Returns the boxed builder as a box of `Any`.
- fn into_box_any(self: Box<Self>) -> Box<Any> {
- self
- }
-
- /// Returns the number of array slots in the builder
- fn len(&self) -> usize {
- self.values_builder.len
- }
-
- /// Returns whether the number of array slots is zero
- fn is_empty(&self) -> bool {
- self.values_builder.is_empty()
- }
-
- /// Builds the array and reset this builder.
- fn finish(&mut self) -> ArrayRef {
- Arc::new(self.finish())
- }
-}
-
-/// Array builder for fixed-width primitive types
-#[derive(Debug)]
-pub struct PrimitiveBuilder<T: ArrowPrimitiveType> {
- values_builder: BufferBuilder<T::Native>,
- /// We only materialize the builder when we add `false`.
- /// This optimization is **very** important for performance of `StringBuilder`.
- bitmap_builder: Option<BooleanBufferBuilder>,
-}
-
-impl<T: ArrowPrimitiveType> ArrayBuilder for PrimitiveBuilder<T> {
- /// Returns the builder as a non-mutable `Any` reference.
- fn as_any(&self) -> &Any {
- self
- }
-
- /// Returns the builder as a mutable `Any` reference.
- fn as_any_mut(&mut self) -> &mut Any {
- self
- }
-
- /// Returns the boxed builder as a box of `Any`.
- fn into_box_any(self: Box<Self>) -> Box<Any> {
- self
- }
-
- /// Returns the number of array slots in the builder
- fn len(&self) -> usize {
- self.values_builder.len
- }
-
- /// Returns whether the number of array slots is zero
- fn is_empty(&self) -> bool {
- self.values_builder.is_empty()
- }
-
- /// Builds the array and reset this builder.
- fn finish(&mut self) -> ArrayRef {
- Arc::new(self.finish())
- }
-}
-
-impl<T: ArrowPrimitiveType> PrimitiveBuilder<T> {
- /// Creates a new primitive array builder
- pub fn new(capacity: usize) -> Self {
- Self {
- values_builder: BufferBuilder::<T::Native>::new(capacity),
- bitmap_builder: None,
- }
- }
-
- /// Returns the capacity of this builder measured in slots of type `T`
- pub fn capacity(&self) -> usize {
- self.values_builder.capacity()
- }
-
- /// Appends a value of type `T` into the builder
- #[inline]
- pub fn append_value(&mut self, v: T::Native) -> Result<()> {
- if let Some(b) = self.bitmap_builder.as_mut() {
- b.append(true);
- }
- self.values_builder.append(v);
- Ok(())
- }
-
- /// Appends a null slot into the builder
- #[inline]
- pub fn append_null(&mut self) -> Result<()> {
- self.materialize_bitmap_builder();
- self.bitmap_builder.as_mut().unwrap().append(false);
- self.values_builder.advance(1);
- Ok(())
- }
-
- /// Appends an `Option<T>` into the builder
- #[inline]
- pub fn append_option(&mut self, v: Option<T::Native>) -> Result<()> {
- match v {
- None => self.append_null()?,
- Some(v) => self.append_value(v)?,
- };
- Ok(())
- }
-
- /// Appends a slice of type `T` into the builder
- #[inline]
- pub fn append_slice(&mut self, v: &[T::Native]) -> Result<()> {
- if let Some(b) = self.bitmap_builder.as_mut() {
- b.append_n(v.len(), true);
- }
- self.values_builder.append_slice(v);
- Ok(())
- }
-
- /// Appends values from a slice of type `T` and a validity boolean slice
- #[inline]
- pub fn append_values(
- &mut self,
- values: &[T::Native],
- is_valid: &[bool],
- ) -> Result<()> {
- if values.len() != is_valid.len() {
- return Err(ArrowError::InvalidArgumentError(
- "Value and validity lengths must be equal".to_string(),
- ));
- }
- if is_valid.iter().any(|v| !*v) {
- self.materialize_bitmap_builder();
- }
- if let Some(b) = self.bitmap_builder.as_mut() {
- b.append_slice(is_valid);
- }
- self.values_builder.append_slice(values);
- Ok(())
- }
-
- /// Builds the `PrimitiveArray` and reset this builder.
- pub fn finish(&mut self) -> PrimitiveArray<T> {
- let len = self.len();
- let null_bit_buffer = self.bitmap_builder.as_mut().map(|b| b.finish());
- let null_count = len
- - null_bit_buffer
- .as_ref()
- .map(|b| b.count_set_bits())
- .unwrap_or(len);
- let mut builder = ArrayData::builder(T::DATA_TYPE)
- .len(len)
- .add_buffer(self.values_builder.finish());
- if null_count > 0 {
- builder = builder.null_bit_buffer(null_bit_buffer.unwrap());
- }
- let data = builder.build();
- PrimitiveArray::<T>::from(data)
- }
-
- /// Builds the `DictionaryArray` and reset this builder.
- pub fn finish_dict(&mut self, values: ArrayRef) -> DictionaryArray<T> {
- let len = self.len();
- let null_bit_buffer = self.bitmap_builder.as_mut().map(|b| b.finish());
- let null_count = len
- - null_bit_buffer
- .as_ref()
- .map(|b| b.count_set_bits())
- .unwrap_or(len);
- let data_type = DataType::Dictionary(
- Box::new(T::DATA_TYPE),
- Box::new(values.data_type().clone()),
- );
- let mut builder = ArrayData::builder(data_type)
- .len(len)
- .add_buffer(self.values_builder.finish());
- if null_count > 0 {
- builder = builder.null_bit_buffer(null_bit_buffer.unwrap());
- }
- builder = builder.add_child_data(values.data().clone());
- DictionaryArray::<T>::from(builder.build())
- }
-
- fn materialize_bitmap_builder(&mut self) {
- if self.bitmap_builder.is_some() {
- return;
- }
- let mut b = BooleanBufferBuilder::new(0);
- b.reserve(self.values_builder.capacity());
- b.append_n(self.values_builder.len, true);
- self.bitmap_builder = Some(b);
- }
-}
-
-/// Array builder for `ListArray`
-#[derive(Debug)]
-pub struct GenericListBuilder<OffsetSize: OffsetSizeTrait, T: ArrayBuilder> {
- offsets_builder: BufferBuilder<OffsetSize>,
- bitmap_builder: BooleanBufferBuilder,
- values_builder: T,
- len: OffsetSize,
-}
-
-impl<OffsetSize: OffsetSizeTrait, T: ArrayBuilder> GenericListBuilder<OffsetSize, T> {
- /// Creates a new `ListArrayBuilder` from a given values array builder
- pub fn new(values_builder: T) -> Self {
- let capacity = values_builder.len();
- Self::with_capacity(values_builder, capacity)
- }
-
- /// Creates a new `ListArrayBuilder` from a given values array builder
- /// `capacity` is the number of items to pre-allocate space for in this builder
- pub fn with_capacity(values_builder: T, capacity: usize) -> Self {
- let mut offsets_builder = BufferBuilder::<OffsetSize>::new(capacity + 1);
- let len = OffsetSize::zero();
- offsets_builder.append(len);
- Self {
- offsets_builder,
- bitmap_builder: BooleanBufferBuilder::new(capacity),
- values_builder,
- len,
- }
- }
-}
-
-impl<OffsetSize: OffsetSizeTrait, T: ArrayBuilder> ArrayBuilder
- for GenericListBuilder<OffsetSize, T>
-where
- T: 'static,
-{
- /// Returns the builder as a non-mutable `Any` reference.
- fn as_any(&self) -> &Any {
- self
- }
-
- /// Returns the builder as a mutable `Any` reference.
- fn as_any_mut(&mut self) -> &mut Any {
- self
- }
-
- /// Returns the boxed builder as a box of `Any`.
- fn into_box_any(self: Box<Self>) -> Box<Any> {
- self
- }
-
- /// Returns the number of array slots in the builder
- fn len(&self) -> usize {
- self.len.to_usize().unwrap()
- }
-
- /// Returns whether the number of array slots is zero
- fn is_empty(&self) -> bool {
- self.len == OffsetSize::zero()
- }
-
- /// Builds the array and reset this builder.
- fn finish(&mut self) -> ArrayRef {
- Arc::new(self.finish())
- }
-}
-
-impl<OffsetSize: OffsetSizeTrait, T: ArrayBuilder> GenericListBuilder<OffsetSize, T>
-where
- T: 'static,
-{
- /// Returns the child array builder as a mutable reference.
- ///
- /// This mutable reference can be used to append values into the child array builder,
- /// but you must call `append` to delimit each distinct list value.
- pub fn values(&mut self) -> &mut T {
- &mut self.values_builder
- }
-
- /// Finish the current variable-length list array slot
- #[inline]
- pub fn append(&mut self, is_valid: bool) -> Result<()> {
- self.offsets_builder
- .append(OffsetSize::from_usize(self.values_builder.len()).unwrap());
- self.bitmap_builder.append(is_valid);
- self.len += OffsetSize::one();
- Ok(())
- }
-
- /// Builds the `ListArray` and reset this builder.
- pub fn finish(&mut self) -> GenericListArray<OffsetSize> {
- let len = self.len();
- self.len = OffsetSize::zero();
- let values_arr = self
- .values_builder
- .as_any_mut()
- .downcast_mut::<T>()
- .unwrap()
- .finish();
- let values_data = values_arr.data();
-
- let offset_buffer = self.offsets_builder.finish();
- let null_bit_buffer = self.bitmap_builder.finish();
- self.offsets_builder.append(self.len);
- let field = Box::new(Field::new(
- "item",
- values_data.data_type().clone(),
- true, // TODO: find a consistent way of getting this
- ));
- let data_type = if OffsetSize::is_large() {
- DataType::LargeList(field)
- } else {
- DataType::List(field)
- };
- let data = ArrayData::builder(data_type)
- .len(len)
- .add_buffer(offset_buffer)
- .add_child_data(values_data.clone())
- .null_bit_buffer(null_bit_buffer)
- .build();
-
- GenericListArray::<OffsetSize>::from(data)
- }
-}
-
-pub type ListBuilder<T> = GenericListBuilder<i32, T>;
-pub type LargeListBuilder<T> = GenericListBuilder<i64, T>;
-
-/// Array builder for `ListArray`
-#[derive(Debug)]
-pub struct FixedSizeListBuilder<T: ArrayBuilder> {
- bitmap_builder: BooleanBufferBuilder,
- values_builder: T,
- len: usize,
- list_len: i32,
-}
-
-impl<T: ArrayBuilder> FixedSizeListBuilder<T> {
- /// Creates a new `FixedSizeListBuilder` from a given values array builder
- /// `length` is the number of values within each array
- pub fn new(values_builder: T, length: i32) -> Self {
- let capacity = values_builder.len();
- Self::with_capacity(values_builder, length, capacity)
- }
-
- /// Creates a new `FixedSizeListBuilder` from a given values array builder
- /// `length` is the number of values within each array
- /// `capacity` is the number of items to pre-allocate space for in this builder
- pub fn with_capacity(values_builder: T, length: i32, capacity: usize) -> Self {
- let mut offsets_builder = Int32BufferBuilder::new(capacity + 1);
- offsets_builder.append(0);
- Self {
- bitmap_builder: BooleanBufferBuilder::new(capacity),
- values_builder,
- len: 0,
- list_len: length,
- }
- }
-}
-
-impl<T: ArrayBuilder> ArrayBuilder for FixedSizeListBuilder<T>
-where
- T: 'static,
-{
- /// Returns the builder as a non-mutable `Any` reference.
- fn as_any(&self) -> &Any {
- self
- }
-
- /// Returns the builder as a mutable `Any` reference.
- fn as_any_mut(&mut self) -> &mut Any {
- self
- }
-
- /// Returns the boxed builder as a box of `Any`.
- fn into_box_any(self: Box<Self>) -> Box<Any> {
- self
- }
-
- /// Returns the number of array slots in the builder
- fn len(&self) -> usize {
- self.len
- }
-
- /// Returns whether the number of array slots is zero
- fn is_empty(&self) -> bool {
- self.len == 0
- }
-
- /// Builds the array and reset this builder.
- fn finish(&mut self) -> ArrayRef {
- Arc::new(self.finish())
- }
-}
-
-impl<T: ArrayBuilder> FixedSizeListBuilder<T>
-where
- T: 'static,
-{
- /// Returns the child array builder as a mutable reference.
- ///
- /// This mutable reference can be used to append values into the child array builder,
- /// but you must call `append` to delimit each distinct list value.
- pub fn values(&mut self) -> &mut T {
- &mut self.values_builder
- }
-
- pub fn value_length(&self) -> i32 {
- self.list_len
- }
-
- /// Finish the current variable-length list array slot
- #[inline]
- pub fn append(&mut self, is_valid: bool) -> Result<()> {
- self.bitmap_builder.append(is_valid);
- self.len += 1;
- Ok(())
- }
-
- /// Builds the `FixedSizeListBuilder` and reset this builder.
- pub fn finish(&mut self) -> FixedSizeListArray {
- let len = self.len();
- self.len = 0;
- let values_arr = self
- .values_builder
- .as_any_mut()
- .downcast_mut::<T>()
- .unwrap()
- .finish();
- let values_data = values_arr.data();
-
- // check that values_data length is multiple of len if we have data
- if len != 0 {
- assert!(
- values_data.len() / len == self.list_len as usize,
- "Values of FixedSizeList must have equal lengths, values have length {} and list has {}",
- values_data.len() / len,
- self.list_len
- );
- }
-
- let null_bit_buffer = self.bitmap_builder.finish();
- let data = ArrayData::builder(DataType::FixedSizeList(
- Box::new(Field::new("item", values_data.data_type().clone(), true)),
- self.list_len,
- ))
- .len(len)
- .add_child_data(values_data.clone())
- .null_bit_buffer(null_bit_buffer)
- .build();
-
- FixedSizeListArray::from(data)
- }
-}
-
-/// Array builder for `BinaryArray`
-#[derive(Debug)]
-pub struct GenericBinaryBuilder<OffsetSize: OffsetSizeTrait> {
- builder: GenericListBuilder<OffsetSize, UInt8Builder>,
-}
-
-pub type BinaryBuilder = GenericBinaryBuilder<i32>;
-pub type LargeBinaryBuilder = GenericBinaryBuilder<i64>;
-
-#[derive(Debug)]
-pub struct GenericStringBuilder<OffsetSize: OffsetSizeTrait> {
- builder: GenericListBuilder<OffsetSize, UInt8Builder>,
-}
-
-pub type StringBuilder = GenericStringBuilder<i32>;
-pub type LargeStringBuilder = GenericStringBuilder<i64>;
-
-#[derive(Debug)]
-pub struct FixedSizeBinaryBuilder {
- builder: FixedSizeListBuilder<UInt8Builder>,
-}
-
-#[derive(Debug)]
-pub struct DecimalBuilder {
- builder: FixedSizeListBuilder<UInt8Builder>,
- precision: usize,
- scale: usize,
-}
-
-impl<OffsetSize: BinaryOffsetSizeTrait> ArrayBuilder
- for GenericBinaryBuilder<OffsetSize>
-{
- /// Returns the builder as a non-mutable `Any` reference.
- fn as_any(&self) -> &Any {
- self
- }
-
- /// Returns the builder as a mutable `Any` reference.
- fn as_any_mut(&mut self) -> &mut Any {
- self
- }
-
- /// Returns the boxed builder as a box of `Any`.
- fn into_box_any(self: Box<Self>) -> Box<Any> {
- self
- }
-
- /// Returns the number of array slots in the builder
- fn len(&self) -> usize {
- self.builder.len()
- }
-
- /// Returns whether the number of array slots is zero
- fn is_empty(&self) -> bool {
- self.builder.is_empty()
- }
-
- /// Builds the array and reset this builder.
- fn finish(&mut self) -> ArrayRef {
- Arc::new(self.finish())
- }
-}
-
-impl<OffsetSize: StringOffsetSizeTrait> ArrayBuilder
- for GenericStringBuilder<OffsetSize>
-{
- /// Returns the builder as a non-mutable `Any` reference.
- fn as_any(&self) -> &Any {
- self
- }
-
- /// Returns the builder as a mutable `Any` reference.
- fn as_any_mut(&mut self) -> &mut Any {
- self
- }
-
- /// Returns the boxed builder as a box of `Any`.
- fn into_box_any(self: Box<Self>) -> Box<Any> {
- self
- }
-
- /// Returns the number of array slots in the builder
- fn len(&self) -> usize {
- self.builder.len()
- }
-
- /// Returns whether the number of array slots is zero
- fn is_empty(&self) -> bool {
- self.builder.is_empty()
- }
-
- /// Builds the array and reset this builder.
- fn finish(&mut self) -> ArrayRef {
- let a = GenericStringBuilder::<OffsetSize>::finish(self);
- Arc::new(a)
- }
-}
-
-impl ArrayBuilder for FixedSizeBinaryBuilder {
- /// Returns the builder as a non-mutable `Any` reference.
- fn as_any(&self) -> &Any {
- self
- }
-
- /// Returns the builder as a mutable `Any` reference.
- fn as_any_mut(&mut self) -> &mut Any {
- self
- }
-
- /// Returns the boxed builder as a box of `Any`.
- fn into_box_any(self: Box<Self>) -> Box<Any> {
- self
- }
-
- /// Returns the number of array slots in the builder
- fn len(&self) -> usize {
- self.builder.len()
- }
-
- /// Returns whether the number of array slots is zero
- fn is_empty(&self) -> bool {
- self.builder.is_empty()
- }
-
- /// Builds the array and reset this builder.
- fn finish(&mut self) -> ArrayRef {
- Arc::new(self.finish())
- }
-}
-
-impl ArrayBuilder for DecimalBuilder {
- /// Returns the builder as a non-mutable `Any` reference.
- fn as_any(&self) -> &Any {
- self
- }
-
- /// Returns the builder as a mutable `Any` reference.
- fn as_any_mut(&mut self) -> &mut Any {
- self
- }
-
- /// Returns the boxed builder as a box of `Any`.
- fn into_box_any(self: Box<Self>) -> Box<Any> {
- self
- }
-
- /// Returns the number of array slots in the builder
- fn len(&self) -> usize {
- self.builder.len()
- }
-
- /// Returns whether the number of array slots is zero
- fn is_empty(&self) -> bool {
- self.builder.is_empty()
- }
-
- /// Builds the array and reset this builder.
- fn finish(&mut self) -> ArrayRef {
- Arc::new(self.finish())
- }
-}
-
-impl<OffsetSize: BinaryOffsetSizeTrait> GenericBinaryBuilder<OffsetSize> {
- /// Creates a new `GenericBinaryBuilder`, `capacity` is the number of bytes in the values
- /// array
- pub fn new(capacity: usize) -> Self {
- let values_builder = UInt8Builder::new(capacity);
- Self {
- builder: GenericListBuilder::new(values_builder),
- }
- }
-
- /// Appends a single byte value into the builder's values array.
- ///
- /// Note, when appending individual byte values you must call `append` to delimit each
- /// distinct list value.
- #[inline]
- pub fn append_byte(&mut self, value: u8) -> Result<()> {
- self.builder.values().append_value(value)?;
- Ok(())
- }
-
- /// Appends a byte slice into the builder.
- ///
- /// Automatically calls the `append` method to delimit the slice appended in as a
- /// distinct array element.
- #[inline]
- pub fn append_value(&mut self, value: impl AsRef<[u8]>) -> Result<()> {
- self.builder.values().append_slice(value.as_ref())?;
- self.builder.append(true)?;
- Ok(())
- }
-
- /// Finish the current variable-length list array slot.
- #[inline]
- pub fn append(&mut self, is_valid: bool) -> Result<()> {
- self.builder.append(is_valid)
- }
-
- /// Append a null value to the array.
- #[inline]
- pub fn append_null(&mut self) -> Result<()> {
- self.append(false)
- }
-
- /// Builds the `BinaryArray` and reset this builder.
- pub fn finish(&mut self) -> GenericBinaryArray<OffsetSize> {
- GenericBinaryArray::<OffsetSize>::from(self.builder.finish())
- }
-}
-
-impl<OffsetSize: StringOffsetSizeTrait> GenericStringBuilder<OffsetSize> {
- /// Creates a new `StringBuilder`,
- /// `capacity` is the number of bytes of string data to pre-allocate space for in this builder
- pub fn new(capacity: usize) -> Self {
- let values_builder = UInt8Builder::new(capacity);
- Self {
- builder: GenericListBuilder::new(values_builder),
- }
- }
-
- /// Creates a new `StringBuilder`,
- /// `data_capacity` is the number of bytes of string data to pre-allocate space for in this builder
- /// `item_capacity` is the number of items to pre-allocate space for in this builder
- pub fn with_capacity(item_capacity: usize, data_capacity: usize) -> Self {
- let values_builder = UInt8Builder::new(data_capacity);
- Self {
- builder: GenericListBuilder::with_capacity(values_builder, item_capacity),
- }
- }
-
- /// Appends a string into the builder.
- ///
- /// Automatically calls the `append` method to delimit the string appended in as a
- /// distinct array element.
- #[inline]
- pub fn append_value(&mut self, value: impl AsRef<str>) -> Result<()> {
- self.builder
- .values()
- .append_slice(value.as_ref().as_bytes())?;
- self.builder.append(true)?;
- Ok(())
- }
-
- /// Finish the current variable-length list array slot.
- #[inline]
- pub fn append(&mut self, is_valid: bool) -> Result<()> {
- self.builder.append(is_valid)
- }
-
- /// Append a null value to the array.
- #[inline]
- pub fn append_null(&mut self) -> Result<()> {
- self.append(false)
- }
-
- /// Builds the `StringArray` and reset this builder.
- pub fn finish(&mut self) -> GenericStringArray<OffsetSize> {
- GenericStringArray::<OffsetSize>::from(self.builder.finish())
- }
-}
-
-impl FixedSizeBinaryBuilder {
- /// Creates a new `BinaryBuilder`, `capacity` is the number of bytes in the values
- /// array
- pub fn new(capacity: usize, byte_width: i32) -> Self {
- let values_builder = UInt8Builder::new(capacity);
- Self {
- builder: FixedSizeListBuilder::new(values_builder, byte_width),
- }
- }
-
- /// Appends a byte slice into the builder.
- ///
- /// Automatically calls the `append` method to delimit the slice appended in as a
- /// distinct array element.
- #[inline]
- pub fn append_value(&mut self, value: impl AsRef<[u8]>) -> Result<()> {
- if self.builder.value_length() != value.as_ref().len() as i32 {
- return Err(ArrowError::InvalidArgumentError(
- "Byte slice does not have the same length as FixedSizeBinaryBuilder value lengths".to_string()
- ));
- }
- self.builder.values().append_slice(value.as_ref())?;
- self.builder.append(true)
- }
-
- /// Append a null value to the array.
- #[inline]
- pub fn append_null(&mut self) -> Result<()> {
- let length: usize = self.builder.value_length() as usize;
- self.builder.values().append_slice(&vec![0u8; length][..])?;
- self.builder.append(false)
- }
-
- /// Builds the `FixedSizeBinaryArray` and reset this builder.
- pub fn finish(&mut self) -> FixedSizeBinaryArray {
- FixedSizeBinaryArray::from(self.builder.finish())
- }
-}
-
-impl DecimalBuilder {
- /// Creates a new `BinaryBuilder`, `capacity` is the number of bytes in the values
- /// array
- pub fn new(capacity: usize, precision: usize, scale: usize) -> Self {
- let values_builder = UInt8Builder::new(capacity);
- let byte_width = 16;
- Self {
- builder: FixedSizeListBuilder::new(values_builder, byte_width),
- precision,
- scale,
- }
- }
-
- /// Appends a byte slice into the builder.
- ///
- /// Automatically calls the `append` method to delimit the slice appended in as a
- /// distinct array element.
- #[inline]
- pub fn append_value(&mut self, value: i128) -> Result<()> {
- let value_as_bytes = Self::from_i128_to_fixed_size_bytes(
- value,
- self.builder.value_length() as usize,
- )?;
- if self.builder.value_length() != value_as_bytes.len() as i32 {
- return Err(ArrowError::InvalidArgumentError(
- "Byte slice does not have the same length as DecimalBuilder value lengths".to_string()
- ));
- }
- self.builder
- .values()
- .append_slice(value_as_bytes.as_slice())?;
- self.builder.append(true)
- }
-
- fn from_i128_to_fixed_size_bytes(v: i128, size: usize) -> Result<Vec<u8>> {
- if size > 16 {
- return Err(ArrowError::InvalidArgumentError(
- "DecimalBuilder only supports values up to 16 bytes.".to_string(),
- ));
- }
- let res = v.to_le_bytes();
- let start_byte = 16 - size;
- Ok(res[start_byte..16].to_vec())
- }
-
- /// Append a null value to the array.
- #[inline]
- pub fn append_null(&mut self) -> Result<()> {
- let length: usize = self.builder.value_length() as usize;
- self.builder.values().append_slice(&vec![0u8; length][..])?;
- self.builder.append(false)
- }
-
- /// Builds the `DecimalArray` and reset this builder.
- pub fn finish(&mut self) -> DecimalArray {
- DecimalArray::from_fixed_size_list_array(
- self.builder.finish(),
- self.precision,
- self.scale,
- )
- }
-}
-
-/// Array builder for Struct types.
-///
-/// Note that callers should make sure that methods of all the child field builders are
-/// properly called to maintain the consistency of the data structure.
-pub struct StructBuilder {
- fields: Vec<Field>,
- field_builders: Vec<Box<ArrayBuilder>>,
- bitmap_builder: BooleanBufferBuilder,
- len: usize,
-}
-
-impl fmt::Debug for StructBuilder {
- fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
- f.debug_struct("StructBuilder")
- .field("fields", &self.fields)
- .field("bitmap_builder", &self.bitmap_builder)
- .field("len", &self.len)
- .finish()
- }
-}
-
-impl ArrayBuilder for StructBuilder {
- /// Returns the number of array slots in the builder.
- ///
- /// Note that this always return the first child field builder's length, and it is
- /// the caller's responsibility to maintain the consistency that all the child field
- /// builder should have the equal number of elements.
- fn len(&self) -> usize {
- self.len
- }
-
- /// Returns whether the number of array slots is zero
- fn is_empty(&self) -> bool {
- self.len == 0
- }
-
- /// Builds the array.
- fn finish(&mut self) -> ArrayRef {
- Arc::new(self.finish())
- }
-
- /// Returns the builder as a non-mutable `Any` reference.
- ///
- /// This is most useful when one wants to call non-mutable APIs on a specific builder
- /// type. In this case, one can first cast this into a `Any`, and then use
- /// `downcast_ref` to get a reference on the specific builder.
- fn as_any(&self) -> &Any {
- self
- }
-
- /// Returns the builder as a mutable `Any` reference.
- ///
- /// This is most useful when one wants to call mutable APIs on a specific builder
- /// type. In this case, one can first cast this into a `Any`, and then use
- /// `downcast_mut` to get a reference on the specific builder.
- fn as_any_mut(&mut self) -> &mut Any {
- self
- }
-
- /// Returns the boxed builder as a box of `Any`.
- fn into_box_any(self: Box<Self>) -> Box<Any> {
- self
- }
-}
-
-/// Returns a builder with capacity `capacity` that corresponds to the datatype `DataType`
-/// This function is useful to construct arrays from an arbitrary vectors with known/expected
-/// schema.
-pub fn make_builder(datatype: &DataType, capacity: usize) -> Box<ArrayBuilder> {
- match datatype {
- DataType::Null => unimplemented!(),
- DataType::Boolean => Box::new(BooleanBuilder::new(capacity)),
- DataType::Int8 => Box::new(Int8Builder::new(capacity)),
- DataType::Int16 => Box::new(Int16Builder::new(capacity)),
- DataType::Int32 => Box::new(Int32Builder::new(capacity)),
- DataType::Int64 => Box::new(Int64Builder::new(capacity)),
- DataType::UInt8 => Box::new(UInt8Builder::new(capacity)),
- DataType::UInt16 => Box::new(UInt16Builder::new(capacity)),
- DataType::UInt32 => Box::new(UInt32Builder::new(capacity)),
- DataType::UInt64 => Box::new(UInt64Builder::new(capacity)),
- DataType::Float32 => Box::new(Float32Builder::new(capacity)),
- DataType::Float64 => Box::new(Float64Builder::new(capacity)),
- DataType::Binary => Box::new(BinaryBuilder::new(capacity)),
- DataType::FixedSizeBinary(len) => {
- Box::new(FixedSizeBinaryBuilder::new(capacity, *len))
- }
- DataType::Decimal(precision, scale) => {
- Box::new(DecimalBuilder::new(capacity, *precision, *scale))
- }
- DataType::Utf8 => Box::new(StringBuilder::new(capacity)),
- DataType::Date32 => Box::new(Date32Builder::new(capacity)),
- DataType::Date64 => Box::new(Date64Builder::new(capacity)),
- DataType::Time32(TimeUnit::Second) => {
- Box::new(Time32SecondBuilder::new(capacity))
- }
- DataType::Time32(TimeUnit::Millisecond) => {
- Box::new(Time32MillisecondBuilder::new(capacity))
- }
- DataType::Time64(TimeUnit::Microsecond) => {
- Box::new(Time64MicrosecondBuilder::new(capacity))
- }
- DataType::Time64(TimeUnit::Nanosecond) => {
- Box::new(Time64NanosecondBuilder::new(capacity))
- }
- DataType::Timestamp(TimeUnit::Second, _) => {
- Box::new(TimestampSecondBuilder::new(capacity))
- }
- DataType::Timestamp(TimeUnit::Millisecond, _) => {
- Box::new(TimestampMillisecondBuilder::new(capacity))
- }
- DataType::Timestamp(TimeUnit::Microsecond, _) => {
- Box::new(TimestampMicrosecondBuilder::new(capacity))
- }
- DataType::Timestamp(TimeUnit::Nanosecond, _) => {
- Box::new(TimestampNanosecondBuilder::new(capacity))
- }
- DataType::Interval(IntervalUnit::YearMonth) => {
- Box::new(IntervalYearMonthBuilder::new(capacity))
- }
- DataType::Interval(IntervalUnit::DayTime) => {
- Box::new(IntervalDayTimeBuilder::new(capacity))
- }
- DataType::Duration(TimeUnit::Second) => {
- Box::new(DurationSecondBuilder::new(capacity))
- }
- DataType::Duration(TimeUnit::Millisecond) => {
- Box::new(DurationMillisecondBuilder::new(capacity))
- }
- DataType::Duration(TimeUnit::Microsecond) => {
- Box::new(DurationMicrosecondBuilder::new(capacity))
- }
- DataType::Duration(TimeUnit::Nanosecond) => {
- Box::new(DurationNanosecondBuilder::new(capacity))
- }
- DataType::Struct(fields) => {
- Box::new(StructBuilder::from_fields(fields.clone(), capacity))
- }
- t => panic!("Data type {:?} is not currently supported", t),
- }
-}
-
-impl StructBuilder {
- pub fn new(fields: Vec<Field>, field_builders: Vec<Box<ArrayBuilder>>) -> Self {
- Self {
- fields,
- field_builders,
- bitmap_builder: BooleanBufferBuilder::new(0),
- len: 0,
- }
- }
-
- pub fn from_fields(fields: Vec<Field>, capacity: usize) -> Self {
- let mut builders = Vec::with_capacity(fields.len());
- for field in &fields {
- builders.push(make_builder(field.data_type(), capacity));
- }
- Self::new(fields, builders)
- }
-
- /// Returns a mutable reference to the child field builder at index `i`.
- /// Result will be `None` if the input type `T` provided doesn't match the actual
- /// field builder's type.
- pub fn field_builder<T: ArrayBuilder>(&mut self, i: usize) -> Option<&mut T> {
- self.field_builders[i].as_any_mut().downcast_mut::<T>()
- }
-
- /// Returns the number of fields for the struct this builder is building.
- pub fn num_fields(&self) -> usize {
- self.field_builders.len()
- }
-
- /// Appends an element (either null or non-null) to the struct. The actual elements
- /// should be appended for each child sub-array in a consistent way.
- #[inline]
- pub fn append(&mut self, is_valid: bool) -> Result<()> {
- self.bitmap_builder.append(is_valid);
- self.len += 1;
- Ok(())
- }
-
- /// Appends a null element to the struct.
- #[inline]
- pub fn append_null(&mut self) -> Result<()> {
- self.append(false)
- }
-
- /// Builds the `StructArray` and reset this builder.
- pub fn finish(&mut self) -> StructArray {
- let mut child_data = Vec::with_capacity(self.field_builders.len());
- for f in &mut self.field_builders {
- let arr = f.finish();
- child_data.push(arr.data().clone());
- }
-
- let null_bit_buffer = self.bitmap_builder.finish();
- let null_count = self.len - null_bit_buffer.count_set_bits();
- let mut builder = ArrayData::builder(DataType::Struct(self.fields.clone()))
- .len(self.len)
- .child_data(child_data);
- if null_count > 0 {
- builder = builder.null_bit_buffer(null_bit_buffer);
- }
-
- self.len = 0;
-
- StructArray::from(builder.build())
- }
-}
-
-/// `FieldData` is a helper struct to track the state of the fields in the `UnionBuilder`.
-#[derive(Debug)]
-struct FieldData {
- /// The type id for this field
- type_id: i8,
- /// The Arrow data type represented in the `values_buffer`, which is untyped
- data_type: DataType,
- /// A buffer containing the values for this field in raw bytes
- values_buffer: Option<MutableBuffer>,
- /// The number of array slots represented by the buffer
- slots: usize,
- /// A builder for the bitmap if required (for Sparse Unions)
- bitmap_builder: Option<BooleanBufferBuilder>,
-}
-
-impl FieldData {
- /// Creates a new `FieldData`.
- fn new(
- type_id: i8,
- data_type: DataType,
- bitmap_builder: Option<BooleanBufferBuilder>,
- ) -> Self {
- Self {
- type_id,
- data_type,
- values_buffer: Some(MutableBuffer::new(1)),
- slots: 0,
- bitmap_builder,
- }
- }
-
- /// Appends a single value to this `FieldData`'s `values_buffer`.
- #[allow(clippy::unnecessary_wraps)]
- fn append_to_values_buffer<T: ArrowPrimitiveType>(
- &mut self,
- v: T::Native,
- ) -> Result<()> {
- let values_buffer = self
- .values_buffer
- .take()
- .expect("Values buffer was never created");
- let mut builder: BufferBuilder<T::Native> =
- mutable_buffer_to_builder(values_buffer, self.slots);
- builder.append(v);
- let mutable_buffer = builder_to_mutable_buffer(builder);
- self.values_buffer = Some(mutable_buffer);
-
- self.slots += 1;
- if let Some(b) = &mut self.bitmap_builder {
- b.append(true)
- };
- Ok(())
- }
-
- /// Appends a null to this `FieldData`.
- #[allow(clippy::unnecessary_wraps)]
- fn append_null<T: ArrowPrimitiveType>(&mut self) -> Result<()> {
- if let Some(b) = &mut self.bitmap_builder {
- let values_buffer = self
- .values_buffer
- .take()
- .expect("Values buffer was never created");
- let mut builder: BufferBuilder<T::Native> =
- mutable_buffer_to_builder(values_buffer, self.slots);
- builder.advance(1);
- let mutable_buffer = builder_to_mutable_buffer(builder);
- self.values_buffer = Some(mutable_buffer);
- self.slots += 1;
- b.append(false);
- };
- Ok(())
- }
-
- /// Appends a null to this `FieldData` when the type is not known at compile time.
- ///
- /// As the main `append` method of `UnionBuilder` is generic, we need a way to append null
- /// slots to the fields that are not being appended to in the case of sparse unions. This
- /// method solves this problem by appending dynamically based on `DataType`.
- ///
- /// Note, this method does **not** update the length of the `UnionArray` (this is done by the
- /// main append operation) and assumes that it is called from a method that is generic over `T`
- /// where `T` satisfies the bound `ArrowPrimitiveType`.
- fn append_null_dynamic(&mut self) -> Result<()> {
- match self.data_type {
- DataType::Null => unimplemented!(),
- DataType::Int8 => self.append_null::<Int8Type>()?,
- DataType::Int16 => self.append_null::<Int16Type>()?,
- DataType::Int32
- | DataType::Date32
- | DataType::Time32(_)
- | DataType::Interval(IntervalUnit::YearMonth) => {
- self.append_null::<Int32Type>()?
- }
- DataType::Int64
- | DataType::Timestamp(_, _)
- | DataType::Date64
- | DataType::Time64(_)
- | DataType::Interval(IntervalUnit::DayTime)
- | DataType::Duration(_) => self.append_null::<Int64Type>()?,
- DataType::UInt8 => self.append_null::<UInt8Type>()?,
- DataType::UInt16 => self.append_null::<UInt16Type>()?,
- DataType::UInt32 => self.append_null::<UInt32Type>()?,
- DataType::UInt64 => self.append_null::<UInt64Type>()?,
- DataType::Float32 => self.append_null::<Float32Type>()?,
- DataType::Float64 => self.append_null::<Float64Type>()?,
- _ => unreachable!("All cases of types that satisfy the trait bounds over T are covered above."),
- };
- Ok(())
- }
-}
-
-/// Builder type for creating a new `UnionArray`.
-#[derive(Debug)]
-pub struct UnionBuilder {
- /// The current number of slots in the array
- len: usize,
- /// Maps field names to `FieldData` instances which track the builders for that field
- fields: HashMap<String, FieldData>,
- /// Builder to keep track of type ids
- type_id_builder: Int8BufferBuilder,
- /// Builder to keep track of offsets (`None` for sparse unions)
- value_offset_builder: Option<Int32BufferBuilder>,
- /// Optional builder for null slots
- bitmap_builder: Option<BooleanBufferBuilder>,
-}
-
-impl UnionBuilder {
- /// Creates a new dense array builder.
- pub fn new_dense(capacity: usize) -> Self {
- Self {
- len: 0,
- fields: HashMap::default(),
- type_id_builder: Int8BufferBuilder::new(capacity),
- value_offset_builder: Some(Int32BufferBuilder::new(capacity)),
- bitmap_builder: None,
- }
- }
-
- /// Creates a new sparse array builder.
- pub fn new_sparse(capacity: usize) -> Self {
- Self {
- len: 0,
- fields: HashMap::default(),
- type_id_builder: Int8BufferBuilder::new(capacity),
- value_offset_builder: None,
- bitmap_builder: None,
- }
- }
-
- /// Appends a null to this builder.
- #[inline]
- pub fn append_null(&mut self) -> Result<()> {
- if self.bitmap_builder.is_none() {
- let mut builder = BooleanBufferBuilder::new(self.len + 1);
- for _ in 0..self.len {
- builder.append(true);
- }
- self.bitmap_builder = Some(builder)
- }
- self.bitmap_builder
- .as_mut()
- .expect("Cannot be None")
- .append(false);
-
- self.type_id_builder.append(i8::default());
-
- // Handle sparse union
- if self.value_offset_builder.is_none() {
- for (_, fd) in self.fields.iter_mut() {
- fd.append_null_dynamic()?;
- }
- }
- self.len += 1;
- Ok(())
- }
-
- /// Appends a value to this builder.
- #[inline]
- pub fn append<T: ArrowPrimitiveType>(
- &mut self,
- type_name: &str,
- v: T::Native,
- ) -> Result<()> {
- let type_name = type_name.to_string();
-
- let mut field_data = match self.fields.remove(&type_name) {
- Some(data) => data,
- None => match self.value_offset_builder {
- Some(_) => FieldData::new(self.fields.len() as i8, T::DATA_TYPE, None),
- None => {
- let mut fd = FieldData::new(
- self.fields.len() as i8,
- T::DATA_TYPE,
- Some(BooleanBufferBuilder::new(1)),
- );
- for _ in 0..self.len {
- fd.append_null::<T>()?;
- }
- fd
- }
- },
- };
- self.type_id_builder.append(field_data.type_id);
-
- match &mut self.value_offset_builder {
- // Dense Union
- Some(offset_builder) => {
- offset_builder.append(field_data.slots as i32);
- }
- // Sparse Union
- None => {
- for (name, fd) in self.fields.iter_mut() {
- if name != &type_name {
- fd.append_null_dynamic()?;
- }
- }
- }
- }
- field_data.append_to_values_buffer::<T>(v)?;
- self.fields.insert(type_name, field_data);
-
- // Update the bitmap builder if it exists
- if let Some(b) = &mut self.bitmap_builder {
- b.append(true);
- }
- self.len += 1;
- Ok(())
- }
-
- /// Builds this builder creating a new `UnionArray`.
- pub fn build(mut self) -> Result<UnionArray> {
- let type_id_buffer = self.type_id_builder.finish();
- let value_offsets_buffer = self.value_offset_builder.map(|mut b| b.finish());
- let mut children = Vec::new();
- for (
- name,
- FieldData {
- type_id,
- data_type,
- values_buffer,
- slots,
- bitmap_builder,
- },
- ) in self.fields.into_iter()
- {
- let buffer = values_buffer
- .expect("The `values_buffer` should only ever be None inside the `append` method.")
- .into();
- let arr_data_builder = ArrayDataBuilder::new(data_type.clone())
- .add_buffer(buffer)
- .len(slots);
- // .build();
- let arr_data_ref = match bitmap_builder {
- Some(mut bb) => arr_data_builder.null_bit_buffer(bb.finish()).build(),
- None => arr_data_builder.build(),
- };
- let array_ref = make_array(arr_data_ref);
- children.push((type_id, (Field::new(&name, data_type, false), array_ref)))
- }
-
- children.sort_by(|a, b| {
- a.0.partial_cmp(&b.0)
- .expect("This will never be None as type ids are always i8 values.")
- });
- let children: Vec<_> = children.into_iter().map(|(_, b)| b).collect();
- let bitmap = self.bitmap_builder.map(|mut b| b.finish());
-
- UnionArray::try_new(type_id_buffer, value_offsets_buffer, children, bitmap)
- }
-}
-
-/// Array builder for `DictionaryArray`. For example to map a set of byte indices
-/// to f32 values. Note that the use of a `HashMap` here will not scale to very large
-/// arrays or result in an ordered dictionary.
-#[derive(Debug)]
-pub struct PrimitiveDictionaryBuilder<K, V>
-where
- K: ArrowPrimitiveType,
- V: ArrowPrimitiveType,
-{
- keys_builder: PrimitiveBuilder<K>,
- values_builder: PrimitiveBuilder<V>,
- map: HashMap<Box<[u8]>, K::Native>,
-}
-
-impl<K, V> PrimitiveDictionaryBuilder<K, V>
-where
- K: ArrowPrimitiveType,
- V: ArrowPrimitiveType,
-{
- /// Creates a new `PrimitiveDictionaryBuilder` from a keys builder and a value builder.
- pub fn new(
- keys_builder: PrimitiveBuilder<K>,
- values_builder: PrimitiveBuilder<V>,
- ) -> Self {
- Self {
- keys_builder,
- values_builder,
- map: HashMap::new(),
- }
- }
-}
-
-impl<K, V> ArrayBuilder for PrimitiveDictionaryBuilder<K, V>
-where
- K: ArrowPrimitiveType,
- V: ArrowPrimitiveType,
-{
- /// Returns the builder as an non-mutable `Any` reference.
- fn as_any(&self) -> &Any {
- self
- }
-
- /// Returns the builder as an mutable `Any` reference.
- fn as_any_mut(&mut self) -> &mut Any {
- self
- }
-
- /// Returns the boxed builder as a box of `Any`.
- fn into_box_any(self: Box<Self>) -> Box<Any> {
- self
- }
-
- /// Returns the number of array slots in the builder
- fn len(&self) -> usize {
- self.keys_builder.len()
- }
-
- /// Returns whether the number of array slots is zero
- fn is_empty(&self) -> bool {
- self.keys_builder.is_empty()
- }
-
- /// Builds the array and reset this builder.
- fn finish(&mut self) -> ArrayRef {
- Arc::new(self.finish())
- }
-}
-
-impl<K, V> PrimitiveDictionaryBuilder<K, V>
-where
- K: ArrowPrimitiveType,
- V: ArrowPrimitiveType,
-{
- /// Append a primitive value to the array. Return an existing index
- /// if already present in the values array or a new index if the
- /// value is appended to the values array.
- #[inline]
- pub fn append(&mut self, value: V::Native) -> Result<K::Native> {
- if let Some(&key) = self.map.get(value.to_byte_slice()) {
- // Append existing value.
- self.keys_builder.append_value(key)?;
- Ok(key)
- } else {
- // Append new value.
- let key = K::Native::from_usize(self.values_builder.len())
- .ok_or(ArrowError::DictionaryKeyOverflowError)?;
- self.values_builder.append_value(value)?;
- self.keys_builder.append_value(key as K::Native)?;
- self.map.insert(value.to_byte_slice().into(), key);
- Ok(key)
- }
- }
-
- #[inline]
- pub fn append_null(&mut self) -> Result<()> {
- self.keys_builder.append_null()
- }
-
- /// Builds the `DictionaryArray` and reset this builder.
- pub fn finish(&mut self) -> DictionaryArray<K> {
- self.map.clear();
- let value_ref: ArrayRef = Arc::new(self.values_builder.finish());
- self.keys_builder.finish_dict(value_ref)
- }
-}
-
-/// Array builder for `DictionaryArray` that stores Strings. For example to map a set of byte indices
-/// to String values. Note that the use of a `HashMap` here will not scale to very large
-/// arrays or result in an ordered dictionary.
-///
-/// ```
-/// use arrow::{
-/// array::{
-/// Int8Array, StringArray,
-/// PrimitiveBuilder, StringBuilder, StringDictionaryBuilder,
-/// },
-/// datatypes::Int8Type,
-/// };
-///
-/// // Create a dictionary array indexed by bytes whose values are Strings.
-/// // It can thus hold up to 256 distinct string values.
-///
-/// let key_builder = PrimitiveBuilder::<Int8Type>::new(100);
-/// let value_builder = StringBuilder::new(100);
-/// let mut builder = StringDictionaryBuilder::new(key_builder, value_builder);
-///
-/// // The builder builds the dictionary value by value
-/// builder.append("abc").unwrap();
-/// builder.append_null().unwrap();
-/// builder.append("def").unwrap();
-/// builder.append("def").unwrap();
-/// builder.append("abc").unwrap();
-/// let array = builder.finish();
-///
-/// assert_eq!(
-/// array.keys(),
-/// &Int8Array::from(vec![Some(0), None, Some(1), Some(1), Some(0)])
-/// );
-///
-/// // Values are polymorphic and so require a downcast.
-/// let av = array.values();
-/// let ava: &StringArray = av.as_any().downcast_ref::<StringArray>().unwrap();
-///
-/// assert_eq!(ava.value(0), "abc");
-/// assert_eq!(ava.value(1), "def");
-///
-/// ```
-#[derive(Debug)]
-pub struct StringDictionaryBuilder<K>
-where
- K: ArrowDictionaryKeyType,
-{
- keys_builder: PrimitiveBuilder<K>,
- values_builder: StringBuilder,
- map: HashMap<Box<[u8]>, K::Native>,
-}
-
-impl<K> StringDictionaryBuilder<K>
-where
- K: ArrowDictionaryKeyType,
-{
- /// Creates a new `StringDictionaryBuilder` from a keys builder and a value builder.
- pub fn new(keys_builder: PrimitiveBuilder<K>, values_builder: StringBuilder) -> Self {
- Self {
- keys_builder,
- values_builder,
- map: HashMap::new(),
- }
- }
-
- /// Creates a new `StringDictionaryBuilder` from a keys builder and a dictionary
- /// which is initialized with the given values.
- /// The indices of those dictionary values are used as keys.
- ///
- /// # Example
- ///
- /// ```
- /// use arrow::datatypes::Int16Type;
- /// use arrow::array::{StringArray, StringDictionaryBuilder, PrimitiveBuilder, Int16Array};
- /// use std::convert::TryFrom;
- ///
- /// let dictionary_values = StringArray::from(vec![None, Some("abc"), Some("def")]);
- ///
- /// let mut builder = StringDictionaryBuilder::new_with_dictionary(PrimitiveBuilder::<Int16Type>::new(3), &dictionary_values).unwrap();
- /// builder.append("def").unwrap();
- /// builder.append_null().unwrap();
- /// builder.append("abc").unwrap();
- ///
- /// let dictionary_array = builder.finish();
- ///
- /// let keys = dictionary_array.keys();
- ///
- /// assert_eq!(keys, &Int16Array::from(vec![Some(2), None, Some(1)]));
- /// ```
- pub fn new_with_dictionary(
- keys_builder: PrimitiveBuilder<K>,
- dictionary_values: &StringArray,
- ) -> Result<Self> {
- let dict_len = dictionary_values.len();
- let mut values_builder =
- StringBuilder::with_capacity(dict_len, dictionary_values.value_data().len());
- let mut map: HashMap<Box<[u8]>, K::Native> = HashMap::with_capacity(dict_len);
- for i in 0..dict_len {
- if dictionary_values.is_valid(i) {
- let value = dictionary_values.value(i);
- map.insert(
- value.as_bytes().into(),
- K::Native::from_usize(i)
- .ok_or(ArrowError::DictionaryKeyOverflowError)?,
- );
- values_builder.append_value(value)?;
- } else {
- values_builder.append_null()?;
- }
- }
- Ok(Self {
- keys_builder,
- values_builder,
- map,
- })
- }
-}
-
-impl<K> ArrayBuilder for StringDictionaryBuilder<K>
-where
- K: ArrowDictionaryKeyType,
-{
- /// Returns the builder as an non-mutable `Any` reference.
- fn as_any(&self) -> &Any {
- self
- }
-
- /// Returns the builder as an mutable `Any` reference.
- fn as_any_mut(&mut self) -> &mut Any {
- self
- }
-
- /// Returns the boxed builder as a box of `Any`.
- fn into_box_any(self: Box<Self>) -> Box<Any> {
- self
- }
-
- /// Returns the number of array slots in the builder
- fn len(&self) -> usize {
- self.keys_builder.len()
- }
-
- /// Returns whether the number of array slots is zero
- fn is_empty(&self) -> bool {
- self.keys_builder.is_empty()
- }
-
- /// Builds the array and reset this builder.
- fn finish(&mut self) -> ArrayRef {
- Arc::new(self.finish())
- }
-}
-
-impl<K> StringDictionaryBuilder<K>
-where
- K: ArrowDictionaryKeyType,
-{
- /// Append a primitive value to the array. Return an existing index
- /// if already present in the values array or a new index if the
- /// value is appended to the values array.
- pub fn append(&mut self, value: impl AsRef<str>) -> Result<K::Native> {
- if let Some(&key) = self.map.get(value.as_ref().as_bytes()) {
- // Append existing value.
- self.keys_builder.append_value(key)?;
- Ok(key)
- } else {
- // Append new value.
- let key = K::Native::from_usize(self.values_builder.len())
- .ok_or(ArrowError::DictionaryKeyOverflowError)?;
- self.values_builder.append_value(value.as_ref())?;
- self.keys_builder.append_value(key as K::Native)?;
- self.map.insert(value.as_ref().as_bytes().into(), key);
- Ok(key)
- }
- }
-
- #[inline]
- pub fn append_null(&mut self) -> Result<()> {
- self.keys_builder.append_null()
- }
-
- /// Builds the `DictionaryArray` and reset this builder.
- pub fn finish(&mut self) -> DictionaryArray<K> {
- self.map.clear();
- let value_ref: ArrayRef = Arc::new(self.values_builder.finish());
- self.keys_builder.finish_dict(value_ref)
- }
-}
-
-#[cfg(test)]
-mod tests {
- use super::*;
-
- use crate::array::Array;
- use crate::bitmap::Bitmap;
-
- #[test]
- fn test_builder_i32_empty() {
- let mut b = Int32BufferBuilder::new(5);
- assert_eq!(0, b.len());
- assert_eq!(16, b.capacity());
- let a = b.finish();
- assert_eq!(0, a.len());
- }
-
- #[test]
- fn test_builder_i32_alloc_zero_bytes() {
- let mut b = Int32BufferBuilder::new(0);
- b.append(123);
- let a = b.finish();
- assert_eq!(4, a.len());
- }
-
- #[test]
- fn test_builder_i32() {
- let mut b = Int32BufferBuilder::new(5);
- for i in 0..5 {
- b.append(i);
- }
- assert_eq!(16, b.capacity());
- let a = b.finish();
- assert_eq!(20, a.len());
- }
-
- #[test]
- fn test_builder_i32_grow_buffer() {
- let mut b = Int32BufferBuilder::new(2);
- assert_eq!(16, b.capacity());
- for i in 0..20 {
- b.append(i);
- }
- assert_eq!(32, b.capacity());
- let a = b.finish();
- assert_eq!(80, a.len());
- }
-
- #[test]
- fn test_builder_finish() {
- let mut b = Int32BufferBuilder::new(5);
- assert_eq!(16, b.capacity());
- for i in 0..10 {
- b.append(i);
- }
- let mut a = b.finish();
- assert_eq!(40, a.len());
- assert_eq!(0, b.len());
- assert_eq!(0, b.capacity());
-
- // Try build another buffer after cleaning up.
- for i in 0..20 {
- b.append(i)
- }
- assert_eq!(32, b.capacity());
- a = b.finish();
- assert_eq!(80, a.len());
- }
-
- #[test]
- fn test_reserve() {
- let mut b = UInt8BufferBuilder::new(2);
- assert_eq!(64, b.capacity());
- b.reserve(64);
- assert_eq!(64, b.capacity());
- b.reserve(65);
- assert_eq!(128, b.capacity());
-
- let mut b = Int32BufferBuilder::new(2);
- assert_eq!(16, b.capacity());
- b.reserve(16);
- assert_eq!(16, b.capacity());
- b.reserve(17);
- assert_eq!(32, b.capacity());
- }
-
- #[test]
- fn test_append_slice() {
- let mut b = UInt8BufferBuilder::new(0);
- b.append_slice(b"Hello, ");
- b.append_slice(b"World!");
- let buffer = b.finish();
- assert_eq!(13, buffer.len());
-
- let mut b = Int32BufferBuilder::new(0);
- b.append_slice(&[32, 54]);
- let buffer = b.finish();
- assert_eq!(8, buffer.len());
- }
-
- #[test]
- fn test_append_values() -> Result<()> {
- let mut a = Int8Builder::new(0);
- a.append_value(1)?;
- a.append_null()?;
- a.append_value(-2)?;
- assert_eq!(a.len(), 3);
-
- // append values
- let values = &[1, 2, 3, 4];
- let is_valid = &[true, true, false, true];
- a.append_values(values, is_valid)?;
-
- assert_eq!(a.len(), 7);
- let array = a.finish();
- assert_eq!(array.value(0), 1);
- assert_eq!(array.is_null(1), true);
- assert_eq!(array.value(2), -2);
- assert_eq!(array.value(3), 1);
- assert_eq!(array.value(4), 2);
- assert_eq!(array.is_null(5), true);
- assert_eq!(array.value(6), 4);
-
- Ok(())
- }
-
- #[test]
- fn test_boolean_buffer_builder_write_bytes() {
- let mut b = BooleanBufferBuilder::new(4);
- b.append(false);
- b.append(true);
- b.append(false);
- b.append(true);
- assert_eq!(4, b.len());
- assert_eq!(512, b.capacity());
- let buffer = b.finish();
- assert_eq!(1, buffer.len());
-
- let mut b = BooleanBufferBuilder::new(4);
- b.append_slice(&[false, true, false, true]);
- assert_eq!(4, b.len());
- assert_eq!(512, b.capacity());
- let buffer = b.finish();
- assert_eq!(1, buffer.len());
- }
-
- #[test]
- fn test_boolean_buffer_builder_unset_first_bit() {
- let mut buffer = BooleanBufferBuilder::new(4);
- buffer.append(true);
- buffer.append(true);
- buffer.append(false);
- buffer.append(true);
- buffer.set_bit(0, false);
- assert_eq!(buffer.len(), 4);
- assert_eq!(buffer.finish().as_slice(), &[0b1010_u8]);
- }
-
- #[test]
- fn test_boolean_buffer_builder_unset_last_bit() {
- let mut buffer = BooleanBufferBuilder::new(4);
- buffer.append(true);
- buffer.append(true);
- buffer.append(false);
- buffer.append(true);
- buffer.set_bit(3, false);
- assert_eq!(buffer.len(), 4);
- assert_eq!(buffer.finish().as_slice(), &[0b0011_u8]);
- }
-
- #[test]
- fn test_boolean_buffer_builder_unset_an_inner_bit() {
- let mut buffer = BooleanBufferBuilder::new(5);
- buffer.append(true);
- buffer.append(true);
- buffer.append(false);
- buffer.append(true);
- buffer.set_bit(1, false);
- assert_eq!(buffer.len(), 4);
- assert_eq!(buffer.finish().as_slice(), &[0b1001_u8]);
- }
-
- #[test]
- fn test_boolean_buffer_builder_unset_several_bits() {
- let mut buffer = BooleanBufferBuilder::new(5);
- buffer.append(true);
- buffer.append(true);
- buffer.append(true);
- buffer.append(false);
- buffer.append(true);
- buffer.set_bit(1, false);
- buffer.set_bit(2, false);
- assert_eq!(buffer.len(), 5);
- assert_eq!(buffer.finish().as_slice(), &[0b10001_u8]);
- }
-
- #[test]
- fn test_boolean_buffer_builder_unset_several_bits_bigger_than_one_byte() {
- let mut buffer = BooleanBufferBuilder::new(16);
- buffer.append_n(10, true);
- buffer.set_bit(0, false);
- buffer.set_bit(3, false);
- buffer.set_bit(9, false);
- assert_eq!(buffer.len(), 10);
- assert_eq!(buffer.finish().as_slice(), &[0b11110110_u8, 0b01_u8]);
- }
-
- #[test]
- fn test_boolean_buffer_builder_flip_several_bits_bigger_than_one_byte() {
- let mut buffer = BooleanBufferBuilder::new(16);
- buffer.append_n(5, true);
- buffer.append_n(5, false);
- buffer.append_n(5, true);
- buffer.set_bit(0, false);
- buffer.set_bit(3, false);
- buffer.set_bit(9, false);
- buffer.set_bit(6, true);
- buffer.set_bit(14, true);
- buffer.set_bit(13, false);
- assert_eq!(buffer.len(), 15);
- assert_eq!(buffer.finish().as_slice(), &[0b01010110_u8, 0b1011100_u8]);
- }
-
- #[test]
- fn test_boolean_array_builder_append_slice() {
- let arr1 =
- BooleanArray::from(vec![Some(true), Some(false), None, None, Some(false)]);
-
- let mut builder = BooleanArray::builder(0);
- builder.append_slice(&[true, false]).unwrap();
- builder.append_null().unwrap();
- builder.append_null().unwrap();
- builder.append_value(false).unwrap();
- let arr2 = builder.finish();
-
- assert_eq!(arr1, arr2);
- }
-
- #[test]
- fn test_boolean_array_builder_append_slice_large() {
- let arr1 = BooleanArray::from(vec![true; 513]);
-
- let mut builder = BooleanArray::builder(512);
- builder.append_slice(&[true; 513]).unwrap();
- let arr2 = builder.finish();
-
- assert_eq!(arr1, arr2);
- }
-
- #[test]
- fn test_boolean_builder_increases_buffer_len() {
- // 00000010 01001000
- let buf = Buffer::from([72_u8, 2_u8]);
- let mut builder = BooleanBufferBuilder::new(8);
-
- for i in 0..16 {
- if i == 3 || i == 6 || i == 9 {
- builder.append(true);
- } else {
- builder.append(false);
- }
- }
- let buf2 = builder.finish();
-
- assert_eq!(buf.len(), buf2.len());
- assert_eq!(buf.as_slice(), buf2.as_slice());
- }
-
- #[test]
- fn test_primitive_array_builder_i32() {
- let mut builder = Int32Array::builder(5);
- for i in 0..5 {
- builder.append_value(i).unwrap();
- }
- let arr = builder.finish();
- assert_eq!(5, arr.len());
- assert_eq!(0, arr.offset());
- assert_eq!(0, arr.null_count());
- for i in 0..5 {
- assert!(!arr.is_null(i));
- assert!(arr.is_valid(i));
- assert_eq!(i as i32, arr.value(i));
- }
- }
-
- #[test]
- fn test_primitive_array_builder_date32() {
- let mut builder = Date32Array::builder(5);
- for i in 0..5 {
- builder.append_value(i).unwrap();
- }
- let arr = builder.finish();
- assert_eq!(5, arr.len());
- assert_eq!(0, arr.offset());
- assert_eq!(0, arr.null_count());
- for i in 0..5 {
- assert!(!arr.is_null(i));
- assert!(arr.is_valid(i));
- assert_eq!(i as i32, arr.value(i));
- }
- }
-
- #[test]
- fn test_primitive_array_builder_timestamp_second() {
- let mut builder = TimestampSecondArray::builder(5);
- for i in 0..5 {
- builder.append_value(i).unwrap();
- }
- let arr = builder.finish();
- assert_eq!(5, arr.len());
- assert_eq!(0, arr.offset());
- assert_eq!(0, arr.null_count());
- for i in 0..5 {
- assert!(!arr.is_null(i));
- assert!(arr.is_valid(i));
- assert_eq!(i as i64, arr.value(i));
- }
- }
-
- #[test]
- fn test_primitive_array_builder_bool() {
- // 00000010 01001000
- let buf = Buffer::from([72_u8, 2_u8]);
- let mut builder = BooleanArray::builder(10);
- for i in 0..10 {
- if i == 3 || i == 6 || i == 9 {
- builder.append_value(true).unwrap();
- } else {
- builder.append_value(false).unwrap();
- }
- }
-
- let arr = builder.finish();
- assert_eq!(&buf, arr.values());
- assert_eq!(10, arr.len());
- assert_eq!(0, arr.offset());
- assert_eq!(0, arr.null_count());
- for i in 0..10 {
- assert!(!arr.is_null(i));
- assert!(arr.is_valid(i));
- assert_eq!(i == 3 || i == 6 || i == 9, arr.value(i), "failed at {}", i)
- }
- }
-
- #[test]
- fn test_primitive_array_builder_append_option() {
- let arr1 = Int32Array::from(vec![Some(0), None, Some(2), None, Some(4)]);
-
- let mut builder = Int32Array::builder(5);
- builder.append_option(Some(0)).unwrap();
- builder.append_option(None).unwrap();
- builder.append_option(Some(2)).unwrap();
- builder.append_option(None).unwrap();
- builder.append_option(Some(4)).unwrap();
- let arr2 = builder.finish();
-
- assert_eq!(arr1.len(), arr2.len());
- assert_eq!(arr1.offset(), arr2.offset());
- assert_eq!(arr1.null_count(), arr2.null_count());
- for i in 0..5 {
- assert_eq!(arr1.is_null(i), arr2.is_null(i));
- assert_eq!(arr1.is_valid(i), arr2.is_valid(i));
- if arr1.is_valid(i) {
- assert_eq!(arr1.value(i), arr2.value(i));
- }
- }
- }
-
- #[test]
- fn test_primitive_array_builder_append_null() {
- let arr1 = Int32Array::from(vec![Some(0), Some(2), None, None, Some(4)]);
-
- let mut builder = Int32Array::builder(5);
- builder.append_value(0).unwrap();
- builder.append_value(2).unwrap();
- builder.append_null().unwrap();
- builder.append_null().unwrap();
- builder.append_value(4).unwrap();
- let arr2 = builder.finish();
-
- assert_eq!(arr1.len(), arr2.len());
- assert_eq!(arr1.offset(), arr2.offset());
- assert_eq!(arr1.null_count(), arr2.null_count());
- for i in 0..5 {
- assert_eq!(arr1.is_null(i), arr2.is_null(i));
- assert_eq!(arr1.is_valid(i), arr2.is_valid(i));
- if arr1.is_valid(i) {
- assert_eq!(arr1.value(i), arr2.value(i));
- }
- }
- }
-
- #[test]
- fn test_primitive_array_builder_append_slice() {
- let arr1 = Int32Array::from(vec![Some(0), Some(2), None, None, Some(4)]);
-
- let mut builder = Int32Array::builder(5);
- builder.append_slice(&[0, 2]).unwrap();
- builder.append_null().unwrap();
- builder.append_null().unwrap();
- builder.append_value(4).unwrap();
- let arr2 = builder.finish();
-
- assert_eq!(arr1.len(), arr2.len());
- assert_eq!(arr1.offset(), arr2.offset());
- assert_eq!(arr1.null_count(), arr2.null_count());
- for i in 0..5 {
- assert_eq!(arr1.is_null(i), arr2.is_null(i));
- assert_eq!(arr1.is_valid(i), arr2.is_valid(i));
- if arr1.is_valid(i) {
- assert_eq!(arr1.value(i), arr2.value(i));
- }
- }
- }
-
- #[test]
- fn test_primitive_array_builder_finish() {
- let mut builder = Int32Builder::new(5);
- builder.append_slice(&[2, 4, 6, 8]).unwrap();
- let mut arr = builder.finish();
- assert_eq!(4, arr.len());
- assert_eq!(0, builder.len());
-
- builder.append_slice(&[1, 3, 5, 7, 9]).unwrap();
- arr = builder.finish();
- assert_eq!(5, arr.len());
- assert_eq!(0, builder.len());
- }
-
- #[test]
- fn test_list_array_builder() {
- let values_builder = Int32Builder::new(10);
- let mut builder = ListBuilder::new(values_builder);
-
- // [[0, 1, 2], [3, 4, 5], [6, 7]]
- builder.values().append_value(0).unwrap();
- builder.values().append_value(1).unwrap();
- builder.values().append_value(2).unwrap();
- builder.append(true).unwrap();
- builder.values().append_value(3).unwrap();
- builder.values().append_value(4).unwrap();
- builder.values().append_value(5).unwrap();
- builder.append(true).unwrap();
- builder.values().append_value(6).unwrap();
- builder.values().append_value(7).unwrap();
- builder.append(true).unwrap();
- let list_array = builder.finish();
-
- let values = list_array.values().data().buffers()[0].clone();
- assert_eq!(Buffer::from_slice_ref(&[0, 1, 2, 3, 4, 5, 6, 7]), values);
- assert_eq!(
- Buffer::from_slice_ref(&[0, 3, 6, 8]),
- list_array.data().buffers()[0].clone()
- );
- assert_eq!(DataType::Int32, list_array.value_type());
- assert_eq!(3, list_array.len());
- assert_eq!(0, list_array.null_count());
- assert_eq!(6, list_array.value_offsets()[2]);
- assert_eq!(2, list_array.value_length(2));
- for i in 0..3 {
- assert!(list_array.is_valid(i));
- assert!(!list_array.is_null(i));
- }
- }
-
- #[test]
- fn test_large_list_array_builder() {
- let values_builder = Int32Builder::new(10);
- let mut builder = LargeListBuilder::new(values_builder);
-
- // [[0, 1, 2], [3, 4, 5], [6, 7]]
- builder.values().append_value(0).unwrap();
- builder.values().append_value(1).unwrap();
- builder.values().append_value(2).unwrap();
- builder.append(true).unwrap();
- builder.values().append_value(3).unwrap();
- builder.values().append_value(4).unwrap();
- builder.values().append_value(5).unwrap();
- builder.append(true).unwrap();
- builder.values().append_value(6).unwrap();
- builder.values().append_value(7).unwrap();
- builder.append(true).unwrap();
- let list_array = builder.finish();
-
- let values = list_array.values().data().buffers()[0].clone();
- assert_eq!(Buffer::from_slice_ref(&[0, 1, 2, 3, 4, 5, 6, 7]), values);
- assert_eq!(
- Buffer::from_slice_ref(&[0i64, 3, 6, 8]),
- list_array.data().buffers()[0].clone()
- );
- assert_eq!(DataType::Int32, list_array.value_type());
- assert_eq!(3, list_array.len());
- assert_eq!(0, list_array.null_count());
- assert_eq!(6, list_array.value_offsets()[2]);
- assert_eq!(2, list_array.value_length(2));
- for i in 0..3 {
- assert!(list_array.is_valid(i));
- assert!(!list_array.is_null(i));
- }
- }
-
- #[test]
- fn test_list_array_builder_nulls() {
- let values_builder = Int32Builder::new(10);
- let mut builder = ListBuilder::new(values_builder);
-
- // [[0, 1, 2], null, [3, null, 5], [6, 7]]
- builder.values().append_value(0).unwrap();
- builder.values().append_value(1).unwrap();
- builder.values().append_value(2).unwrap();
- builder.append(true).unwrap();
- builder.append(false).unwrap();
- builder.values().append_value(3).unwrap();
- builder.values().append_null().unwrap();
- builder.values().append_value(5).unwrap();
- builder.append(true).unwrap();
- builder.values().append_value(6).unwrap();
- builder.values().append_value(7).unwrap();
- builder.append(true).unwrap();
- let list_array = builder.finish();
-
- assert_eq!(DataType::Int32, list_array.value_type());
- assert_eq!(4, list_array.len());
- assert_eq!(1, list_array.null_count());
- assert_eq!(3, list_array.value_offsets()[2]);
- assert_eq!(3, list_array.value_length(2));
- }
-
- #[test]
- fn test_large_list_array_builder_nulls() {
- let values_builder = Int32Builder::new(10);
- let mut builder = LargeListBuilder::new(values_builder);
-
- // [[0, 1, 2], null, [3, null, 5], [6, 7]]
- builder.values().append_value(0).unwrap();
- builder.values().append_value(1).unwrap();
- builder.values().append_value(2).unwrap();
- builder.append(true).unwrap();
- builder.append(false).unwrap();
- builder.values().append_value(3).unwrap();
- builder.values().append_null().unwrap();
- builder.values().append_value(5).unwrap();
- builder.append(true).unwrap();
- builder.values().append_value(6).unwrap();
- builder.values().append_value(7).unwrap();
- builder.append(true).unwrap();
- let list_array = builder.finish();
-
- assert_eq!(DataType::Int32, list_array.value_type());
- assert_eq!(4, list_array.len());
- assert_eq!(1, list_array.null_count());
- assert_eq!(3, list_array.value_offsets()[2]);
- assert_eq!(3, list_array.value_length(2));
- }
-
- #[test]
- fn test_fixed_size_list_array_builder() {
- let values_builder = Int32Builder::new(10);
- let mut builder = FixedSizeListBuilder::new(values_builder, 3);
-
- // [[0, 1, 2], null, [3, null, 5], [6, 7, null]]
- builder.values().append_value(0).unwrap();
- builder.values().append_value(1).unwrap();
- builder.values().append_value(2).unwrap();
- builder.append(true).unwrap();
- builder.values().append_null().unwrap();
- builder.values().append_null().unwrap();
- builder.values().append_null().unwrap();
- builder.append(false).unwrap();
- builder.values().append_value(3).unwrap();
- builder.values().append_null().unwrap();
- builder.values().append_value(5).unwrap();
- builder.append(true).unwrap();
- builder.values().append_value(6).unwrap();
- builder.values().append_value(7).unwrap();
- builder.values().append_null().unwrap();
- builder.append(true).unwrap();
- let list_array = builder.finish();
-
- assert_eq!(DataType::Int32, list_array.value_type());
- assert_eq!(4, list_array.len());
- assert_eq!(1, list_array.null_count());
- assert_eq!(6, list_array.value_offset(2));
- assert_eq!(3, list_array.value_length());
- }
-
- #[test]
- fn test_list_array_builder_finish() {
- let values_builder = Int32Array::builder(5);
- let mut builder = ListBuilder::new(values_builder);
-
- builder.values().append_slice(&[1, 2, 3]).unwrap();
- builder.append(true).unwrap();
- builder.values().append_slice(&[4, 5, 6]).unwrap();
- builder.append(true).unwrap();
-
- let mut arr = builder.finish();
- assert_eq!(2, arr.len());
- assert_eq!(0, builder.len());
-
- builder.values().append_slice(&[7, 8, 9]).unwrap();
- builder.append(true).unwrap();
- arr = builder.finish();
- assert_eq!(1, arr.len());
- assert_eq!(0, builder.len());
- }
-
- #[test]
- fn test_fixed_size_list_array_builder_empty() {
- let values_builder = Int32Array::builder(5);
- let mut builder = FixedSizeListBuilder::new(values_builder, 3);
-
- let arr = builder.finish();
- assert_eq!(0, arr.len());
- assert_eq!(0, builder.len());
- }
-
- #[test]
- fn test_fixed_size_list_array_builder_finish() {
- let values_builder = Int32Array::builder(5);
- let mut builder = FixedSizeListBuilder::new(values_builder, 3);
-
- builder.values().append_slice(&[1, 2, 3]).unwrap();
- builder.append(true).unwrap();
- builder.values().append_slice(&[4, 5, 6]).unwrap();
- builder.append(true).unwrap();
-
- let mut arr = builder.finish();
- assert_eq!(2, arr.len());
- assert_eq!(0, builder.len());
-
- builder.values().append_slice(&[7, 8, 9]).unwrap();
- builder.append(true).unwrap();
- arr = builder.finish();
- assert_eq!(1, arr.len());
- assert_eq!(0, builder.len());
- }
-
- #[test]
- fn test_list_list_array_builder() {
- let primitive_builder = Int32Builder::new(10);
- let values_builder = ListBuilder::new(primitive_builder);
- let mut builder = ListBuilder::new(values_builder);
-
- // [[[1, 2], [3, 4]], [[5, 6, 7], null, [8]], null, [[9, 10]]]
- builder.values().values().append_value(1).unwrap();
- builder.values().values().append_value(2).unwrap();
- builder.values().append(true).unwrap();
- builder.values().values().append_value(3).unwrap();
- builder.values().values().append_value(4).unwrap();
- builder.values().append(true).unwrap();
- builder.append(true).unwrap();
-
- builder.values().values().append_value(5).unwrap();
- builder.values().values().append_value(6).unwrap();
- builder.values().values().append_value(7).unwrap();
- builder.values().append(true).unwrap();
- builder.values().append(false).unwrap();
- builder.values().values().append_value(8).unwrap();
- builder.values().append(true).unwrap();
- builder.append(true).unwrap();
-
- builder.append(false).unwrap();
-
- builder.values().values().append_value(9).unwrap();
- builder.values().values().append_value(10).unwrap();
- builder.values().append(true).unwrap();
- builder.append(true).unwrap();
-
- let list_array = builder.finish();
-
- assert_eq!(4, list_array.len());
- assert_eq!(1, list_array.null_count());
- assert_eq!(
- Buffer::from_slice_ref(&[0, 2, 5, 5, 6]),
- list_array.data().buffers()[0].clone()
- );
-
- assert_eq!(6, list_array.values().data().len());
- assert_eq!(1, list_array.values().data().null_count());
- assert_eq!(
- Buffer::from_slice_ref(&[0, 2, 4, 7, 7, 8, 10]),
- list_array.values().data().buffers()[0].clone()
- );
-
- assert_eq!(10, list_array.values().data().child_data()[0].len());
- assert_eq!(0, list_array.values().data().child_data()[0].null_count());
- assert_eq!(
- Buffer::from_slice_ref(&[1, 2, 3, 4, 5, 6, 7, 8, 9, 10]),
- list_array.values().data().child_data()[0].buffers()[0].clone()
- );
- }
-
- #[test]
- fn test_binary_array_builder() {
- let mut builder = BinaryBuilder::new(20);
-
- builder.append_byte(b'h').unwrap();
- builder.append_byte(b'e').unwrap();
- builder.append_byte(b'l').unwrap();
- builder.append_byte(b'l').unwrap();
- builder.append_byte(b'o').unwrap();
- builder.append(true).unwrap();
- builder.append(true).unwrap();
- builder.append_byte(b'w').unwrap();
- builder.append_byte(b'o').unwrap();
- builder.append_byte(b'r').unwrap();
- builder.append_byte(b'l').unwrap();
- builder.append_byte(b'd').unwrap();
- builder.append(true).unwrap();
-
- let binary_array = builder.finish();
-
- assert_eq!(3, binary_array.len());
- assert_eq!(0, binary_array.null_count());
- assert_eq!([b'h', b'e', b'l', b'l', b'o'], binary_array.value(0));
- assert_eq!([] as [u8; 0], binary_array.value(1));
- assert_eq!([b'w', b'o', b'r', b'l', b'd'], binary_array.value(2));
- assert_eq!(5, binary_array.value_offsets()[2]);
- assert_eq!(5, binary_array.value_length(2));
- }
-
- #[test]
- fn test_large_binary_array_builder() {
- let mut builder = LargeBinaryBuilder::new(20);
-
- builder.append_byte(b'h').unwrap();
- builder.append_byte(b'e').unwrap();
- builder.append_byte(b'l').unwrap();
- builder.append_byte(b'l').unwrap();
- builder.append_byte(b'o').unwrap();
- builder.append(true).unwrap();
- builder.append(true).unwrap();
- builder.append_byte(b'w').unwrap();
- builder.append_byte(b'o').unwrap();
- builder.append_byte(b'r').unwrap();
- builder.append_byte(b'l').unwrap();
- builder.append_byte(b'd').unwrap();
- builder.append(true).unwrap();
-
- let binary_array = builder.finish();
-
- assert_eq!(3, binary_array.len());
- assert_eq!(0, binary_array.null_count());
- assert_eq!([b'h', b'e', b'l', b'l', b'o'], binary_array.value(0));
- assert_eq!([] as [u8; 0], binary_array.value(1));
- assert_eq!([b'w', b'o', b'r', b'l', b'd'], binary_array.value(2));
- assert_eq!(5, binary_array.value_offsets()[2]);
- assert_eq!(5, binary_array.value_length(2));
- }
-
- #[test]
- fn test_string_array_builder() {
- let mut builder = StringBuilder::new(20);
-
- builder.append_value("hello").unwrap();
- builder.append(true).unwrap();
- builder.append_value("world").unwrap();
-
- let string_array = builder.finish();
-
- assert_eq!(3, string_array.len());
- assert_eq!(0, string_array.null_count());
- assert_eq!("hello", string_array.value(0));
- assert_eq!("", string_array.value(1));
- assert_eq!("world", string_array.value(2));
- assert_eq!(5, string_array.value_offsets()[2]);
- assert_eq!(5, string_array.value_length(2));
- }
-
- #[test]
- fn test_fixed_size_binary_builder() {
- let mut builder = FixedSizeBinaryBuilder::new(15, 5);
-
- // [b"hello", null, "arrow"]
- builder.append_value(b"hello").unwrap();
- builder.append_null().unwrap();
- builder.append_value(b"arrow").unwrap();
- let fixed_size_binary_array: FixedSizeBinaryArray = builder.finish();
-
- assert_eq!(
- &DataType::FixedSizeBinary(5),
- fixed_size_binary_array.data_type()
- );
- assert_eq!(3, fixed_size_binary_array.len());
- assert_eq!(1, fixed_size_binary_array.null_count());
- assert_eq!(10, fixed_size_binary_array.value_offset(2));
- assert_eq!(5, fixed_size_binary_array.value_length());
- }
-
- #[test]
- fn test_decimal_builder() {
- let mut builder = DecimalBuilder::new(30, 23, 6);
-
- builder.append_value(8_887_000_000).unwrap();
- builder.append_null().unwrap();
- builder.append_value(-8_887_000_000).unwrap();
- let decimal_array: DecimalArray = builder.finish();
-
- assert_eq!(&DataType::Decimal(23, 6), decimal_array.data_type());
- assert_eq!(3, decimal_array.len());
- assert_eq!(1, decimal_array.null_count());
- assert_eq!(32, decimal_array.value_offset(2));
- assert_eq!(16, decimal_array.value_length());
- }
-
- #[test]
- fn test_string_array_builder_finish() {
- let mut builder = StringBuilder::new(10);
-
- builder.append_value("hello").unwrap();
- builder.append_value("world").unwrap();
-
- let mut arr = builder.finish();
- assert_eq!(2, arr.len());
- assert_eq!(0, builder.len());
-
- builder.append_value("arrow").unwrap();
- arr = builder.finish();
- assert_eq!(1, arr.len());
- assert_eq!(0, builder.len());
- }
-
- #[test]
- fn test_string_array_builder_append_string() {
- let mut builder = StringBuilder::new(20);
-
- let var = "hello".to_owned();
- builder.append_value(&var).unwrap();
- builder.append(true).unwrap();
- builder.append_value("world").unwrap();
-
- let string_array = builder.finish();
-
- assert_eq!(3, string_array.len());
- assert_eq!(0, string_array.null_count());
- assert_eq!("hello", string_array.value(0));
- assert_eq!("", string_array.value(1));
- assert_eq!("world", string_array.value(2));
- assert_eq!(5, string_array.value_offsets()[2]);
- assert_eq!(5, string_array.value_length(2));
- }
-
- #[test]
- fn test_struct_array_builder() {
- let string_builder = StringBuilder::new(4);
- let int_builder = Int32Builder::new(4);
-
- let mut fields = Vec::new();
- let mut field_builders = Vec::new();
- fields.push(Field::new("f1", DataType::Utf8, false));
- field_builders.push(Box::new(string_builder) as Box<ArrayBuilder>);
- fields.push(Field::new("f2", DataType::Int32, false));
- field_builders.push(Box::new(int_builder) as Box<ArrayBuilder>);
-
- let mut builder = StructBuilder::new(fields, field_builders);
- assert_eq!(2, builder.num_fields());
-
- let string_builder = builder
- .field_builder::<StringBuilder>(0)
- .expect("builder at field 0 should be string builder");
- string_builder.append_value("joe").unwrap();
- string_builder.append_null().unwrap();
- string_builder.append_null().unwrap();
- string_builder.append_value("mark").unwrap();
-
- let int_builder = builder
- .field_builder::<Int32Builder>(1)
- .expect("builder at field 1 should be int builder");
- int_builder.append_value(1).unwrap();
- int_builder.append_value(2).unwrap();
- int_builder.append_null().unwrap();
- int_builder.append_value(4).unwrap();
-
- builder.append(true).unwrap();
- builder.append(true).unwrap();
- builder.append_null().unwrap();
- builder.append(true).unwrap();
-
- let arr = builder.finish();
-
- let struct_data = arr.data();
- assert_eq!(4, struct_data.len());
- assert_eq!(1, struct_data.null_count());
- assert_eq!(
- &Some(Bitmap::from(Buffer::from(&[11_u8]))),
- struct_data.null_bitmap()
- );
-
- let expected_string_data = ArrayData::builder(DataType::Utf8)
- .len(4)
- .null_bit_buffer(Buffer::from(&[9_u8]))
- .add_buffer(Buffer::from_slice_ref(&[0, 3, 3, 3, 7]))
- .add_buffer(Buffer::from_slice_ref(b"joemark"))
- .build();
-
- let expected_int_data = ArrayData::builder(DataType::Int32)
- .len(4)
- .null_bit_buffer(Buffer::from_slice_ref(&[11_u8]))
- .add_buffer(Buffer::from_slice_ref(&[1, 2, 0, 4]))
- .build();
-
- assert_eq!(&expected_string_data, arr.column(0).data());
-
- // TODO: implement equality for ArrayData
- assert_eq!(expected_int_data.len(), arr.column(1).data().len());
- assert_eq!(
- expected_int_data.null_count(),
- arr.column(1).data().null_count()
- );
- assert_eq!(
- expected_int_data.null_bitmap(),
- arr.column(1).data().null_bitmap()
- );
- let expected_value_buf = expected_int_data.buffers()[0].clone();
- let actual_value_buf = arr.column(1).data().buffers()[0].clone();
- for i in 0..expected_int_data.len() {
- if !expected_int_data.is_null(i) {
- assert_eq!(
- expected_value_buf.as_slice()[i * 4..(i + 1) * 4],
- actual_value_buf.as_slice()[i * 4..(i + 1) * 4]
- );
- }
- }
- }
-
- #[test]
- fn test_struct_array_builder_finish() {
- let int_builder = Int32Builder::new(10);
- let bool_builder = BooleanBuilder::new(10);
-
- let mut fields = Vec::new();
- let mut field_builders = Vec::new();
- fields.push(Field::new("f1", DataType::Int32, false));
- field_builders.push(Box::new(int_builder) as Box<ArrayBuilder>);
- fields.push(Field::new("f2", DataType::Boolean, false));
- field_builders.push(Box::new(bool_builder) as Box<ArrayBuilder>);
-
- let mut builder = StructBuilder::new(fields, field_builders);
- builder
- .field_builder::<Int32Builder>(0)
- .unwrap()
- .append_slice(&[0, 1, 2, 3, 4, 5, 6, 7, 8, 9])
- .unwrap();
- builder
- .field_builder::<BooleanBuilder>(1)
- .unwrap()
- .append_slice(&[
- false, true, false, true, false, true, false, true, false, true,
- ])
- .unwrap();
-
- // Append slot values - all are valid.
- for _ in 0..10 {
- assert!(builder.append(true).is_ok())
- }
-
- assert_eq!(10, builder.len());
-
- let arr = builder.finish();
-
- assert_eq!(10, arr.len());
- assert_eq!(0, builder.len());
-
- builder
- .field_builder::<Int32Builder>(0)
- .unwrap()
- .append_slice(&[1, 3, 5, 7, 9])
- .unwrap();
- builder
- .field_builder::<BooleanBuilder>(1)
- .unwrap()
- .append_slice(&[false, true, false, true, false])
- .unwrap();
-
- // Append slot values - all are valid.
- for _ in 0..5 {
- assert!(builder.append(true).is_ok())
- }
-
- assert_eq!(5, builder.len());
-
- let arr = builder.finish();
-
- assert_eq!(5, arr.len());
- assert_eq!(0, builder.len());
- }
-
- #[test]
- fn test_struct_array_builder_from_schema() {
- let mut fields = Vec::new();
- fields.push(Field::new("f1", DataType::Float32, false));
- fields.push(Field::new("f2", DataType::Utf8, false));
- let mut sub_fields = Vec::new();
- sub_fields.push(Field::new("g1", DataType::Int32, false));
- sub_fields.push(Field::new("g2", DataType::Boolean, false));
- let struct_type = DataType::Struct(sub_fields);
- fields.push(Field::new("f3", struct_type, false));
-
- let mut builder = StructBuilder::from_fields(fields, 5);
- assert_eq!(3, builder.num_fields());
- assert!(builder.field_builder::<Float32Builder>(0).is_some());
- assert!(builder.field_builder::<StringBuilder>(1).is_some());
- assert!(builder.field_builder::<StructBuilder>(2).is_some());
- }
-
- #[test]
- #[should_panic(
- expected = "Data type List(Field { name: \"item\", data_type: Int64, nullable: true, dict_id: 0, dict_is_ordered: false, metadata: None }) is not currently supported"
- )]
- fn test_struct_array_builder_from_schema_unsupported_type() {
- let mut fields = Vec::new();
- fields.push(Field::new("f1", DataType::Int16, false));
- let list_type =
- DataType::List(Box::new(Field::new("item", DataType::Int64, true)));
- fields.push(Field::new("f2", list_type, false));
-
- let _ = StructBuilder::from_fields(fields, 5);
- }
-
- #[test]
- fn test_struct_array_builder_field_builder_type_mismatch() {
- let int_builder = Int32Builder::new(10);
-
- let mut fields = Vec::new();
- let mut field_builders = Vec::new();
- fields.push(Field::new("f1", DataType::Int32, false));
- field_builders.push(Box::new(int_builder) as Box<ArrayBuilder>);
-
- let mut builder = StructBuilder::new(fields, field_builders);
- assert!(builder.field_builder::<BinaryBuilder>(0).is_none());
- }
-
- #[test]
- fn test_primitive_dictionary_builder() {
- let key_builder = PrimitiveBuilder::<UInt8Type>::new(3);
- let value_builder = PrimitiveBuilder::<UInt32Type>::new(2);
- let mut builder = PrimitiveDictionaryBuilder::new(key_builder, value_builder);
- builder.append(12345678).unwrap();
- builder.append_null().unwrap();
- builder.append(22345678).unwrap();
- let array = builder.finish();
-
- assert_eq!(
- array.keys(),
- &UInt8Array::from(vec![Some(0), None, Some(1)])
- );
-
- // Values are polymorphic and so require a downcast.
- let av = array.values();
- let ava: &UInt32Array = av.as_any().downcast_ref::<UInt32Array>().unwrap();
- let avs: &[u32] = ava.values();
-
- assert_eq!(array.is_null(0), false);
- assert_eq!(array.is_null(1), true);
- assert_eq!(array.is_null(2), false);
-
- assert_eq!(avs, &[12345678, 22345678]);
- }
-
- #[test]
- fn test_string_dictionary_builder() {
- let key_builder = PrimitiveBuilder::<Int8Type>::new(5);
- let value_builder = StringBuilder::new(2);
- let mut builder = StringDictionaryBuilder::new(key_builder, value_builder);
- builder.append("abc").unwrap();
- builder.append_null().unwrap();
- builder.append("def").unwrap();
- builder.append("def").unwrap();
- builder.append("abc").unwrap();
- let array = builder.finish();
-
- assert_eq!(
- array.keys(),
- &Int8Array::from(vec![Some(0), None, Some(1), Some(1), Some(0)])
- );
-
- // Values are polymorphic and so require a downcast.
- let av = array.values();
- let ava: &StringArray = av.as_any().downcast_ref::<StringArray>().unwrap();
-
- assert_eq!(ava.value(0), "abc");
- assert_eq!(ava.value(1), "def");
- }
-
- #[test]
- fn test_string_dictionary_builder_with_existing_dictionary() {
- let dictionary = StringArray::from(vec![None, Some("def"), Some("abc")]);
-
- let key_builder = PrimitiveBuilder::<Int8Type>::new(6);
- let mut builder =
- StringDictionaryBuilder::new_with_dictionary(key_builder, &dictionary)
- .unwrap();
- builder.append("abc").unwrap();
- builder.append_null().unwrap();
- builder.append("def").unwrap();
- builder.append("def").unwrap();
- builder.append("abc").unwrap();
- builder.append("ghi").unwrap();
- let array = builder.finish();
-
- assert_eq!(
- array.keys(),
- &Int8Array::from(vec![Some(2), None, Some(1), Some(1), Some(2), Some(3)])
- );
-
- // Values are polymorphic and so require a downcast.
- let av = array.values();
- let ava: &StringArray = av.as_any().downcast_ref::<StringArray>().unwrap();
-
- assert_eq!(ava.is_valid(0), false);
- assert_eq!(ava.value(1), "def");
- assert_eq!(ava.value(2), "abc");
- assert_eq!(ava.value(3), "ghi");
- }
-
- #[test]
- fn test_string_dictionary_builder_with_reserved_null_value() {
- let dictionary: Vec<Option<&str>> = vec![None];
- let dictionary = StringArray::from(dictionary);
-
- let key_builder = PrimitiveBuilder::<Int16Type>::new(4);
- let mut builder =
- StringDictionaryBuilder::new_with_dictionary(key_builder, &dictionary)
- .unwrap();
- builder.append("abc").unwrap();
- builder.append_null().unwrap();
- builder.append("def").unwrap();
- builder.append("abc").unwrap();
- let array = builder.finish();
-
- assert_eq!(array.is_null(1), true);
- assert_eq!(array.is_valid(1), false);
-
- let keys = array.keys();
-
- assert_eq!(keys.value(0), 1);
- assert_eq!(keys.is_null(1), true);
- // zero initialization is currently guaranteed by Buffer allocation and resizing
- assert_eq!(keys.value(1), 0);
- assert_eq!(keys.value(2), 2);
- assert_eq!(keys.value(3), 1);
- }
-
- #[test]
- #[should_panic(expected = "DictionaryKeyOverflowError")]
- fn test_primitive_dictionary_overflow() {
- let key_builder = PrimitiveBuilder::<UInt8Type>::new(257);
- let value_builder = PrimitiveBuilder::<UInt32Type>::new(257);
- let mut builder = PrimitiveDictionaryBuilder::new(key_builder, value_builder);
- // 256 unique keys.
- for i in 0..256 {
- builder.append(i + 1000).unwrap();
- }
- // Special error if the key overflows (256th entry)
- builder.append(1257).unwrap();
- }
-}
diff --git a/arrow/src/array/cast.rs b/arrow/src/array/cast.rs
deleted file mode 100644
index 0477f28..0000000
--- a/arrow/src/array/cast.rs
+++ /dev/null
@@ -1,84 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-//! Defines helper functions for force Array type downcast
-
-use crate::array::*;
-use crate::datatypes::*;
-
-/// Force downcast ArrayRef to PrimitiveArray<T>
-pub fn as_primitive_array<T>(arr: &ArrayRef) -> &PrimitiveArray<T>
-where
- T: ArrowPrimitiveType,
-{
- arr.as_any()
- .downcast_ref::<PrimitiveArray<T>>()
- .expect("Unable to downcast to primitive array")
-}
-
-/// Force downcast ArrayRef to DictionaryArray<T>
-pub fn as_dictionary_array<T>(arr: &ArrayRef) -> &DictionaryArray<T>
-where
- T: ArrowDictionaryKeyType,
-{
- arr.as_any()
- .downcast_ref::<DictionaryArray<T>>()
- .expect("Unable to downcast to dictionary array")
-}
-
-#[doc = "Force downcast ArrayRef to GenericListArray"]
-pub fn as_generic_list_array<S: OffsetSizeTrait>(arr: &ArrayRef) -> &GenericListArray<S> {
- arr.as_any()
- .downcast_ref::<GenericListArray<S>>()
- .expect("Unable to downcast to list array")
-}
-
-#[doc = "Force downcast ArrayRef to ListArray"]
-#[inline]
-pub fn as_list_array(arr: &ArrayRef) -> &ListArray {
- as_generic_list_array::<i32>(arr)
-}
-
-#[doc = "Force downcast ArrayRef to LargeListArray"]
-#[inline]
-pub fn as_large_list_array(arr: &ArrayRef) -> &LargeListArray {
- as_generic_list_array::<i64>(arr)
-}
-
-macro_rules! array_downcast_fn {
- ($name: ident, $arrty: ty, $arrty_str:expr) => {
- #[doc = "Force downcast ArrayRef to "]
- #[doc = $arrty_str]
- pub fn $name(arr: &ArrayRef) -> &$arrty {
- arr.as_any().downcast_ref::<$arrty>().expect(concat!(
- "Unable to downcast to typed array through ",
- stringify!($name)
- ))
- }
- };
-
- // use recursive macro to generate dynamic doc string for a given array type
- ($name: ident, $arrty: ty) => {
- array_downcast_fn!($name, $arrty, stringify!($arrty));
- };
-}
-
-array_downcast_fn!(as_string_array, StringArray);
-array_downcast_fn!(as_largestring_array, LargeStringArray);
-array_downcast_fn!(as_boolean_array, BooleanArray);
-array_downcast_fn!(as_null_array, NullArray);
-array_downcast_fn!(as_struct_array, StructArray);
diff --git a/arrow/src/array/data.rs b/arrow/src/array/data.rs
deleted file mode 100644
index 64767e3..0000000
--- a/arrow/src/array/data.rs
+++ /dev/null
@@ -1,685 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-//! Contains `ArrayData`, a generic representation of Arrow array data which encapsulates
-//! common attributes and operations for Arrow array.
-
-use std::mem;
-use std::sync::Arc;
-
-use crate::datatypes::{DataType, IntervalUnit};
-use crate::{bitmap::Bitmap, datatypes::ArrowNativeType};
-use crate::{
- buffer::{Buffer, MutableBuffer},
- util::bit_util,
-};
-
-use super::equal::equal;
-
-#[inline]
-pub(crate) fn count_nulls(
- null_bit_buffer: Option<&Buffer>,
- offset: usize,
- len: usize,
-) -> usize {
- if let Some(buf) = null_bit_buffer {
- len.checked_sub(buf.count_set_bits_offset(offset, len))
- .unwrap()
- } else {
- 0
- }
-}
-
-/// creates 2 [`MutableBuffer`]s with a given `capacity` (in slots).
-#[inline]
-pub(crate) fn new_buffers(data_type: &DataType, capacity: usize) -> [MutableBuffer; 2] {
- let empty_buffer = MutableBuffer::new(0);
- match data_type {
- DataType::Null => [empty_buffer, MutableBuffer::new(0)],
- DataType::Boolean => {
- let bytes = bit_util::ceil(capacity, 8);
- let buffer = MutableBuffer::new(bytes);
- [buffer, empty_buffer]
- }
- DataType::UInt8 => [
- MutableBuffer::new(capacity * mem::size_of::<u8>()),
- empty_buffer,
- ],
- DataType::UInt16 => [
- MutableBuffer::new(capacity * mem::size_of::<u16>()),
- empty_buffer,
- ],
- DataType::UInt32 => [
- MutableBuffer::new(capacity * mem::size_of::<u32>()),
- empty_buffer,
- ],
- DataType::UInt64 => [
- MutableBuffer::new(capacity * mem::size_of::<u64>()),
- empty_buffer,
- ],
- DataType::Int8 => [
- MutableBuffer::new(capacity * mem::size_of::<i8>()),
- empty_buffer,
- ],
- DataType::Int16 => [
- MutableBuffer::new(capacity * mem::size_of::<i16>()),
- empty_buffer,
- ],
- DataType::Int32 => [
- MutableBuffer::new(capacity * mem::size_of::<i32>()),
- empty_buffer,
- ],
- DataType::Int64 => [
- MutableBuffer::new(capacity * mem::size_of::<i64>()),
- empty_buffer,
- ],
- DataType::Float32 => [
- MutableBuffer::new(capacity * mem::size_of::<f32>()),
- empty_buffer,
- ],
- DataType::Float64 => [
- MutableBuffer::new(capacity * mem::size_of::<f64>()),
- empty_buffer,
- ],
- DataType::Date32 | DataType::Time32(_) => [
- MutableBuffer::new(capacity * mem::size_of::<i32>()),
- empty_buffer,
- ],
- DataType::Date64
- | DataType::Time64(_)
- | DataType::Duration(_)
- | DataType::Timestamp(_, _) => [
- MutableBuffer::new(capacity * mem::size_of::<i64>()),
- empty_buffer,
- ],
- DataType::Interval(IntervalUnit::YearMonth) => [
- MutableBuffer::new(capacity * mem::size_of::<i32>()),
- empty_buffer,
- ],
- DataType::Interval(IntervalUnit::DayTime) => [
- MutableBuffer::new(capacity * mem::size_of::<i64>()),
- empty_buffer,
- ],
- DataType::Utf8 | DataType::Binary => {
- let mut buffer = MutableBuffer::new((1 + capacity) * mem::size_of::<i32>());
- // safety: `unsafe` code assumes that this buffer is initialized with one element
- buffer.push(0i32);
- [buffer, MutableBuffer::new(capacity * mem::size_of::<u8>())]
- }
- DataType::LargeUtf8 | DataType::LargeBinary => {
- let mut buffer = MutableBuffer::new((1 + capacity) * mem::size_of::<i64>());
- // safety: `unsafe` code assumes that this buffer is initialized with one element
- buffer.push(0i64);
- [buffer, MutableBuffer::new(capacity * mem::size_of::<u8>())]
- }
- DataType::List(_) => {
- // offset buffer always starts with a zero
- let mut buffer = MutableBuffer::new((1 + capacity) * mem::size_of::<i32>());
- buffer.push(0i32);
- [buffer, empty_buffer]
- }
- DataType::LargeList(_) => {
- // offset buffer always starts with a zero
- let mut buffer = MutableBuffer::new((1 + capacity) * mem::size_of::<i64>());
- buffer.push(0i64);
- [buffer, empty_buffer]
- }
- DataType::FixedSizeBinary(size) => {
- [MutableBuffer::new(capacity * *size as usize), empty_buffer]
- }
- DataType::Dictionary(child_data_type, _) => match child_data_type.as_ref() {
- DataType::UInt8 => [
- MutableBuffer::new(capacity * mem::size_of::<u8>()),
- empty_buffer,
- ],
- DataType::UInt16 => [
- MutableBuffer::new(capacity * mem::size_of::<u16>()),
- empty_buffer,
- ],
- DataType::UInt32 => [
- MutableBuffer::new(capacity * mem::size_of::<u32>()),
- empty_buffer,
- ],
- DataType::UInt64 => [
- MutableBuffer::new(capacity * mem::size_of::<u64>()),
- empty_buffer,
- ],
- DataType::Int8 => [
- MutableBuffer::new(capacity * mem::size_of::<i8>()),
- empty_buffer,
- ],
- DataType::Int16 => [
- MutableBuffer::new(capacity * mem::size_of::<i16>()),
- empty_buffer,
- ],
- DataType::Int32 => [
- MutableBuffer::new(capacity * mem::size_of::<i32>()),
- empty_buffer,
- ],
- DataType::Int64 => [
- MutableBuffer::new(capacity * mem::size_of::<i64>()),
- empty_buffer,
- ],
- _ => unreachable!(),
- },
- DataType::Float16 => unreachable!(),
- DataType::FixedSizeList(_, _) | DataType::Struct(_) => {
- [empty_buffer, MutableBuffer::new(0)]
- }
- DataType::Decimal(_, _) => [
- MutableBuffer::new(capacity * mem::size_of::<u8>()),
- empty_buffer,
- ],
- DataType::Union(_) => unimplemented!(),
- }
-}
-
-/// Maps 2 [`MutableBuffer`]s into a vector of [Buffer]s whose size depends on `data_type`.
-#[inline]
-pub(crate) fn into_buffers(
- data_type: &DataType,
- buffer1: MutableBuffer,
- buffer2: MutableBuffer,
-) -> Vec<Buffer> {
- match data_type {
- DataType::Null | DataType::Struct(_) => vec![],
- DataType::Utf8
- | DataType::Binary
- | DataType::LargeUtf8
- | DataType::LargeBinary => vec![buffer1.into(), buffer2.into()],
- _ => vec![buffer1.into()],
- }
-}
-
-/// An generic representation of Arrow array data which encapsulates common attributes and
-/// operations for Arrow array. Specific operations for different arrays types (e.g.,
-/// primitive, list, struct) are implemented in `Array`.
-#[derive(Debug, Clone)]
-pub struct ArrayData {
- /// The data type for this array data
- data_type: DataType,
-
- /// The number of elements in this array data
- len: usize,
-
- /// The number of null elements in this array data
- null_count: usize,
-
- /// The offset into this array data, in number of items
- offset: usize,
-
- /// The buffers for this array data. Note that depending on the array types, this
- /// could hold different kinds of buffers (e.g., value buffer, value offset buffer)
- /// at different positions.
- buffers: Vec<Buffer>,
-
- /// The child(ren) of this array. Only non-empty for nested types, currently
- /// `ListArray` and `StructArray`.
- child_data: Vec<ArrayData>,
-
- /// The null bitmap. A `None` value for this indicates all values are non-null in
- /// this array.
- null_bitmap: Option<Bitmap>,
-}
-
-pub type ArrayDataRef = Arc<ArrayData>;
-
-impl ArrayData {
- pub fn new(
- data_type: DataType,
- len: usize,
- null_count: Option<usize>,
- null_bit_buffer: Option<Buffer>,
- offset: usize,
- buffers: Vec<Buffer>,
- child_data: Vec<ArrayData>,
- ) -> Self {
- let null_count = match null_count {
- None => count_nulls(null_bit_buffer.as_ref(), offset, len),
- Some(null_count) => null_count,
- };
- let null_bitmap = null_bit_buffer.map(Bitmap::from);
- Self {
- data_type,
- len,
- null_count,
- offset,
- buffers,
- child_data,
- null_bitmap,
- }
- }
-
- /// Returns a builder to construct a `ArrayData` instance.
- #[inline]
- pub const fn builder(data_type: DataType) -> ArrayDataBuilder {
- ArrayDataBuilder::new(data_type)
- }
-
- /// Returns a reference to the data type of this array data
- #[inline]
- pub const fn data_type(&self) -> &DataType {
- &self.data_type
- }
-
- /// Returns a slice of buffers for this array data
- pub fn buffers(&self) -> &[Buffer] {
- &self.buffers[..]
- }
-
- /// Returns a slice of children data arrays
- pub fn child_data(&self) -> &[ArrayData] {
- &self.child_data[..]
- }
-
- /// Returns whether the element at index `i` is null
- pub fn is_null(&self, i: usize) -> bool {
- if let Some(ref b) = self.null_bitmap {
- return !b.is_set(self.offset + i);
- }
- false
- }
-
- /// Returns a reference to the null bitmap of this array data
- #[inline]
- pub const fn null_bitmap(&self) -> &Option<Bitmap> {
- &self.null_bitmap
- }
-
- /// Returns a reference to the null buffer of this array data.
- pub fn null_buffer(&self) -> Option<&Buffer> {
- self.null_bitmap().as_ref().map(|b| b.buffer_ref())
- }
-
- /// Returns whether the element at index `i` is not null
- pub fn is_valid(&self, i: usize) -> bool {
- if let Some(ref b) = self.null_bitmap {
- return b.is_set(self.offset + i);
- }
- true
- }
-
- /// Returns the length (i.e., number of elements) of this array
- #[inline]
- pub const fn len(&self) -> usize {
- self.len
- }
-
- // Returns whether array data is empty
- #[inline]
- pub const fn is_empty(&self) -> bool {
- self.len == 0
- }
-
- /// Returns the offset of this array
- #[inline]
- pub const fn offset(&self) -> usize {
- self.offset
- }
-
- /// Returns the total number of nulls in this array
- #[inline]
- pub const fn null_count(&self) -> usize {
- self.null_count
- }
-
- /// Returns the total number of bytes of memory occupied by the buffers owned by this [ArrayData].
- pub fn get_buffer_memory_size(&self) -> usize {
- let mut size = 0;
- for buffer in &self.buffers {
- size += buffer.capacity();
- }
- if let Some(bitmap) = &self.null_bitmap {
- size += bitmap.get_buffer_memory_size()
- }
- for child in &self.child_data {
- size += child.get_buffer_memory_size();
- }
- size
- }
-
- /// Returns the total number of bytes of memory occupied physically by this [ArrayData].
- pub fn get_array_memory_size(&self) -> usize {
- let mut size = 0;
- // Calculate size of the fields that don't have [get_array_memory_size] method internally.
- size += mem::size_of_val(self)
- - mem::size_of_val(&self.buffers)
- - mem::size_of_val(&self.null_bitmap)
- - mem::size_of_val(&self.child_data);
-
- // Calculate rest of the fields top down which contain actual data
- for buffer in &self.buffers {
- size += mem::size_of_val(&buffer);
- size += buffer.capacity();
- }
- if let Some(bitmap) = &self.null_bitmap {
- size += bitmap.get_array_memory_size()
- }
- for child in &self.child_data {
- size += child.get_array_memory_size();
- }
-
- size
- }
-
- /// Creates a zero-copy slice of itself. This creates a new [ArrayData]
- /// with a different offset, len and a shifted null bitmap.
- ///
- /// # Panics
- ///
- /// Panics if `offset + length > self.len()`.
- pub fn slice(&self, offset: usize, length: usize) -> ArrayData {
- assert!((offset + length) <= self.len());
-
- let mut new_data = self.clone();
-
- new_data.len = length;
- new_data.offset = offset + self.offset;
-
- new_data.null_count =
- count_nulls(new_data.null_buffer(), new_data.offset, new_data.len);
-
- new_data
- }
-
- /// Returns the `buffer` as a slice of type `T` starting at self.offset
- /// # Panics
- /// This function panics if:
- /// * the buffer is not byte-aligned with type T, or
- /// * the datatype is `Boolean` (it corresponds to a bit-packed buffer where the offset is not applicable)
- #[inline]
- pub(crate) fn buffer<T: ArrowNativeType>(&self, buffer: usize) -> &[T] {
- let values = unsafe { self.buffers[buffer].as_slice().align_to::<T>() };
- if !values.0.is_empty() || !values.2.is_empty() {
- panic!("The buffer is not byte-aligned with its interpretation")
- };
- assert_ne!(self.data_type, DataType::Boolean);
- &values.1[self.offset..]
- }
-
- /// Returns a new empty [ArrayData] valid for `data_type`.
- pub fn new_empty(data_type: &DataType) -> Self {
- let buffers = new_buffers(data_type, 0);
- let [buffer1, buffer2] = buffers;
- let buffers = into_buffers(data_type, buffer1, buffer2);
-
- let child_data = match data_type {
- DataType::Null
- | DataType::Boolean
- | DataType::UInt8
- | DataType::UInt16
- | DataType::UInt32
- | DataType::UInt64
- | DataType::Int8
- | DataType::Int16
- | DataType::Int32
- | DataType::Int64
- | DataType::Float32
- | DataType::Float64
- | DataType::Date32
- | DataType::Date64
- | DataType::Time32(_)
- | DataType::Time64(_)
- | DataType::Duration(_)
- | DataType::Timestamp(_, _)
- | DataType::Utf8
- | DataType::Binary
- | DataType::LargeUtf8
- | DataType::LargeBinary
- | DataType::Interval(_)
- | DataType::FixedSizeBinary(_)
- | DataType::Decimal(_, _) => vec![],
- DataType::List(field) => {
- vec![Self::new_empty(field.data_type())]
- }
- DataType::FixedSizeList(field, _) => {
- vec![Self::new_empty(field.data_type())]
- }
- DataType::LargeList(field) => {
- vec![Self::new_empty(field.data_type())]
- }
- DataType::Struct(fields) => fields
- .iter()
- .map(|field| Self::new_empty(field.data_type()))
- .collect(),
- DataType::Union(_) => unimplemented!(),
- DataType::Dictionary(_, data_type) => {
- vec![Self::new_empty(data_type)]
- }
- DataType::Float16 => unreachable!(),
- };
-
- Self::new(data_type.clone(), 0, Some(0), None, 0, buffers, child_data)
- }
-}
-
-impl PartialEq for ArrayData {
- fn eq(&self, other: &Self) -> bool {
- equal(self, other)
- }
-}
-
-/// Builder for `ArrayData` type
-#[derive(Debug)]
-pub struct ArrayDataBuilder {
- data_type: DataType,
- len: usize,
- null_count: Option<usize>,
- null_bit_buffer: Option<Buffer>,
- offset: usize,
- buffers: Vec<Buffer>,
- child_data: Vec<ArrayData>,
-}
-
-impl ArrayDataBuilder {
- #[inline]
- pub const fn new(data_type: DataType) -> Self {
- Self {
- data_type,
- len: 0,
- null_count: None,
- null_bit_buffer: None,
- offset: 0,
- buffers: vec![],
- child_data: vec![],
- }
- }
-
- #[inline]
- #[allow(clippy::len_without_is_empty)]
- pub const fn len(mut self, n: usize) -> Self {
- self.len = n;
- self
- }
-
- pub fn null_count(mut self, null_count: usize) -> Self {
- self.null_count = Some(null_count);
- self
- }
-
- pub fn null_bit_buffer(mut self, buf: Buffer) -> Self {
- self.null_bit_buffer = Some(buf);
- self
- }
-
- #[inline]
- pub const fn offset(mut self, n: usize) -> Self {
- self.offset = n;
- self
- }
-
- pub fn buffers(mut self, v: Vec<Buffer>) -> Self {
- self.buffers = v;
- self
- }
-
- pub fn add_buffer(mut self, b: Buffer) -> Self {
- self.buffers.push(b);
- self
- }
-
- pub fn child_data(mut self, v: Vec<ArrayData>) -> Self {
- self.child_data = v;
- self
- }
-
- pub fn add_child_data(mut self, r: ArrayData) -> Self {
- self.child_data.push(r);
- self
- }
-
- pub fn build(self) -> ArrayData {
- ArrayData::new(
- self.data_type,
- self.len,
- self.null_count,
- self.null_bit_buffer,
- self.offset,
- self.buffers,
- self.child_data,
- )
- }
-}
-
-#[cfg(test)]
-mod tests {
- use super::*;
-
- use crate::buffer::Buffer;
- use crate::util::bit_util;
-
- #[test]
- fn test_new() {
- let arr_data =
- ArrayData::new(DataType::Boolean, 10, Some(1), None, 2, vec![], vec![]);
- assert_eq!(10, arr_data.len());
- assert_eq!(1, arr_data.null_count());
- assert_eq!(2, arr_data.offset());
- assert_eq!(0, arr_data.buffers().len());
- assert_eq!(0, arr_data.child_data().len());
- }
-
- #[test]
- fn test_builder() {
- let child_arr_data = ArrayData::new(
- DataType::Int32,
- 5,
- Some(0),
- None,
- 0,
- vec![Buffer::from_slice_ref(&[1i32, 2, 3, 4, 5])],
- vec![],
- );
- let v = vec![0, 1, 2, 3];
- let b1 = Buffer::from(&v[..]);
- let arr_data = ArrayData::builder(DataType::Int32)
- .len(20)
- .offset(5)
- .add_buffer(b1)
- .null_bit_buffer(Buffer::from(vec![
- 0b01011111, 0b10110101, 0b01100011, 0b00011110,
- ]))
- .add_child_data(child_arr_data.clone())
- .build();
-
- assert_eq!(20, arr_data.len());
- assert_eq!(10, arr_data.null_count());
- assert_eq!(5, arr_data.offset());
- assert_eq!(1, arr_data.buffers().len());
- assert_eq!(&[0, 1, 2, 3], arr_data.buffers()[0].as_slice());
- assert_eq!(1, arr_data.child_data().len());
- assert_eq!(child_arr_data, arr_data.child_data()[0]);
- }
-
- #[test]
- fn test_null_count() {
- let mut bit_v: [u8; 2] = [0; 2];
- bit_util::set_bit(&mut bit_v, 0);
- bit_util::set_bit(&mut bit_v, 3);
- bit_util::set_bit(&mut bit_v, 10);
- let arr_data = ArrayData::builder(DataType::Int32)
- .len(16)
- .null_bit_buffer(Buffer::from(bit_v))
- .build();
- assert_eq!(13, arr_data.null_count());
-
- // Test with offset
- let mut bit_v: [u8; 2] = [0; 2];
- bit_util::set_bit(&mut bit_v, 0);
- bit_util::set_bit(&mut bit_v, 3);
- bit_util::set_bit(&mut bit_v, 10);
- let arr_data = ArrayData::builder(DataType::Int32)
- .len(12)
- .offset(2)
- .null_bit_buffer(Buffer::from(bit_v))
- .build();
- assert_eq!(10, arr_data.null_count());
- }
-
- #[test]
- fn test_null_buffer_ref() {
- let mut bit_v: [u8; 2] = [0; 2];
- bit_util::set_bit(&mut bit_v, 0);
- bit_util::set_bit(&mut bit_v, 3);
- bit_util::set_bit(&mut bit_v, 10);
- let arr_data = ArrayData::builder(DataType::Int32)
- .len(16)
- .null_bit_buffer(Buffer::from(bit_v))
- .build();
- assert!(arr_data.null_buffer().is_some());
- assert_eq!(&bit_v, arr_data.null_buffer().unwrap().as_slice());
- }
-
- #[test]
- fn test_slice() {
- let mut bit_v: [u8; 2] = [0; 2];
- bit_util::set_bit(&mut bit_v, 0);
- bit_util::set_bit(&mut bit_v, 3);
- bit_util::set_bit(&mut bit_v, 10);
- let data = ArrayData::builder(DataType::Int32)
- .len(16)
- .null_bit_buffer(Buffer::from(bit_v))
- .build();
- let new_data = data.slice(1, 15);
- assert_eq!(data.len() - 1, new_data.len());
- assert_eq!(1, new_data.offset());
- assert_eq!(data.null_count(), new_data.null_count());
-
- // slice of a slice (removes one null)
- let new_data = new_data.slice(1, 14);
- assert_eq!(data.len() - 2, new_data.len());
- assert_eq!(2, new_data.offset());
- assert_eq!(data.null_count() - 1, new_data.null_count());
- }
-
- #[test]
- fn test_equality() {
- let int_data = ArrayData::builder(DataType::Int32).build();
- let float_data = ArrayData::builder(DataType::Float32).build();
- assert_ne!(int_data, float_data);
- }
-
- #[test]
- fn test_count_nulls() {
- let null_buffer = Some(Buffer::from(vec![0b00010110, 0b10011111]));
- let count = count_nulls(null_buffer.as_ref(), 0, 16);
- assert_eq!(count, 7);
-
- let count = count_nulls(null_buffer.as_ref(), 4, 8);
- assert_eq!(count, 3);
- }
-}
diff --git a/arrow/src/array/equal/boolean.rs b/arrow/src/array/equal/boolean.rs
deleted file mode 100644
index 35c9786..0000000
--- a/arrow/src/array/equal/boolean.rs
+++ /dev/null
@@ -1,93 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-use crate::array::{data::count_nulls, ArrayData};
-use crate::buffer::Buffer;
-use crate::util::bit_util::get_bit;
-
-use super::utils::{equal_bits, equal_len};
-
-pub(super) fn boolean_equal(
- lhs: &ArrayData,
- rhs: &ArrayData,
- lhs_nulls: Option<&Buffer>,
- rhs_nulls: Option<&Buffer>,
- mut lhs_start: usize,
- mut rhs_start: usize,
- mut len: usize,
-) -> bool {
- let lhs_values = lhs.buffers()[0].as_slice();
- let rhs_values = rhs.buffers()[0].as_slice();
-
- let lhs_null_count = count_nulls(lhs_nulls, lhs_start, len);
- let rhs_null_count = count_nulls(rhs_nulls, rhs_start, len);
-
- if lhs_null_count == 0 && rhs_null_count == 0 {
- // Optimize performance for starting offset at u8 boundary.
- if lhs_start % 8 == 0 && rhs_start % 8 == 0 {
- let quot = len / 8;
- if quot > 0
- && !equal_len(
- lhs_values,
- rhs_values,
- lhs_start / 8 + lhs.offset(),
- rhs_start / 8 + rhs.offset(),
- quot,
- )
- {
- return false;
- }
-
- // Calculate for suffix bits.
- let rem = len % 8;
- if rem == 0 {
- return true;
- } else {
- let aligned_bits = len - rem;
- lhs_start += aligned_bits;
- rhs_start += aligned_bits;
- len = rem
- }
- }
-
- equal_bits(
- lhs_values,
- rhs_values,
- lhs_start + lhs.offset(),
- rhs_start + rhs.offset(),
- len,
- )
- } else {
- // get a ref of the null buffer bytes, to use in testing for nullness
- let lhs_null_bytes = lhs_nulls.as_ref().unwrap().as_slice();
- let rhs_null_bytes = rhs_nulls.as_ref().unwrap().as_slice();
-
- let lhs_start = lhs.offset() + lhs_start;
- let rhs_start = rhs.offset() + rhs_start;
-
- (0..len).all(|i| {
- let lhs_pos = lhs_start + i;
- let rhs_pos = rhs_start + i;
- let lhs_is_null = !get_bit(lhs_null_bytes, lhs_pos);
- let rhs_is_null = !get_bit(rhs_null_bytes, rhs_pos);
-
- lhs_is_null
- || (lhs_is_null == rhs_is_null)
- && equal_bits(lhs_values, rhs_values, lhs_pos, rhs_pos, 1)
- })
- }
-}
diff --git a/arrow/src/array/equal/decimal.rs b/arrow/src/array/equal/decimal.rs
deleted file mode 100644
index 1ee6ec9..0000000
--- a/arrow/src/array/equal/decimal.rs
+++ /dev/null
@@ -1,76 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-use crate::array::{data::count_nulls, ArrayData};
-use crate::buffer::Buffer;
-use crate::datatypes::DataType;
-use crate::util::bit_util::get_bit;
-
-use super::utils::equal_len;
-
-pub(super) fn decimal_equal(
- lhs: &ArrayData,
- rhs: &ArrayData,
- lhs_nulls: Option<&Buffer>,
- rhs_nulls: Option<&Buffer>,
- lhs_start: usize,
- rhs_start: usize,
- len: usize,
-) -> bool {
- let size = match lhs.data_type() {
- DataType::Decimal(_, _) => 16,
- _ => unreachable!(),
- };
-
- let lhs_values = &lhs.buffers()[0].as_slice()[lhs.offset() * size..];
- let rhs_values = &rhs.buffers()[0].as_slice()[rhs.offset() * size..];
-
- let lhs_null_count = count_nulls(lhs_nulls, lhs_start, len);
- let rhs_null_count = count_nulls(rhs_nulls, rhs_start, len);
-
- if lhs_null_count == 0 && rhs_null_count == 0 {
- equal_len(
- lhs_values,
- rhs_values,
- size * lhs_start,
- size * rhs_start,
- size * len,
- )
- } else {
- // get a ref of the null buffer bytes, to use in testing for nullness
- let lhs_null_bytes = lhs_nulls.as_ref().unwrap().as_slice();
- let rhs_null_bytes = rhs_nulls.as_ref().unwrap().as_slice();
- // with nulls, we need to compare item by item whenever it is not null
- (0..len).all(|i| {
- let lhs_pos = lhs_start + i;
- let rhs_pos = rhs_start + i;
-
- let lhs_is_null = !get_bit(lhs_null_bytes, lhs_pos + lhs.offset());
- let rhs_is_null = !get_bit(rhs_null_bytes, rhs_pos + rhs.offset());
-
- lhs_is_null
- || (lhs_is_null == rhs_is_null)
- && equal_len(
- lhs_values,
- rhs_values,
- lhs_pos * size,
- rhs_pos * size,
- size, // 1 * size since we are comparing a single entry
- )
- })
- }
-}
diff --git a/arrow/src/array/equal/dictionary.rs b/arrow/src/array/equal/dictionary.rs
deleted file mode 100644
index 22add24..0000000
--- a/arrow/src/array/equal/dictionary.rs
+++ /dev/null
@@ -1,82 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-use crate::array::{data::count_nulls, ArrayData};
-use crate::buffer::Buffer;
-use crate::datatypes::ArrowNativeType;
-use crate::util::bit_util::get_bit;
-
-use super::equal_range;
-
-pub(super) fn dictionary_equal<T: ArrowNativeType>(
- lhs: &ArrayData,
- rhs: &ArrayData,
- lhs_nulls: Option<&Buffer>,
- rhs_nulls: Option<&Buffer>,
- lhs_start: usize,
- rhs_start: usize,
- len: usize,
-) -> bool {
- let lhs_keys = lhs.buffer::<T>(0);
- let rhs_keys = rhs.buffer::<T>(0);
-
- let lhs_values = &lhs.child_data()[0];
- let rhs_values = &rhs.child_data()[0];
-
- let lhs_null_count = count_nulls(lhs_nulls, lhs_start, len);
- let rhs_null_count = count_nulls(rhs_nulls, rhs_start, len);
-
- if lhs_null_count == 0 && rhs_null_count == 0 {
- (0..len).all(|i| {
- let lhs_pos = lhs_start + i;
- let rhs_pos = rhs_start + i;
-
- equal_range(
- lhs_values,
- rhs_values,
- lhs_values.null_buffer(),
- rhs_values.null_buffer(),
- lhs_keys[lhs_pos].to_usize().unwrap(),
- rhs_keys[rhs_pos].to_usize().unwrap(),
- 1,
- )
- })
- } else {
- // get a ref of the null buffer bytes, to use in testing for nullness
- let lhs_null_bytes = lhs_nulls.as_ref().unwrap().as_slice();
- let rhs_null_bytes = rhs_nulls.as_ref().unwrap().as_slice();
- (0..len).all(|i| {
- let lhs_pos = lhs_start + i;
- let rhs_pos = rhs_start + i;
-
- let lhs_is_null = !get_bit(lhs_null_bytes, lhs_pos + lhs.offset());
- let rhs_is_null = !get_bit(rhs_null_bytes, rhs_pos + rhs.offset());
-
- lhs_is_null
- || (lhs_is_null == rhs_is_null)
- && equal_range(
- lhs_values,
- rhs_values,
- lhs_values.null_buffer(),
- rhs_values.null_buffer(),
- lhs_keys[lhs_pos].to_usize().unwrap(),
- rhs_keys[rhs_pos].to_usize().unwrap(),
- 1,
- )
- })
- }
-}
diff --git a/arrow/src/array/equal/fixed_binary.rs b/arrow/src/array/equal/fixed_binary.rs
deleted file mode 100644
index 5f8f932..0000000
--- a/arrow/src/array/equal/fixed_binary.rs
+++ /dev/null
@@ -1,76 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-use crate::array::{data::count_nulls, ArrayData};
-use crate::buffer::Buffer;
-use crate::datatypes::DataType;
-use crate::util::bit_util::get_bit;
-
-use super::utils::equal_len;
-
-pub(super) fn fixed_binary_equal(
- lhs: &ArrayData,
- rhs: &ArrayData,
- lhs_nulls: Option<&Buffer>,
- rhs_nulls: Option<&Buffer>,
- lhs_start: usize,
- rhs_start: usize,
- len: usize,
-) -> bool {
- let size = match lhs.data_type() {
- DataType::FixedSizeBinary(i) => *i as usize,
- _ => unreachable!(),
- };
-
- let lhs_values = &lhs.buffers()[0].as_slice()[lhs.offset() * size..];
- let rhs_values = &rhs.buffers()[0].as_slice()[rhs.offset() * size..];
-
- let lhs_null_count = count_nulls(lhs_nulls, lhs_start, len);
- let rhs_null_count = count_nulls(rhs_nulls, rhs_start, len);
-
- if lhs_null_count == 0 && rhs_null_count == 0 {
- equal_len(
- lhs_values,
- rhs_values,
- size * lhs_start,
- size * rhs_start,
- size * len,
- )
- } else {
- // get a ref of the null buffer bytes, to use in testing for nullness
- let lhs_null_bytes = lhs_nulls.as_ref().unwrap().as_slice();
- let rhs_null_bytes = rhs_nulls.as_ref().unwrap().as_slice();
- // with nulls, we need to compare item by item whenever it is not null
- (0..len).all(|i| {
- let lhs_pos = lhs_start + i;
- let rhs_pos = rhs_start + i;
-
- let lhs_is_null = !get_bit(lhs_null_bytes, lhs_pos + lhs.offset());
- let rhs_is_null = !get_bit(rhs_null_bytes, rhs_pos + rhs.offset());
-
- lhs_is_null
- || (lhs_is_null == rhs_is_null)
- && equal_len(
- lhs_values,
- rhs_values,
- lhs_pos * size,
- rhs_pos * size,
- size, // 1 * size since we are comparing a single entry
- )
- })
- }
-}
diff --git a/arrow/src/array/equal/fixed_list.rs b/arrow/src/array/equal/fixed_list.rs
deleted file mode 100644
index e708a06..0000000
--- a/arrow/src/array/equal/fixed_list.rs
+++ /dev/null
@@ -1,80 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-use crate::array::{data::count_nulls, ArrayData};
-use crate::buffer::Buffer;
-use crate::datatypes::DataType;
-use crate::util::bit_util::get_bit;
-
-use super::equal_range;
-
-pub(super) fn fixed_list_equal(
- lhs: &ArrayData,
- rhs: &ArrayData,
- lhs_nulls: Option<&Buffer>,
- rhs_nulls: Option<&Buffer>,
- lhs_start: usize,
- rhs_start: usize,
- len: usize,
-) -> bool {
- let size = match lhs.data_type() {
- DataType::FixedSizeList(_, i) => *i as usize,
- _ => unreachable!(),
- };
-
- let lhs_values = &lhs.child_data()[0];
- let rhs_values = &rhs.child_data()[0];
-
- let lhs_null_count = count_nulls(lhs_nulls, lhs_start, len);
- let rhs_null_count = count_nulls(rhs_nulls, rhs_start, len);
-
- if lhs_null_count == 0 && rhs_null_count == 0 {
- equal_range(
- lhs_values,
- rhs_values,
- lhs_values.null_buffer(),
- rhs_values.null_buffer(),
- size * lhs_start,
- size * rhs_start,
- size * len,
- )
- } else {
- // get a ref of the null buffer bytes, to use in testing for nullness
- let lhs_null_bytes = lhs_nulls.as_ref().unwrap().as_slice();
- let rhs_null_bytes = rhs_nulls.as_ref().unwrap().as_slice();
- // with nulls, we need to compare item by item whenever it is not null
- (0..len).all(|i| {
- let lhs_pos = lhs_start + i;
- let rhs_pos = rhs_start + i;
-
- let lhs_is_null = !get_bit(lhs_null_bytes, lhs_pos + lhs.offset());
- let rhs_is_null = !get_bit(rhs_null_bytes, rhs_pos + rhs.offset());
-
- lhs_is_null
- || (lhs_is_null == rhs_is_null)
- && equal_range(
- lhs_values,
- rhs_values,
- lhs_values.null_buffer(),
- rhs_values.null_buffer(),
- lhs_pos * size,
- rhs_pos * size,
- size, // 1 * size since we are comparing a single entry
- )
- })
- }
-}
diff --git a/arrow/src/array/equal/list.rs b/arrow/src/array/equal/list.rs
deleted file mode 100644
index 331cdc7..0000000
--- a/arrow/src/array/equal/list.rs
+++ /dev/null
@@ -1,172 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-use crate::{
- array::ArrayData,
- array::{data::count_nulls, OffsetSizeTrait},
- buffer::Buffer,
- util::bit_util::get_bit,
-};
-
-use super::{equal_range, utils::child_logical_null_buffer};
-
-fn lengths_equal<T: OffsetSizeTrait>(lhs: &[T], rhs: &[T]) -> bool {
- // invariant from `base_equal`
- debug_assert_eq!(lhs.len(), rhs.len());
-
- if lhs.is_empty() {
- return true;
- }
-
- if lhs[0] == T::zero() && rhs[0] == T::zero() {
- return lhs == rhs;
- };
-
- // The expensive case, e.g.
- // [0, 2, 4, 6, 9] == [4, 6, 8, 10, 13]
- lhs.windows(2)
- .zip(rhs.windows(2))
- .all(|(lhs_offsets, rhs_offsets)| {
- // length of left == length of right
- (lhs_offsets[1] - lhs_offsets[0]) == (rhs_offsets[1] - rhs_offsets[0])
- })
-}
-
-#[allow(clippy::too_many_arguments)]
-#[inline]
-fn offset_value_equal<T: OffsetSizeTrait>(
- lhs_values: &ArrayData,
- rhs_values: &ArrayData,
- lhs_nulls: Option<&Buffer>,
- rhs_nulls: Option<&Buffer>,
- lhs_offsets: &[T],
- rhs_offsets: &[T],
- lhs_pos: usize,
- rhs_pos: usize,
- len: usize,
-) -> bool {
- let lhs_start = lhs_offsets[lhs_pos].to_usize().unwrap();
- let rhs_start = rhs_offsets[rhs_pos].to_usize().unwrap();
- let lhs_len = lhs_offsets[lhs_pos + len] - lhs_offsets[lhs_pos];
- let rhs_len = rhs_offsets[rhs_pos + len] - rhs_offsets[rhs_pos];
-
- lhs_len == rhs_len
- && equal_range(
- lhs_values,
- rhs_values,
- lhs_nulls,
- rhs_nulls,
- lhs_start,
- rhs_start,
- lhs_len.to_usize().unwrap(),
- )
-}
-
-pub(super) fn list_equal<T: OffsetSizeTrait>(
- lhs: &ArrayData,
- rhs: &ArrayData,
- lhs_nulls: Option<&Buffer>,
- rhs_nulls: Option<&Buffer>,
- lhs_start: usize,
- rhs_start: usize,
- len: usize,
-) -> bool {
- let lhs_offsets = lhs.buffer::<T>(0);
- let rhs_offsets = rhs.buffer::<T>(0);
-
- // There is an edge-case where a n-length list that has 0 children, results in panics.
- // For example; an array with offsets [0, 0, 0, 0, 0] has 4 slots, but will have
- // no valid children.
- // Under logical equality, the child null bitmap will be an empty buffer, as there are
- // no child values. This causes panics when trying to count set bits.
- //
- // We caught this by chance from an accidental test-case, but due to the nature of this
- // crash only occuring on list equality checks, we are adding a check here, instead of
- // on the buffer/bitmap utilities, as a length check would incur a penalty for almost all
- // other use-cases.
- //
- // The solution is to check the number of child values from offsets, and return `true` if
- // they = 0. Empty arrays are equal, so this is correct.
- //
- // It's unlikely that one would create a n-length list array with no values, where n > 0,
- // however, one is more likely to slice into a list array and get a region that has 0
- // child values.
- // The test that triggered this behaviour had [4, 4] as a slice of 1 value slot.
- let lhs_child_length = lhs_offsets.get(len).unwrap().to_usize().unwrap()
- - lhs_offsets.first().unwrap().to_usize().unwrap();
- let rhs_child_length = rhs_offsets.get(len).unwrap().to_usize().unwrap()
- - rhs_offsets.first().unwrap().to_usize().unwrap();
-
- if lhs_child_length == 0 && lhs_child_length == rhs_child_length {
- return true;
- }
-
- let lhs_values = &lhs.child_data()[0];
- let rhs_values = &rhs.child_data()[0];
-
- let lhs_null_count = count_nulls(lhs_nulls, lhs_start, len);
- let rhs_null_count = count_nulls(rhs_nulls, rhs_start, len);
-
- // compute the child logical bitmap
- let child_lhs_nulls =
- child_logical_null_buffer(lhs, lhs_nulls, lhs.child_data().get(0).unwrap());
- let child_rhs_nulls =
- child_logical_null_buffer(rhs, rhs_nulls, rhs.child_data().get(0).unwrap());
-
- if lhs_null_count == 0 && rhs_null_count == 0 {
- lengths_equal(
- &lhs_offsets[lhs_start..lhs_start + len],
- &rhs_offsets[rhs_start..rhs_start + len],
- ) && equal_range(
- lhs_values,
- rhs_values,
- child_lhs_nulls.as_ref(),
- child_rhs_nulls.as_ref(),
- lhs_offsets[lhs_start].to_usize().unwrap(),
- rhs_offsets[rhs_start].to_usize().unwrap(),
- (lhs_offsets[len] - lhs_offsets[lhs_start])
- .to_usize()
- .unwrap(),
- )
- } else {
- // get a ref of the parent null buffer bytes, to use in testing for nullness
- let lhs_null_bytes = lhs_nulls.unwrap().as_slice();
- let rhs_null_bytes = rhs_nulls.unwrap().as_slice();
- // with nulls, we need to compare item by item whenever it is not null
- (0..len).all(|i| {
- let lhs_pos = lhs_start + i;
- let rhs_pos = rhs_start + i;
-
- let lhs_is_null = !get_bit(lhs_null_bytes, lhs_pos + lhs.offset());
- let rhs_is_null = !get_bit(rhs_null_bytes, rhs_pos + rhs.offset());
-
- lhs_is_null
- || (lhs_is_null == rhs_is_null)
- && offset_value_equal::<T>(
- lhs_values,
- rhs_values,
- child_lhs_nulls.as_ref(),
- child_rhs_nulls.as_ref(),
- lhs_offsets,
- rhs_offsets,
- lhs_pos,
- rhs_pos,
- 1,
- )
- })
- }
-}
diff --git a/arrow/src/array/equal/mod.rs b/arrow/src/array/equal/mod.rs
deleted file mode 100644
index 0924fc1..0000000
--- a/arrow/src/array/equal/mod.rs
+++ /dev/null
@@ -1,1277 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-//! Module containing functionality to compute array equality.
-//! This module uses [ArrayData] and does not
-//! depend on dynamic casting of `Array`.
-
-use super::{
- Array, ArrayData, BinaryOffsetSizeTrait, BooleanArray, DecimalArray,
- FixedSizeBinaryArray, FixedSizeListArray, GenericBinaryArray, GenericListArray,
- GenericStringArray, NullArray, OffsetSizeTrait, PrimitiveArray,
- StringOffsetSizeTrait, StructArray,
-};
-
-use crate::{
- buffer::Buffer,
- datatypes::{ArrowPrimitiveType, DataType, IntervalUnit},
-};
-
-mod boolean;
-mod decimal;
-mod dictionary;
-mod fixed_binary;
-mod fixed_list;
-mod list;
-mod null;
-mod primitive;
-mod structure;
-mod utils;
-mod variable_size;
-
-// these methods assume the same type, len and null count.
-// For this reason, they are not exposed and are instead used
-// to build the generic functions below (`equal_range` and `equal`).
-use boolean::boolean_equal;
-use decimal::decimal_equal;
-use dictionary::dictionary_equal;
-use fixed_binary::fixed_binary_equal;
-use fixed_list::fixed_list_equal;
-use list::list_equal;
-use null::null_equal;
-use primitive::primitive_equal;
-use structure::struct_equal;
-use variable_size::variable_sized_equal;
-
-impl PartialEq for dyn Array {
- fn eq(&self, other: &Self) -> bool {
- equal(self.data(), other.data())
- }
-}
-
-impl<T: Array> PartialEq<T> for dyn Array {
- fn eq(&self, other: &T) -> bool {
- equal(self.data(), other.data())
- }
-}
-
-impl PartialEq for NullArray {
- fn eq(&self, other: &NullArray) -> bool {
- equal(self.data(), other.data())
- }
-}
-
-impl<T: ArrowPrimitiveType> PartialEq for PrimitiveArray<T> {
- fn eq(&self, other: &PrimitiveArray<T>) -> bool {
- equal(self.data(), other.data())
- }
-}
-
-impl PartialEq for BooleanArray {
- fn eq(&self, other: &BooleanArray) -> bool {
- equal(self.data(), other.data())
- }
-}
-
-impl<OffsetSize: StringOffsetSizeTrait> PartialEq for GenericStringArray<OffsetSize> {
- fn eq(&self, other: &Self) -> bool {
- equal(self.data(), other.data())
- }
-}
-
-impl<OffsetSize: BinaryOffsetSizeTrait> PartialEq for GenericBinaryArray<OffsetSize> {
- fn eq(&self, other: &Self) -> bool {
- equal(self.data(), other.data())
- }
-}
-
-impl PartialEq for FixedSizeBinaryArray {
- fn eq(&self, other: &Self) -> bool {
- equal(self.data(), other.data())
- }
-}
-
-impl PartialEq for DecimalArray {
- fn eq(&self, other: &Self) -> bool {
- equal(self.data(), other.data())
- }
-}
-
-impl<OffsetSize: OffsetSizeTrait> PartialEq for GenericListArray<OffsetSize> {
- fn eq(&self, other: &Self) -> bool {
- equal(self.data(), other.data())
- }
-}
-
-impl PartialEq for FixedSizeListArray {
- fn eq(&self, other: &Self) -> bool {
- equal(self.data(), other.data())
- }
-}
-
-impl PartialEq for StructArray {
- fn eq(&self, other: &Self) -> bool {
- equal(self.data(), other.data())
- }
-}
-
-/// Compares the values of two [ArrayData] starting at `lhs_start` and `rhs_start` respectively
-/// for `len` slots. The null buffers `lhs_nulls` and `rhs_nulls` inherit parent nullability.
-///
-/// If an array is a child of a struct or list, the array's nulls have to be merged with the parent.
-/// This then affects the null count of the array, thus the merged nulls are passed separately
-/// as `lhs_nulls` and `rhs_nulls` variables to functions.
-/// The nulls are merged with a bitwise AND, and null counts are recomputed where necessary.
-#[inline]
-fn equal_values(
- lhs: &ArrayData,
- rhs: &ArrayData,
- lhs_nulls: Option<&Buffer>,
- rhs_nulls: Option<&Buffer>,
- lhs_start: usize,
- rhs_start: usize,
- len: usize,
-) -> bool {
- match lhs.data_type() {
- DataType::Null => null_equal(lhs, rhs, lhs_start, rhs_start, len),
- DataType::Boolean => {
- boolean_equal(lhs, rhs, lhs_nulls, rhs_nulls, lhs_start, rhs_start, len)
- }
- DataType::UInt8 => primitive_equal::<u8>(
- lhs, rhs, lhs_nulls, rhs_nulls, lhs_start, rhs_start, len,
- ),
- DataType::UInt16 => primitive_equal::<u16>(
- lhs, rhs, lhs_nulls, rhs_nulls, lhs_start, rhs_start, len,
- ),
- DataType::UInt32 => primitive_equal::<u32>(
- lhs, rhs, lhs_nulls, rhs_nulls, lhs_start, rhs_start, len,
- ),
- DataType::UInt64 => primitive_equal::<u64>(
- lhs, rhs, lhs_nulls, rhs_nulls, lhs_start, rhs_start, len,
- ),
- DataType::Int8 => primitive_equal::<i8>(
- lhs, rhs, lhs_nulls, rhs_nulls, lhs_start, rhs_start, len,
- ),
- DataType::Int16 => primitive_equal::<i16>(
- lhs, rhs, lhs_nulls, rhs_nulls, lhs_start, rhs_start, len,
- ),
- DataType::Int32 => primitive_equal::<i32>(
- lhs, rhs, lhs_nulls, rhs_nulls, lhs_start, rhs_start, len,
- ),
- DataType::Int64 => primitive_equal::<i64>(
- lhs, rhs, lhs_nulls, rhs_nulls, lhs_start, rhs_start, len,
- ),
- DataType::Float32 => primitive_equal::<f32>(
- lhs, rhs, lhs_nulls, rhs_nulls, lhs_start, rhs_start, len,
- ),
- DataType::Float64 => primitive_equal::<f64>(
- lhs, rhs, lhs_nulls, rhs_nulls, lhs_start, rhs_start, len,
- ),
- DataType::Date32
- | DataType::Time32(_)
- | DataType::Interval(IntervalUnit::YearMonth) => primitive_equal::<i32>(
- lhs, rhs, lhs_nulls, rhs_nulls, lhs_start, rhs_start, len,
- ),
- DataType::Date64
- | DataType::Interval(IntervalUnit::DayTime)
- | DataType::Time64(_)
- | DataType::Timestamp(_, _)
- | DataType::Duration(_) => primitive_equal::<i64>(
- lhs, rhs, lhs_nulls, rhs_nulls, lhs_start, rhs_start, len,
- ),
- DataType::Utf8 | DataType::Binary => variable_sized_equal::<i32>(
- lhs, rhs, lhs_nulls, rhs_nulls, lhs_start, rhs_start, len,
- ),
- DataType::LargeUtf8 | DataType::LargeBinary => variable_sized_equal::<i64>(
- lhs, rhs, lhs_nulls, rhs_nulls, lhs_start, rhs_start, len,
- ),
- DataType::FixedSizeBinary(_) => {
- fixed_binary_equal(lhs, rhs, lhs_nulls, rhs_nulls, lhs_start, rhs_start, len)
- }
- DataType::Decimal(_, _) => {
- decimal_equal(lhs, rhs, lhs_nulls, rhs_nulls, lhs_start, rhs_start, len)
- }
- DataType::List(_) => {
- list_equal::<i32>(lhs, rhs, lhs_nulls, rhs_nulls, lhs_start, rhs_start, len)
- }
- DataType::LargeList(_) => {
- list_equal::<i64>(lhs, rhs, lhs_nulls, rhs_nulls, lhs_start, rhs_start, len)
- }
- DataType::FixedSizeList(_, _) => {
- fixed_list_equal(lhs, rhs, lhs_nulls, rhs_nulls, lhs_start, rhs_start, len)
- }
- DataType::Struct(_) => {
- struct_equal(lhs, rhs, lhs_nulls, rhs_nulls, lhs_start, rhs_start, len)
- }
- DataType::Union(_) => unimplemented!("See ARROW-8576"),
- DataType::Dictionary(data_type, _) => match data_type.as_ref() {
- DataType::Int8 => dictionary_equal::<i8>(
- lhs, rhs, lhs_nulls, rhs_nulls, lhs_start, rhs_start, len,
- ),
- DataType::Int16 => dictionary_equal::<i16>(
- lhs, rhs, lhs_nulls, rhs_nulls, lhs_start, rhs_start, len,
- ),
- DataType::Int32 => dictionary_equal::<i32>(
- lhs, rhs, lhs_nulls, rhs_nulls, lhs_start, rhs_start, len,
- ),
- DataType::Int64 => dictionary_equal::<i64>(
- lhs, rhs, lhs_nulls, rhs_nulls, lhs_start, rhs_start, len,
- ),
- DataType::UInt8 => dictionary_equal::<u8>(
- lhs, rhs, lhs_nulls, rhs_nulls, lhs_start, rhs_start, len,
- ),
- DataType::UInt16 => dictionary_equal::<u16>(
- lhs, rhs, lhs_nulls, rhs_nulls, lhs_start, rhs_start, len,
- ),
- DataType::UInt32 => dictionary_equal::<u32>(
- lhs, rhs, lhs_nulls, rhs_nulls, lhs_start, rhs_start, len,
- ),
- DataType::UInt64 => dictionary_equal::<u64>(
- lhs, rhs, lhs_nulls, rhs_nulls, lhs_start, rhs_start, len,
- ),
- _ => unreachable!(),
- },
- DataType::Float16 => unreachable!(),
- }
-}
-
-fn equal_range(
- lhs: &ArrayData,
- rhs: &ArrayData,
- lhs_nulls: Option<&Buffer>,
- rhs_nulls: Option<&Buffer>,
- lhs_start: usize,
- rhs_start: usize,
- len: usize,
-) -> bool {
- utils::base_equal(lhs, rhs)
- && utils::equal_nulls(lhs, rhs, lhs_nulls, rhs_nulls, lhs_start, rhs_start, len)
- && equal_values(lhs, rhs, lhs_nulls, rhs_nulls, lhs_start, rhs_start, len)
-}
-
-/// Logically compares two [ArrayData].
-/// Two arrays are logically equal if and only if:
-/// * their data types are equal
-/// * their lengths are equal
-/// * their null counts are equal
-/// * their null bitmaps are equal
-/// * each of their items are equal
-/// two items are equal when their in-memory representation is physically equal (i.e. same bit content).
-/// The physical comparison depend on the data type.
-/// # Panics
-/// This function may panic whenever any of the [ArrayData] does not follow the Arrow specification.
-/// (e.g. wrong number of buffers, buffer `len` does not correspond to the declared `len`)
-pub fn equal(lhs: &ArrayData, rhs: &ArrayData) -> bool {
- let lhs_nulls = lhs.null_buffer();
- let rhs_nulls = rhs.null_buffer();
- utils::base_equal(lhs, rhs)
- && lhs.null_count() == rhs.null_count()
- && utils::equal_nulls(lhs, rhs, lhs_nulls, rhs_nulls, 0, 0, lhs.len())
- && equal_values(lhs, rhs, lhs_nulls, rhs_nulls, 0, 0, lhs.len())
-}
-
-#[cfg(test)]
-mod tests {
- use std::convert::TryFrom;
- use std::sync::Arc;
-
- use crate::array::{
- array::Array, ArrayDataBuilder, ArrayRef, BinaryOffsetSizeTrait, BooleanArray,
- DecimalBuilder, FixedSizeBinaryBuilder, FixedSizeListBuilder, GenericBinaryArray,
- Int32Builder, ListBuilder, NullArray, PrimitiveBuilder, StringArray,
- StringDictionaryBuilder, StringOffsetSizeTrait, StructArray,
- };
- use crate::array::{GenericStringArray, Int32Array};
- use crate::buffer::Buffer;
- use crate::datatypes::{Field, Int16Type, ToByteSlice};
-
- use super::*;
-
- #[test]
- fn test_null_equal() {
- let a = NullArray::new(12);
- let a = a.data();
- let b = NullArray::new(12);
- let b = b.data();
- test_equal(&a, &b, true);
-
- let b = NullArray::new(10);
- let b = b.data();
- test_equal(&a, &b, false);
-
- // Test the case where offset != 0
-
- let a_slice = a.slice(2, 3);
- let b_slice = b.slice(1, 3);
- test_equal(&a_slice, &b_slice, true);
-
- let a_slice = a.slice(5, 4);
- let b_slice = b.slice(3, 3);
- test_equal(&a_slice, &b_slice, false);
- }
-
- #[test]
- fn test_boolean_equal() {
- let a = BooleanArray::from(vec![false, false, true]);
- let a = a.data();
- let b = BooleanArray::from(vec![false, false, true]);
- let b = b.data();
- test_equal(&a, &b, true);
-
- let b = BooleanArray::from(vec![false, false, false]);
- let b = b.data();
- test_equal(&a, &b, false);
- }
-
- #[test]
- fn test_boolean_equal_nulls() {
- let a = BooleanArray::from(vec![Some(false), None, None, Some(true)]);
- let a = a.data();
- let b = BooleanArray::from(vec![Some(false), None, None, Some(true)]);
- let b = b.data();
- test_equal(&a, &b, true);
-
- let b = BooleanArray::from(vec![None, None, None, Some(true)]);
- let b = b.data();
- test_equal(&a, &b, false);
-
- let b = BooleanArray::from(vec![Some(true), None, None, Some(true)]);
- let b = b.data();
- test_equal(&a, &b, false);
- }
-
- #[test]
- fn test_boolean_equal_offset() {
- let a = BooleanArray::from(vec![false, true, false, true, false, false, true]);
- let a = a.data();
- let b =
- BooleanArray::from(vec![true, false, false, false, true, false, true, true]);
- let b = b.data();
- assert_eq!(equal(a, b), false);
- assert_eq!(equal(b, a), false);
-
- let a_slice = a.slice(2, 3);
- let b_slice = b.slice(3, 3);
- assert_eq!(equal(&a_slice, &b_slice), true);
- assert_eq!(equal(&b_slice, &a_slice), true);
-
- let a_slice = a.slice(3, 4);
- let b_slice = b.slice(4, 4);
- assert_eq!(equal(&a_slice, &b_slice), false);
- assert_eq!(equal(&b_slice, &a_slice), false);
-
- // Test the optimization cases where null_count == 0 and starts at 0 and len >= size_of(u8)
-
- // Elements fill in `u8`'s exactly.
- let mut vector = vec![false, false, true, true, true, true, true, true];
- let a = BooleanArray::from(vector.clone());
- let a = a.data();
- let b = BooleanArray::from(vector.clone());
- let b = b.data();
- test_equal(&a, &b, true);
-
- // Elements fill in `u8`s + suffix bits.
- vector.push(true);
- let a = BooleanArray::from(vector.clone());
- let a = a.data();
- let b = BooleanArray::from(vector);
- let b = b.data();
- test_equal(&a, &b, true);
- }
-
- #[test]
- fn test_primitive() {
- let cases = vec![
- (
- vec![Some(1), Some(2), Some(3)],
- vec![Some(1), Some(2), Some(3)],
- true,
- ),
- (
- vec![Some(1), Some(2), Some(3)],
- vec![Some(1), Some(2), Some(4)],
- false,
- ),
- (
- vec![Some(1), Some(2), None],
- vec![Some(1), Some(2), None],
- true,
- ),
- (
- vec![Some(1), None, Some(3)],
- vec![Some(1), Some(2), None],
- false,
- ),
- (
- vec![Some(1), None, None],
- vec![Some(1), Some(2), None],
- false,
- ),
- ];
-
- for (lhs, rhs, expected) in cases {
- let lhs = Int32Array::from(lhs);
- let lhs = lhs.data();
- let rhs = Int32Array::from(rhs);
- let rhs = rhs.data();
- test_equal(&lhs, &rhs, expected);
- }
- }
-
- #[test]
- fn test_primitive_slice() {
- let cases = vec![
- (
- vec![Some(1), Some(2), Some(3)],
- (0, 1),
- vec![Some(1), Some(2), Some(3)],
- (0, 1),
- true,
- ),
- (
- vec![Some(1), Some(2), Some(3)],
- (1, 1),
- vec![Some(1), Some(2), Some(3)],
- (2, 1),
- false,
- ),
- (
- vec![Some(1), Some(2), None],
- (1, 1),
- vec![Some(1), None, Some(2)],
- (2, 1),
- true,
- ),
- (
- vec![None, Some(2), None],
- (1, 1),
- vec![None, None, Some(2)],
- (2, 1),
- true,
- ),
- (
- vec![Some(1), None, Some(2), None, Some(3)],
- (2, 2),
- vec![None, Some(2), None, Some(3)],
- (1, 2),
- true,
- ),
- ];
-
- for (lhs, slice_lhs, rhs, slice_rhs, expected) in cases {
- let lhs = Int32Array::from(lhs);
- let lhs = lhs.data();
- let lhs = lhs.slice(slice_lhs.0, slice_lhs.1);
- let rhs = Int32Array::from(rhs);
- let rhs = rhs.data();
- let rhs = rhs.slice(slice_rhs.0, slice_rhs.1);
-
- test_equal(&lhs, &rhs, expected);
- }
- }
-
- fn test_equal(lhs: &ArrayData, rhs: &ArrayData, expected: bool) {
- // equality is symmetric
- assert_eq!(equal(lhs, lhs), true, "\n{:?}\n{:?}", lhs, lhs);
- assert_eq!(equal(rhs, rhs), true, "\n{:?}\n{:?}", rhs, rhs);
-
- assert_eq!(equal(lhs, rhs), expected, "\n{:?}\n{:?}", lhs, rhs);
- assert_eq!(equal(rhs, lhs), expected, "\n{:?}\n{:?}", rhs, lhs);
- }
-
- fn binary_cases() -> Vec<(Vec<Option<String>>, Vec<Option<String>>, bool)> {
- let base = vec![
- Some("hello".to_owned()),
- None,
- None,
- Some("world".to_owned()),
- None,
- None,
- ];
- let not_base = vec![
- Some("hello".to_owned()),
- Some("foo".to_owned()),
- None,
- Some("world".to_owned()),
- None,
- None,
- ];
- vec![
- (
- vec![Some("hello".to_owned()), Some("world".to_owned())],
- vec![Some("hello".to_owned()), Some("world".to_owned())],
- true,
- ),
- (
- vec![Some("hello".to_owned()), Some("world".to_owned())],
- vec![Some("hello".to_owned()), Some("arrow".to_owned())],
- false,
- ),
- (base.clone(), base.clone(), true),
- (base, not_base, false),
- ]
- }
-
- fn test_generic_string_equal<OffsetSize: StringOffsetSizeTrait>() {
- let cases = binary_cases();
-
- for (lhs, rhs, expected) in cases {
- let lhs = lhs.iter().map(|x| x.as_deref()).collect();
- let rhs = rhs.iter().map(|x| x.as_deref()).collect();
- let lhs = GenericStringArray::<OffsetSize>::from_opt_vec(lhs);
- let lhs = lhs.data();
- let rhs = GenericStringArray::<OffsetSize>::from_opt_vec(rhs);
- let rhs = rhs.data();
- test_equal(lhs, rhs, expected);
- }
- }
-
- #[test]
- fn test_string_equal() {
- test_generic_string_equal::<i32>()
- }
-
- #[test]
- fn test_large_string_equal() {
- test_generic_string_equal::<i64>()
- }
-
- fn test_generic_binary_equal<OffsetSize: BinaryOffsetSizeTrait>() {
- let cases = binary_cases();
-
- for (lhs, rhs, expected) in cases {
- let lhs = lhs
- .iter()
- .map(|x| x.as_deref().map(|x| x.as_bytes()))
- .collect();
- let rhs = rhs
- .iter()
- .map(|x| x.as_deref().map(|x| x.as_bytes()))
- .collect();
- let lhs = GenericBinaryArray::<OffsetSize>::from_opt_vec(lhs);
- let lhs = lhs.data();
- let rhs = GenericBinaryArray::<OffsetSize>::from_opt_vec(rhs);
- let rhs = rhs.data();
- test_equal(lhs, rhs, expected);
- }
- }
-
- #[test]
- fn test_binary_equal() {
- test_generic_binary_equal::<i32>()
- }
-
- #[test]
- fn test_large_binary_equal() {
- test_generic_binary_equal::<i64>()
- }
-
- #[test]
- fn test_string_offset() {
- let a = StringArray::from(vec![Some("a"), None, Some("b")]);
- let a = a.data();
- let a = a.slice(2, 1);
- let b = StringArray::from(vec![Some("b")]);
- let b = b.data();
-
- test_equal(&a, &b, true);
- }
-
- #[test]
- fn test_string_offset_larger() {
- let a = StringArray::from(vec![Some("a"), None, Some("b"), None, Some("c")]);
- let a = a.data();
- let b = StringArray::from(vec![None, Some("b"), None, Some("c")]);
- let b = b.data();
-
- test_equal(&a.slice(2, 2), &b.slice(0, 2), false);
- test_equal(&a.slice(2, 2), &b.slice(1, 2), true);
- test_equal(&a.slice(2, 2), &b.slice(2, 2), false);
- }
-
- #[test]
- fn test_null() {
- let a = NullArray::new(2);
- let a = a.data();
- let b = NullArray::new(2);
- let b = b.data();
- test_equal(&a, &b, true);
-
- let b = NullArray::new(1);
- let b = b.data();
- test_equal(&a, &b, false);
- }
-
- fn create_list_array<U: AsRef<[i32]>, T: AsRef<[Option<U>]>>(data: T) -> ArrayData {
- let mut builder = ListBuilder::new(Int32Builder::new(10));
- for d in data.as_ref() {
- if let Some(v) = d {
- builder.values().append_slice(v.as_ref()).unwrap();
- builder.append(true).unwrap()
- } else {
- builder.append(false).unwrap()
- }
- }
- builder.finish().data().clone()
- }
-
- #[test]
- fn test_list_equal() {
- let a = create_list_array(&[Some(&[1, 2, 3]), Some(&[4, 5, 6])]);
- let b = create_list_array(&[Some(&[1, 2, 3]), Some(&[4, 5, 6])]);
- test_equal(&a, &b, true);
-
- let b = create_list_array(&[Some(&[1, 2, 3]), Some(&[4, 5, 7])]);
- test_equal(&a, &b, false);
- }
-
- // Test the case where null_count > 0
- #[test]
- fn test_list_null() {
- let a =
- create_list_array(&[Some(&[1, 2]), None, None, Some(&[3, 4]), None, None]);
- let b =
- create_list_array(&[Some(&[1, 2]), None, None, Some(&[3, 4]), None, None]);
- test_equal(&a, &b, true);
-
- let b = create_list_array(&[
- Some(&[1, 2]),
- None,
- Some(&[5, 6]),
- Some(&[3, 4]),
- None,
- None,
- ]);
- test_equal(&a, &b, false);
-
- let b =
- create_list_array(&[Some(&[1, 2]), None, None, Some(&[3, 5]), None, None]);
- test_equal(&a, &b, false);
-
- // a list where the nullness of values is determined by the list's bitmap
- let c_values = Int32Array::from(vec![1, 2, -1, -2, 3, 4, -3, -4]);
- let c = ArrayDataBuilder::new(DataType::List(Box::new(Field::new(
- "item",
- DataType::Int32,
- true,
- ))))
- .len(6)
- .add_buffer(Buffer::from(vec![0i32, 2, 3, 4, 6, 7, 8].to_byte_slice()))
- .add_child_data(c_values.data().clone())
- .null_bit_buffer(Buffer::from(vec![0b00001001]))
- .build();
-
- let d_values = Int32Array::from(vec![
- Some(1),
- Some(2),
- None,
- None,
- Some(3),
- Some(4),
- None,
- None,
- ]);
- let d = ArrayDataBuilder::new(DataType::List(Box::new(Field::new(
- "item",
- DataType::Int32,
- true,
- ))))
- .len(6)
- .add_buffer(Buffer::from(vec![0i32, 2, 3, 4, 6, 7, 8].to_byte_slice()))
- .add_child_data(d_values.data().clone())
- .null_bit_buffer(Buffer::from(vec![0b00001001]))
- .build();
- test_equal(&c, &d, true);
- }
-
- // Test the case where offset != 0
- #[test]
- fn test_list_offsets() {
- let a =
- create_list_array(&[Some(&[1, 2]), None, None, Some(&[3, 4]), None, None]);
- let b =
- create_list_array(&[Some(&[1, 2]), None, None, Some(&[3, 5]), None, None]);
-
- let a_slice = a.slice(0, 3);
- let b_slice = b.slice(0, 3);
- test_equal(&a_slice, &b_slice, true);
-
- let a_slice = a.slice(0, 5);
- let b_slice = b.slice(0, 5);
- test_equal(&a_slice, &b_slice, false);
-
- let a_slice = a.slice(4, 1);
- let b_slice = b.slice(4, 1);
- test_equal(&a_slice, &b_slice, true);
- }
-
- fn create_fixed_size_binary_array<U: AsRef<[u8]>, T: AsRef<[Option<U>]>>(
- data: T,
- ) -> ArrayData {
- let mut builder = FixedSizeBinaryBuilder::new(15, 5);
-
- for d in data.as_ref() {
- if let Some(v) = d {
- builder.append_value(v.as_ref()).unwrap();
- } else {
- builder.append_null().unwrap();
- }
- }
- builder.finish().data().clone()
- }
-
- #[test]
- fn test_fixed_size_binary_equal() {
- let a = create_fixed_size_binary_array(&[Some(b"hello"), Some(b"world")]);
- let b = create_fixed_size_binary_array(&[Some(b"hello"), Some(b"world")]);
- test_equal(&a, &b, true);
-
- let b = create_fixed_size_binary_array(&[Some(b"hello"), Some(b"arrow")]);
- test_equal(&a, &b, false);
- }
-
- // Test the case where null_count > 0
- #[test]
- fn test_fixed_size_binary_null() {
- let a = create_fixed_size_binary_array(&[Some(b"hello"), None, Some(b"world")]);
- let b = create_fixed_size_binary_array(&[Some(b"hello"), None, Some(b"world")]);
- test_equal(&a, &b, true);
-
- let b = create_fixed_size_binary_array(&[Some(b"hello"), Some(b"world"), None]);
- test_equal(&a, &b, false);
-
- let b = create_fixed_size_binary_array(&[Some(b"hello"), None, Some(b"arrow")]);
- test_equal(&a, &b, false);
- }
-
- #[test]
- fn test_fixed_size_binary_offsets() {
- // Test the case where offset != 0
- let a = create_fixed_size_binary_array(&[
- Some(b"hello"),
- None,
- None,
- Some(b"world"),
- None,
- None,
- ]);
- let b = create_fixed_size_binary_array(&[
- Some(b"hello"),
- None,
- None,
- Some(b"arrow"),
- None,
- None,
- ]);
-
- let a_slice = a.slice(0, 3);
- let b_slice = b.slice(0, 3);
- test_equal(&a_slice, &b_slice, true);
-
- let a_slice = a.slice(0, 5);
- let b_slice = b.slice(0, 5);
- test_equal(&a_slice, &b_slice, false);
-
- let a_slice = a.slice(4, 1);
- let b_slice = b.slice(4, 1);
- test_equal(&a_slice, &b_slice, true);
-
- let a_slice = a.slice(3, 1);
- let b_slice = b.slice(3, 1);
- test_equal(&a_slice, &b_slice, false);
- }
-
- fn create_decimal_array(data: &[Option<i128>]) -> ArrayData {
- let mut builder = DecimalBuilder::new(20, 23, 6);
-
- for d in data {
- if let Some(v) = d {
- builder.append_value(*v).unwrap();
- } else {
- builder.append_null().unwrap();
- }
- }
- builder.finish().data().clone()
- }
-
- #[test]
- fn test_decimal_equal() {
- let a = create_decimal_array(&[Some(8_887_000_000), Some(-8_887_000_000)]);
- let b = create_decimal_array(&[Some(8_887_000_000), Some(-8_887_000_000)]);
- test_equal(&a, &b, true);
-
- let b = create_decimal_array(&[Some(15_887_000_000), Some(-8_887_000_000)]);
- test_equal(&a, &b, false);
- }
-
- // Test the case where null_count > 0
- #[test]
- fn test_decimal_null() {
- let a = create_decimal_array(&[Some(8_887_000_000), None, Some(-8_887_000_000)]);
- let b = create_decimal_array(&[Some(8_887_000_000), None, Some(-8_887_000_000)]);
- test_equal(&a, &b, true);
-
- let b = create_decimal_array(&[Some(8_887_000_000), Some(-8_887_000_000), None]);
- test_equal(&a, &b, false);
-
- let b = create_decimal_array(&[Some(15_887_000_000), None, Some(-8_887_000_000)]);
- test_equal(&a, &b, false);
- }
-
- #[test]
- fn test_decimal_offsets() {
- // Test the case where offset != 0
- let a = create_decimal_array(&[
- Some(8_887_000_000),
- None,
- None,
- Some(-8_887_000_000),
- None,
- None,
- ]);
- let b = create_decimal_array(&[
- None,
- Some(8_887_000_000),
- None,
- None,
- Some(15_887_000_000),
- None,
- None,
- ]);
-
- let a_slice = a.slice(0, 3);
- let b_slice = b.slice(1, 3);
- test_equal(&a_slice, &b_slice, true);
-
- let a_slice = a.slice(0, 5);
- let b_slice = b.slice(1, 5);
- test_equal(&a_slice, &b_slice, false);
-
- let a_slice = a.slice(4, 1);
- let b_slice = b.slice(5, 1);
- test_equal(&a_slice, &b_slice, true);
-
- let a_slice = a.slice(3, 3);
- let b_slice = b.slice(4, 3);
- test_equal(&a_slice, &b_slice, false);
-
- let a_slice = a.slice(1, 3);
- let b_slice = b.slice(2, 3);
- test_equal(&a_slice, &b_slice, false);
-
- let b = create_decimal_array(&[
- None,
- None,
- None,
- Some(-8_887_000_000),
- Some(-3_000),
- None,
- ]);
- let a_slice = a.slice(1, 3);
- let b_slice = b.slice(1, 3);
- test_equal(&a_slice, &b_slice, true);
- }
-
- /// Create a fixed size list of 2 value lengths
- fn create_fixed_size_list_array<U: AsRef<[i32]>, T: AsRef<[Option<U>]>>(
- data: T,
- ) -> ArrayData {
- let mut builder = FixedSizeListBuilder::new(Int32Builder::new(10), 3);
-
- for d in data.as_ref() {
- if let Some(v) = d {
- builder.values().append_slice(v.as_ref()).unwrap();
- builder.append(true).unwrap()
- } else {
- for _ in 0..builder.value_length() {
- builder.values().append_null().unwrap();
- }
- builder.append(false).unwrap()
- }
- }
- builder.finish().data().clone()
- }
-
- #[test]
- fn test_fixed_size_list_equal() {
- let a = create_fixed_size_list_array(&[Some(&[1, 2, 3]), Some(&[4, 5, 6])]);
- let b = create_fixed_size_list_array(&[Some(&[1, 2, 3]), Some(&[4, 5, 6])]);
- test_equal(&a, &b, true);
-
- let b = create_fixed_size_list_array(&[Some(&[1, 2, 3]), Some(&[4, 5, 7])]);
- test_equal(&a, &b, false);
- }
-
- // Test the case where null_count > 0
- #[test]
- fn test_fixed_list_null() {
- let a = create_fixed_size_list_array(&[
- Some(&[1, 2, 3]),
- None,
- None,
- Some(&[4, 5, 6]),
- None,
- None,
- ]);
- let b = create_fixed_size_list_array(&[
- Some(&[1, 2, 3]),
- None,
- None,
- Some(&[4, 5, 6]),
- None,
- None,
- ]);
- test_equal(&a, &b, true);
-
- let b = create_fixed_size_list_array(&[
- Some(&[1, 2, 3]),
- None,
- Some(&[7, 8, 9]),
- Some(&[4, 5, 6]),
- None,
- None,
- ]);
- test_equal(&a, &b, false);
-
- let b = create_fixed_size_list_array(&[
- Some(&[1, 2, 3]),
- None,
- None,
- Some(&[3, 6, 9]),
- None,
- None,
- ]);
- test_equal(&a, &b, false);
- }
-
- #[test]
- fn test_fixed_list_offsets() {
- // Test the case where offset != 0
- let a = create_fixed_size_list_array(&[
- Some(&[1, 2, 3]),
- None,
- None,
- Some(&[4, 5, 6]),
- None,
- None,
- ]);
- let b = create_fixed_size_list_array(&[
- Some(&[1, 2, 3]),
- None,
- None,
- Some(&[3, 6, 9]),
- None,
- None,
- ]);
-
- let a_slice = a.slice(0, 3);
- let b_slice = b.slice(0, 3);
- test_equal(&a_slice, &b_slice, true);
-
- let a_slice = a.slice(0, 5);
- let b_slice = b.slice(0, 5);
- test_equal(&a_slice, &b_slice, false);
-
- let a_slice = a.slice(4, 1);
- let b_slice = b.slice(4, 1);
- test_equal(&a_slice, &b_slice, true);
- }
-
- #[test]
- fn test_struct_equal() {
- let strings: ArrayRef = Arc::new(StringArray::from(vec![
- Some("joe"),
- None,
- None,
- Some("mark"),
- Some("doe"),
- ]));
- let ints: ArrayRef = Arc::new(Int32Array::from(vec![
- Some(1),
- Some(2),
- None,
- Some(4),
- Some(5),
- ]));
-
- let a =
- StructArray::try_from(vec![("f1", strings.clone()), ("f2", ints.clone())])
- .unwrap();
- let a = a.data();
-
- let b = StructArray::try_from(vec![("f1", strings), ("f2", ints)]).unwrap();
- let b = b.data();
-
- test_equal(&a, &b, true);
- }
-
- #[test]
- fn test_struct_equal_null() {
- let strings: ArrayRef = Arc::new(StringArray::from(vec![
- Some("joe"),
- None,
- None,
- Some("mark"),
- Some("doe"),
- ]));
- let ints: ArrayRef = Arc::new(Int32Array::from(vec![
- Some(1),
- Some(2),
- None,
- Some(4),
- Some(5),
- ]));
- let ints_non_null: ArrayRef = Arc::new(Int32Array::from(vec![1, 2, 3, 4, 0]));
-
- let a = ArrayData::builder(DataType::Struct(vec![
- Field::new("f1", DataType::Utf8, true),
- Field::new("f2", DataType::Int32, true),
- ]))
- .null_bit_buffer(Buffer::from(vec![0b00001011]))
- .len(5)
- .add_child_data(strings.data_ref().clone())
- .add_child_data(ints.data_ref().clone())
- .build();
- let a = crate::array::make_array(a);
-
- let b = ArrayData::builder(DataType::Struct(vec![
- Field::new("f1", DataType::Utf8, true),
- Field::new("f2", DataType::Int32, true),
- ]))
- .null_bit_buffer(Buffer::from(vec![0b00001011]))
- .len(5)
- .add_child_data(strings.data_ref().clone())
- .add_child_data(ints_non_null.data_ref().clone())
- .build();
- let b = crate::array::make_array(b);
-
- test_equal(a.data_ref(), b.data_ref(), true);
-
- // test with arrays that are not equal
- let c_ints_non_null: ArrayRef = Arc::new(Int32Array::from(vec![1, 2, 3, 0, 4]));
- let c = ArrayData::builder(DataType::Struct(vec![
- Field::new("f1", DataType::Utf8, true),
- Field::new("f2", DataType::Int32, true),
- ]))
- .null_bit_buffer(Buffer::from(vec![0b00001011]))
- .len(5)
- .add_child_data(strings.data_ref().clone())
- .add_child_data(c_ints_non_null.data_ref().clone())
- .build();
- let c = crate::array::make_array(c);
-
- test_equal(a.data_ref(), c.data_ref(), false);
-
- // test a nested struct
- let a = ArrayData::builder(DataType::Struct(vec![Field::new(
- "f3",
- a.data_type().clone(),
- true,
- )]))
- .null_bit_buffer(Buffer::from(vec![0b00011110]))
- .len(5)
- .add_child_data(a.data_ref().clone())
- .build();
- let a = crate::array::make_array(a);
-
- // reconstruct b, but with different data where the first struct is null
- let strings: ArrayRef = Arc::new(StringArray::from(vec![
- Some("joanne"), // difference
- None,
- None,
- Some("mark"),
- Some("doe"),
- ]));
- let b = ArrayData::builder(DataType::Struct(vec![
- Field::new("f1", DataType::Utf8, true),
- Field::new("f2", DataType::Int32, true),
- ]))
- .null_bit_buffer(Buffer::from(vec![0b00001011]))
- .len(5)
- .add_child_data(strings.data_ref().clone())
- .add_child_data(ints_non_null.data_ref().clone())
- .build();
-
- let b = ArrayData::builder(DataType::Struct(vec![Field::new(
- "f3",
- b.data_type().clone(),
- true,
- )]))
- .null_bit_buffer(Buffer::from(vec![0b00011110]))
- .len(5)
- .add_child_data(b)
- .build();
- let b = crate::array::make_array(b);
-
- test_equal(a.data_ref(), b.data_ref(), true);
- }
-
- #[test]
- fn test_struct_equal_null_variable_size() {
- // the string arrays differ, but where the struct array is null
- let strings1: ArrayRef = Arc::new(StringArray::from(vec![
- Some("joe"),
- None,
- None,
- Some("mark"),
- Some("doel"),
- ]));
- let strings2: ArrayRef = Arc::new(StringArray::from(vec![
- Some("joel"),
- None,
- None,
- Some("mark"),
- Some("doe"),
- ]));
-
- let a = ArrayData::builder(DataType::Struct(vec![Field::new(
- "f1",
- DataType::Utf8,
- true,
- )]))
- .null_bit_buffer(Buffer::from(vec![0b00001010]))
- .len(5)
- .add_child_data(strings1.data_ref().clone())
- .build();
- let a = crate::array::make_array(a);
-
- let b = ArrayData::builder(DataType::Struct(vec![Field::new(
- "f1",
- DataType::Utf8,
- true,
- )]))
- .null_bit_buffer(Buffer::from(vec![0b00001010]))
- .len(5)
- .add_child_data(strings2.data_ref().clone())
- .build();
- let b = crate::array::make_array(b);
-
- test_equal(a.data_ref(), b.data_ref(), true);
-
- // test with arrays that are not equal
- let strings3: ArrayRef = Arc::new(StringArray::from(vec![
- Some("mark"),
- None,
- None,
- Some("doe"),
- Some("joe"),
- ]));
- let c = ArrayData::builder(DataType::Struct(vec![Field::new(
- "f1",
- DataType::Utf8,
- true,
- )]))
- .null_bit_buffer(Buffer::from(vec![0b00001011]))
- .len(5)
- .add_child_data(strings3.data_ref().clone())
- .build();
- let c = crate::array::make_array(c);
-
- test_equal(a.data_ref(), c.data_ref(), false);
- }
-
- fn create_dictionary_array(values: &[&str], keys: &[Option<&str>]) -> ArrayData {
- let values = StringArray::from(values.to_vec());
- let mut builder = StringDictionaryBuilder::new_with_dictionary(
- PrimitiveBuilder::<Int16Type>::new(3),
- &values,
- )
- .unwrap();
- for key in keys {
- if let Some(v) = key {
- builder.append(v).unwrap();
- } else {
- builder.append_null().unwrap()
- }
- }
- builder.finish().data().clone()
- }
-
- #[test]
- fn test_dictionary_equal() {
- // (a, b, c), (1, 2, 1, 3) => (a, b, a, c)
- let a = create_dictionary_array(
- &["a", "b", "c"],
- &[Some("a"), Some("b"), Some("a"), Some("c")],
- );
- // different representation (values and keys are swapped), same result
- let b = create_dictionary_array(
- &["a", "c", "b"],
- &[Some("a"), Some("b"), Some("a"), Some("c")],
- );
- test_equal(&a, &b, true);
-
- // different len
- let b =
- create_dictionary_array(&["a", "c", "b"], &[Some("a"), Some("b"), Some("a")]);
- test_equal(&a, &b, false);
-
- // different key
- let b = create_dictionary_array(
- &["a", "c", "b"],
- &[Some("a"), Some("b"), Some("a"), Some("a")],
- );
- test_equal(&a, &b, false);
-
- // different values, same keys
- let b = create_dictionary_array(
- &["a", "b", "d"],
- &[Some("a"), Some("b"), Some("a"), Some("d")],
- );
- test_equal(&a, &b, false);
- }
-
- #[test]
- fn test_dictionary_equal_null() {
- // (a, b, c), (1, 2, 1, 3) => (a, b, a, c)
- let a = create_dictionary_array(
- &["a", "b", "c"],
- &[Some("a"), None, Some("a"), Some("c")],
- );
-
- // equal to self
- test_equal(&a, &a, true);
-
- // different representation (values and keys are swapped), same result
- let b = create_dictionary_array(
- &["a", "c", "b"],
- &[Some("a"), None, Some("a"), Some("c")],
- );
- test_equal(&a, &b, true);
-
- // different null position
- let b = create_dictionary_array(
- &["a", "c", "b"],
- &[Some("a"), Some("b"), Some("a"), None],
- );
- test_equal(&a, &b, false);
-
- // different key
- let b = create_dictionary_array(
- &["a", "c", "b"],
- &[Some("a"), None, Some("a"), Some("a")],
- );
- test_equal(&a, &b, false);
-
- // different values, same keys
- let b = create_dictionary_array(
- &["a", "b", "d"],
- &[Some("a"), None, Some("a"), Some("d")],
- );
- test_equal(&a, &b, false);
- }
-}
diff --git a/arrow/src/array/equal/null.rs b/arrow/src/array/equal/null.rs
deleted file mode 100644
index f287a38..0000000
--- a/arrow/src/array/equal/null.rs
+++ /dev/null
@@ -1,31 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-use crate::array::ArrayData;
-
-#[inline]
-pub(super) fn null_equal(
- _lhs: &ArrayData,
- _rhs: &ArrayData,
- _lhs_start: usize,
- _rhs_start: usize,
- _len: usize,
-) -> bool {
- // a null buffer's range is always true, as every element is by definition equal (to null).
- // We only need to compare data_types
- true
-}
diff --git a/arrow/src/array/equal/primitive.rs b/arrow/src/array/equal/primitive.rs
deleted file mode 100644
index db75879..0000000
--- a/arrow/src/array/equal/primitive.rs
+++ /dev/null
@@ -1,73 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-use std::mem::size_of;
-
-use crate::array::{data::count_nulls, ArrayData};
-use crate::buffer::Buffer;
-use crate::util::bit_util::get_bit;
-
-use super::utils::equal_len;
-
-pub(super) fn primitive_equal<T>(
- lhs: &ArrayData,
- rhs: &ArrayData,
- lhs_nulls: Option<&Buffer>,
- rhs_nulls: Option<&Buffer>,
- lhs_start: usize,
- rhs_start: usize,
- len: usize,
-) -> bool {
- let byte_width = size_of::<T>();
- let lhs_values = &lhs.buffers()[0].as_slice()[lhs.offset() * byte_width..];
- let rhs_values = &rhs.buffers()[0].as_slice()[rhs.offset() * byte_width..];
-
- let lhs_null_count = count_nulls(lhs_nulls, lhs_start, len);
- let rhs_null_count = count_nulls(rhs_nulls, rhs_start, len);
-
- if lhs_null_count == 0 && rhs_null_count == 0 {
- // without nulls, we just need to compare slices
- equal_len(
- lhs_values,
- rhs_values,
- lhs_start * byte_width,
- rhs_start * byte_width,
- len * byte_width,
- )
- } else {
- // get a ref of the null buffer bytes, to use in testing for nullness
- let lhs_null_bytes = lhs_nulls.as_ref().unwrap().as_slice();
- let rhs_null_bytes = rhs_nulls.as_ref().unwrap().as_slice();
- // with nulls, we need to compare item by item whenever it is not null
- (0..len).all(|i| {
- let lhs_pos = lhs_start + i;
- let rhs_pos = rhs_start + i;
- let lhs_is_null = !get_bit(lhs_null_bytes, lhs_pos + lhs.offset());
- let rhs_is_null = !get_bit(rhs_null_bytes, rhs_pos + rhs.offset());
-
- lhs_is_null
- || (lhs_is_null == rhs_is_null)
- && equal_len(
- lhs_values,
- rhs_values,
- lhs_pos * byte_width,
- rhs_pos * byte_width,
- byte_width, // 1 * byte_width since we are comparing a single entry
- )
- })
- }
-}
diff --git a/arrow/src/array/equal/structure.rs b/arrow/src/array/equal/structure.rs
deleted file mode 100644
index b3cc402..0000000
--- a/arrow/src/array/equal/structure.rs
+++ /dev/null
@@ -1,90 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-use crate::{
- array::data::count_nulls, array::ArrayData, buffer::Buffer, util::bit_util::get_bit,
-};
-
-use super::{equal_range, utils::child_logical_null_buffer};
-
-/// Compares the values of two [ArrayData] starting at `lhs_start` and `rhs_start` respectively
-/// for `len` slots. The null buffers `lhs_nulls` and `rhs_nulls` inherit parent nullability.
-///
-/// If an array is a child of a struct or list, the array's nulls have to be merged with the parent.
-/// This then affects the null count of the array, thus the merged nulls are passed separately
-/// as `lhs_nulls` and `rhs_nulls` variables to functions.
-/// The nulls are merged with a bitwise AND, and null counts are recomputed where necessary.
-fn equal_values(
- lhs: &ArrayData,
- rhs: &ArrayData,
- lhs_nulls: Option<&Buffer>,
- rhs_nulls: Option<&Buffer>,
- lhs_start: usize,
- rhs_start: usize,
- len: usize,
-) -> bool {
- lhs.child_data()
- .iter()
- .zip(rhs.child_data())
- .all(|(lhs_values, rhs_values)| {
- // merge the null data
- let lhs_merged_nulls = child_logical_null_buffer(lhs, lhs_nulls, lhs_values);
- let rhs_merged_nulls = child_logical_null_buffer(rhs, rhs_nulls, rhs_values);
- equal_range(
- lhs_values,
- rhs_values,
- lhs_merged_nulls.as_ref(),
- rhs_merged_nulls.as_ref(),
- lhs_start,
- rhs_start,
- len,
- )
- })
-}
-
-pub(super) fn struct_equal(
- lhs: &ArrayData,
- rhs: &ArrayData,
- lhs_nulls: Option<&Buffer>,
- rhs_nulls: Option<&Buffer>,
- lhs_start: usize,
- rhs_start: usize,
- len: usize,
-) -> bool {
- // we have to recalculate null counts from the null buffers
- let lhs_null_count = count_nulls(lhs_nulls, lhs_start, len);
- let rhs_null_count = count_nulls(rhs_nulls, rhs_start, len);
- if lhs_null_count == 0 && rhs_null_count == 0 {
- equal_values(lhs, rhs, lhs_nulls, rhs_nulls, lhs_start, rhs_start, len)
- } else {
- // get a ref of the null buffer bytes, to use in testing for nullness
- let lhs_null_bytes = lhs_nulls.as_ref().unwrap().as_slice();
- let rhs_null_bytes = rhs_nulls.as_ref().unwrap().as_slice();
- // with nulls, we need to compare item by item whenever it is not null
- (0..len).all(|i| {
- let lhs_pos = lhs_start + i;
- let rhs_pos = rhs_start + i;
- // if both struct and child had no null buffers,
- let lhs_is_null = !get_bit(lhs_null_bytes, lhs_pos + lhs.offset());
- let rhs_is_null = !get_bit(rhs_null_bytes, rhs_pos + rhs.offset());
-
- lhs_is_null
- || (lhs_is_null == rhs_is_null)
- && equal_values(lhs, rhs, lhs_nulls, rhs_nulls, lhs_pos, rhs_pos, 1)
- })
- }
-}
diff --git a/arrow/src/array/equal/utils.rs b/arrow/src/array/equal/utils.rs
deleted file mode 100644
index d0108d2..0000000
--- a/arrow/src/array/equal/utils.rs
+++ /dev/null
@@ -1,264 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-use crate::array::{data::count_nulls, ArrayData, OffsetSizeTrait};
-use crate::bitmap::Bitmap;
-use crate::buffer::{Buffer, MutableBuffer};
-use crate::datatypes::DataType;
-use crate::util::bit_util;
-
-// whether bits along the positions are equal
-// `lhs_start`, `rhs_start` and `len` are _measured in bits_.
-#[inline]
-pub(super) fn equal_bits(
- lhs_values: &[u8],
- rhs_values: &[u8],
- lhs_start: usize,
- rhs_start: usize,
- len: usize,
-) -> bool {
- (0..len).all(|i| {
- bit_util::get_bit(lhs_values, lhs_start + i)
- == bit_util::get_bit(rhs_values, rhs_start + i)
- })
-}
-
-#[inline]
-pub(super) fn equal_nulls(
- lhs: &ArrayData,
- rhs: &ArrayData,
- lhs_nulls: Option<&Buffer>,
- rhs_nulls: Option<&Buffer>,
- lhs_start: usize,
- rhs_start: usize,
- len: usize,
-) -> bool {
- let lhs_null_count = count_nulls(lhs_nulls, lhs_start, len);
- let rhs_null_count = count_nulls(rhs_nulls, rhs_start, len);
- if lhs_null_count > 0 || rhs_null_count > 0 {
- let lhs_values = lhs_nulls.unwrap().as_slice();
- let rhs_values = rhs_nulls.unwrap().as_slice();
- equal_bits(
- lhs_values,
- rhs_values,
- lhs_start + lhs.offset(),
- rhs_start + rhs.offset(),
- len,
- )
- } else {
- true
- }
-}
-
-#[inline]
-pub(super) fn base_equal(lhs: &ArrayData, rhs: &ArrayData) -> bool {
- lhs.data_type() == rhs.data_type() && lhs.len() == rhs.len()
-}
-
-// whether the two memory regions are equal
-#[inline]
-pub(super) fn equal_len(
- lhs_values: &[u8],
- rhs_values: &[u8],
- lhs_start: usize,
- rhs_start: usize,
- len: usize,
-) -> bool {
- lhs_values[lhs_start..(lhs_start + len)] == rhs_values[rhs_start..(rhs_start + len)]
-}
-
-/// Computes the logical validity bitmap of the array data using the
-/// parent's array data. The parent should be a list or struct, else
-/// the logical bitmap of the array is returned unaltered.
-///
-/// Parent data is passed along with the parent's logical bitmap, as
-/// nested arrays could have a logical bitmap different to the physical
-/// one on the `ArrayData`.
-pub(super) fn child_logical_null_buffer(
- parent_data: &ArrayData,
- logical_null_buffer: Option<&Buffer>,
- child_data: &ArrayData,
-) -> Option<Buffer> {
- let parent_len = parent_data.len();
- let parent_bitmap = logical_null_buffer
- .cloned()
- .map(Bitmap::from)
- .unwrap_or_else(|| {
- let ceil = bit_util::ceil(parent_len, 8);
- Bitmap::from(Buffer::from(vec![0b11111111; ceil]))
- });
- let self_null_bitmap = child_data.null_bitmap().clone().unwrap_or_else(|| {
- let ceil = bit_util::ceil(child_data.len(), 8);
- Bitmap::from(Buffer::from(vec![0b11111111; ceil]))
- });
- match parent_data.data_type() {
- DataType::List(_) => Some(logical_list_bitmap::<i32>(
- parent_data,
- parent_bitmap,
- self_null_bitmap,
- )),
- DataType::LargeList(_) => Some(logical_list_bitmap::<i64>(
- parent_data,
- parent_bitmap,
- self_null_bitmap,
- )),
- DataType::FixedSizeList(_, len) => {
- let len = *len as usize;
- let array_offset = parent_data.offset();
- let bitmap_len = bit_util::ceil(parent_len * len, 8);
- let mut buffer = MutableBuffer::from_len_zeroed(bitmap_len);
- let mut null_slice = buffer.as_slice_mut();
- (array_offset..parent_len + array_offset).for_each(|index| {
- let start = index * len;
- let end = start + len;
- let mask = parent_bitmap.is_set(index);
- (start..end).for_each(|child_index| {
- if mask && self_null_bitmap.is_set(child_index) {
- bit_util::set_bit(&mut null_slice, child_index);
- }
- });
- });
- Some(buffer.into())
- }
- DataType::Struct(_) => {
- // Arrow implementations are free to pad data, which can result in null buffers not
- // having the same length.
- // Rust bitwise comparisons will return an error if left AND right is performed on
- // buffers of different length.
- // This might be a valid case during integration testing, where we read Arrow arrays
- // from IPC data, which has padding.
- //
- // We first perform a bitwise comparison, and if there is an error, we revert to a
- // slower method that indexes into the buffers one-by-one.
- let result = &parent_bitmap & &self_null_bitmap;
- if let Ok(bitmap) = result {
- return Some(bitmap.bits);
- }
- // slow path
- let array_offset = parent_data.offset();
- let mut buffer = MutableBuffer::new_null(parent_len);
- let mut null_slice = buffer.as_slice_mut();
- (0..parent_len).for_each(|index| {
- if parent_bitmap.is_set(index + array_offset)
- && self_null_bitmap.is_set(index + array_offset)
- {
- bit_util::set_bit(&mut null_slice, index);
- }
- });
- Some(buffer.into())
- }
- DataType::Union(_) => {
- unimplemented!("Logical equality not yet implemented for union arrays")
- }
- DataType::Dictionary(_, _) => {
- unimplemented!("Logical equality not yet implemented for nested dictionaries")
- }
- data_type => panic!("Data type {:?} is not a supported nested type", data_type),
- }
-}
-
-// Calculate a list child's logical bitmap/buffer
-#[inline]
-fn logical_list_bitmap<OffsetSize: OffsetSizeTrait>(
- parent_data: &ArrayData,
- parent_bitmap: Bitmap,
- child_bitmap: Bitmap,
-) -> Buffer {
- let offsets = parent_data.buffer::<OffsetSize>(0);
- let offset_start = offsets.first().unwrap().to_usize().unwrap();
- let offset_len = offsets.get(parent_data.len()).unwrap().to_usize().unwrap();
- let mut buffer = MutableBuffer::new_null(offset_len - offset_start);
- let mut null_slice = buffer.as_slice_mut();
-
- offsets
- .windows(2)
- .enumerate()
- .take(offset_len - offset_start)
- .for_each(|(index, window)| {
- let start = window[0].to_usize().unwrap();
- let end = window[1].to_usize().unwrap();
- let mask = parent_bitmap.is_set(index);
- (start..end).for_each(|child_index| {
- if mask && child_bitmap.is_set(child_index) {
- bit_util::set_bit(&mut null_slice, child_index - offset_start);
- }
- });
- });
- buffer.into()
-}
-
-#[cfg(test)]
-mod tests {
- use super::*;
-
- use crate::datatypes::{Field, ToByteSlice};
-
- #[test]
- fn test_logical_null_buffer() {
- let child_data = ArrayData::builder(DataType::Int32)
- .len(11)
- .add_buffer(Buffer::from(
- vec![1i32, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11].to_byte_slice(),
- ))
- .build();
-
- let data = ArrayData::builder(DataType::List(Box::new(Field::new(
- "item",
- DataType::Int32,
- false,
- ))))
- .len(7)
- .add_buffer(Buffer::from(vec![0, 0, 3, 5, 6, 9, 10, 11].to_byte_slice()))
- .null_bit_buffer(Buffer::from(vec![0b01011010]))
- .add_child_data(child_data.clone())
- .build();
-
- // Get the child logical null buffer. The child is non-nullable, but because the list has nulls,
- // we expect the child to logically have some nulls, inherited from the parent:
- // [1, 2, 3, null, null, 6, 7, 8, 9, null, 11]
- let nulls = child_logical_null_buffer(
- &data,
- data.null_buffer(),
- data.child_data().get(0).unwrap(),
- );
- let expected = Some(Buffer::from(vec![0b11100111, 0b00000101]));
- assert_eq!(nulls, expected);
-
- // test with offset
- let data = ArrayData::builder(DataType::List(Box::new(Field::new(
- "item",
- DataType::Int32,
- false,
- ))))
- .len(4)
- .offset(3)
- .add_buffer(Buffer::from(vec![0, 0, 3, 5, 6, 9, 10, 11].to_byte_slice()))
- // the null_bit_buffer doesn't have an offset, i.e. cleared the 3 offset bits 0b[---]01011[010]
- .null_bit_buffer(Buffer::from(vec![0b00001011]))
- .add_child_data(child_data)
- .build();
-
- let nulls = child_logical_null_buffer(
- &data,
- data.null_buffer(),
- data.child_data().get(0).unwrap(),
- );
-
- let expected = Some(Buffer::from(vec![0b00101111]));
- assert_eq!(nulls, expected);
- }
-}
diff --git a/arrow/src/array/equal/variable_size.rs b/arrow/src/array/equal/variable_size.rs
deleted file mode 100644
index ecb3bc2..0000000
--- a/arrow/src/array/equal/variable_size.rs
+++ /dev/null
@@ -1,110 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-use crate::buffer::Buffer;
-use crate::util::bit_util::get_bit;
-use crate::{
- array::data::count_nulls,
- array::{ArrayData, OffsetSizeTrait},
-};
-
-use super::utils::equal_len;
-
-fn offset_value_equal<T: OffsetSizeTrait>(
- lhs_values: &[u8],
- rhs_values: &[u8],
- lhs_offsets: &[T],
- rhs_offsets: &[T],
- lhs_pos: usize,
- rhs_pos: usize,
- len: usize,
-) -> bool {
- let lhs_start = lhs_offsets[lhs_pos].to_usize().unwrap();
- let rhs_start = rhs_offsets[rhs_pos].to_usize().unwrap();
- let lhs_len = lhs_offsets[lhs_pos + len] - lhs_offsets[lhs_pos];
- let rhs_len = rhs_offsets[rhs_pos + len] - rhs_offsets[rhs_pos];
-
- lhs_len == rhs_len
- && equal_len(
- lhs_values,
- rhs_values,
- lhs_start,
- rhs_start,
- lhs_len.to_usize().unwrap(),
- )
-}
-
-pub(super) fn variable_sized_equal<T: OffsetSizeTrait>(
- lhs: &ArrayData,
- rhs: &ArrayData,
- lhs_nulls: Option<&Buffer>,
- rhs_nulls: Option<&Buffer>,
- lhs_start: usize,
- rhs_start: usize,
- len: usize,
-) -> bool {
- let lhs_offsets = lhs.buffer::<T>(0);
- let rhs_offsets = rhs.buffer::<T>(0);
-
- // the offsets of the `ArrayData` are ignored as they are only applied to the offset buffer.
- let lhs_values = lhs.buffers()[1].as_slice();
- let rhs_values = rhs.buffers()[1].as_slice();
-
- let lhs_null_count = count_nulls(lhs_nulls, lhs_start, len);
- let rhs_null_count = count_nulls(rhs_nulls, rhs_start, len);
-
- if lhs_null_count == 0
- && rhs_null_count == 0
- && !lhs_values.is_empty()
- && !rhs_values.is_empty()
- {
- offset_value_equal(
- lhs_values,
- rhs_values,
- lhs_offsets,
- rhs_offsets,
- lhs_start,
- rhs_start,
- len,
- )
- } else {
- (0..len).all(|i| {
- let lhs_pos = lhs_start + i;
- let rhs_pos = rhs_start + i;
-
- // the null bits can still be `None`, so we don't unwrap
- let lhs_is_null = !lhs_nulls
- .map(|v| get_bit(v.as_slice(), lhs.offset() + lhs_pos))
- .unwrap_or(false);
- let rhs_is_null = !rhs_nulls
- .map(|v| get_bit(v.as_slice(), rhs.offset() + rhs_pos))
- .unwrap_or(false);
-
- lhs_is_null
- || (lhs_is_null == rhs_is_null)
- && offset_value_equal(
- lhs_values,
- rhs_values,
- lhs_offsets,
- rhs_offsets,
- lhs_pos,
- rhs_pos,
- 1,
- )
- })
- }
-}
diff --git a/arrow/src/array/equal_json.rs b/arrow/src/array/equal_json.rs
deleted file mode 100644
index 043174b..0000000
--- a/arrow/src/array/equal_json.rs
+++ /dev/null
@@ -1,1113 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-use super::*;
-use crate::datatypes::*;
-use array::Array;
-use hex::FromHex;
-use serde_json::value::Value::{Null as JNull, Object, String as JString};
-use serde_json::Value;
-
-/// Trait for comparing arrow array with json array
-pub trait JsonEqual {
- /// Checks whether arrow array equals to json array.
- fn equals_json(&self, json: &[&Value]) -> bool;
-
- /// Checks whether arrow array equals to json array.
- fn equals_json_values(&self, json: &[Value]) -> bool {
- let refs = json.iter().collect::<Vec<&Value>>();
-
- self.equals_json(&refs)
- }
-}
-
-/// Implement array equals for numeric type
-impl<T: ArrowPrimitiveType> JsonEqual for PrimitiveArray<T> {
- fn equals_json(&self, json: &[&Value]) -> bool {
- self.len() == json.len()
- && (0..self.len()).all(|i| match json[i] {
- Value::Null => self.is_null(i),
- v => {
- self.is_valid(i)
- && Some(v) == self.value(i).into_json_value().as_ref()
- }
- })
- }
-}
-
-/// Implement array equals for numeric type
-impl JsonEqual for BooleanArray {
- fn equals_json(&self, json: &[&Value]) -> bool {
- self.len() == json.len()
- && (0..self.len()).all(|i| match json[i] {
- Value::Null => self.is_null(i),
- v => {
- self.is_valid(i)
- && Some(v) == self.value(i).into_json_value().as_ref()
- }
- })
- }
-}
-
-impl<T: ArrowPrimitiveType> PartialEq<Value> for PrimitiveArray<T> {
- fn eq(&self, json: &Value) -> bool {
- match json {
- Value::Array(array) => self.equals_json_values(&array),
- _ => false,
- }
- }
-}
-
-impl<T: ArrowPrimitiveType> PartialEq<PrimitiveArray<T>> for Value {
- fn eq(&self, arrow: &PrimitiveArray<T>) -> bool {
- match self {
- Value::Array(array) => arrow.equals_json_values(&array),
- _ => false,
- }
- }
-}
-
-impl<OffsetSize: OffsetSizeTrait> JsonEqual for GenericListArray<OffsetSize> {
- fn equals_json(&self, json: &[&Value]) -> bool {
- if self.len() != json.len() {
- return false;
- }
-
- (0..self.len()).all(|i| match json[i] {
- Value::Array(v) => self.is_valid(i) && self.value(i).equals_json_values(v),
- Value::Null => self.is_null(i) || self.value_length(i).is_zero(),
- _ => false,
- })
- }
-}
-
-impl<OffsetSize: OffsetSizeTrait> PartialEq<Value> for GenericListArray<OffsetSize> {
- fn eq(&self, json: &Value) -> bool {
- match json {
- Value::Array(json_array) => self.equals_json_values(json_array),
- _ => false,
- }
- }
-}
-
-impl<OffsetSize: OffsetSizeTrait> PartialEq<GenericListArray<OffsetSize>> for Value {
- fn eq(&self, arrow: &GenericListArray<OffsetSize>) -> bool {
- match self {
- Value::Array(json_array) => arrow.equals_json_values(json_array),
- _ => false,
- }
- }
-}
-
-impl<T: ArrowPrimitiveType> JsonEqual for DictionaryArray<T> {
- fn equals_json(&self, json: &[&Value]) -> bool {
- // todo: this is wrong: we must test the values also
- self.keys().equals_json(json)
- }
-}
-
-impl<T: ArrowPrimitiveType> PartialEq<Value> for DictionaryArray<T> {
- fn eq(&self, json: &Value) -> bool {
- match json {
- Value::Array(json_array) => self.equals_json_values(json_array),
- _ => false,
- }
- }
-}
-
-impl<T: ArrowPrimitiveType> PartialEq<DictionaryArray<T>> for Value {
- fn eq(&self, arrow: &DictionaryArray<T>) -> bool {
- match self {
- Value::Array(json_array) => arrow.equals_json_values(json_array),
- _ => false,
- }
- }
-}
-
-impl JsonEqual for FixedSizeListArray {
- fn equals_json(&self, json: &[&Value]) -> bool {
- if self.len() != json.len() {
- return false;
- }
-
- (0..self.len()).all(|i| match json[i] {
- Value::Array(v) => self.is_valid(i) && self.value(i).equals_json_values(v),
- Value::Null => self.is_null(i) || self.value_length() == 0,
- _ => false,
- })
- }
-}
-
-impl PartialEq<Value> for FixedSizeListArray {
- fn eq(&self, json: &Value) -> bool {
- match json {
- Value::Array(json_array) => self.equals_json_values(json_array),
- _ => false,
- }
- }
-}
-
-impl PartialEq<FixedSizeListArray> for Value {
- fn eq(&self, arrow: &FixedSizeListArray) -> bool {
- match self {
- Value::Array(json_array) => arrow.equals_json_values(json_array),
- _ => false,
- }
- }
-}
-
-impl JsonEqual for StructArray {
- fn equals_json(&self, json: &[&Value]) -> bool {
- if self.len() != json.len() {
- return false;
- }
-
- let all_object = json.iter().all(|v| matches!(v, Object(_) | JNull));
-
- if !all_object {
- return false;
- }
-
- for column_name in self.column_names() {
- let json_values = json
- .iter()
- .map(|obj| obj.get(column_name).unwrap_or(&Value::Null))
- .collect::<Vec<&Value>>();
-
- if !self
- .column_by_name(column_name)
- .map(|arr| arr.equals_json(&json_values))
- .unwrap_or(false)
- {
- return false;
- }
- }
-
- true
- }
-}
-
-impl PartialEq<Value> for StructArray {
- fn eq(&self, json: &Value) -> bool {
- match json {
- Value::Array(json_array) => self.equals_json_values(&json_array),
- _ => false,
- }
- }
-}
-
-impl PartialEq<StructArray> for Value {
- fn eq(&self, arrow: &StructArray) -> bool {
- match self {
- Value::Array(json_array) => arrow.equals_json_values(&json_array),
- _ => false,
- }
- }
-}
-
-impl<OffsetSize: BinaryOffsetSizeTrait> JsonEqual for GenericBinaryArray<OffsetSize> {
- fn equals_json(&self, json: &[&Value]) -> bool {
- if self.len() != json.len() {
- return false;
- }
-
- (0..self.len()).all(|i| match json[i] {
- JString(s) => {
- // binary data is sometimes hex encoded, this checks if bytes are equal,
- // and if not converting to hex is attempted
- self.is_valid(i)
- && (s.as_str().as_bytes() == self.value(i)
- || Vec::from_hex(s.as_str()) == Ok(self.value(i).to_vec()))
- }
- JNull => self.is_null(i),
- _ => false,
- })
- }
-}
-
-impl<OffsetSize: BinaryOffsetSizeTrait> PartialEq<Value>
- for GenericBinaryArray<OffsetSize>
-{
- fn eq(&self, json: &Value) -> bool {
- match json {
- Value::Array(json_array) => self.equals_json_values(&json_array),
- _ => false,
- }
- }
-}
-
-impl<OffsetSize: BinaryOffsetSizeTrait> PartialEq<GenericBinaryArray<OffsetSize>>
- for Value
-{
- fn eq(&self, arrow: &GenericBinaryArray<OffsetSize>) -> bool {
- match self {
- Value::Array(json_array) => arrow.equals_json_values(&json_array),
- _ => false,
- }
- }
-}
-
-impl<OffsetSize: StringOffsetSizeTrait> JsonEqual for GenericStringArray<OffsetSize> {
- fn equals_json(&self, json: &[&Value]) -> bool {
- if self.len() != json.len() {
- return false;
- }
-
- (0..self.len()).all(|i| match json[i] {
- JString(s) => self.is_valid(i) && s.as_str() == self.value(i),
- JNull => self.is_null(i),
- _ => false,
- })
- }
-}
-
-impl<OffsetSize: StringOffsetSizeTrait> PartialEq<Value>
- for GenericStringArray<OffsetSize>
-{
- fn eq(&self, json: &Value) -> bool {
- match json {
- Value::Array(json_array) => self.equals_json_values(&json_array),
- _ => false,
- }
- }
-}
-
-impl<OffsetSize: StringOffsetSizeTrait> PartialEq<GenericStringArray<OffsetSize>>
- for Value
-{
- fn eq(&self, arrow: &GenericStringArray<OffsetSize>) -> bool {
- match self {
- Value::Array(json_array) => arrow.equals_json_values(&json_array),
- _ => false,
- }
- }
-}
-
-impl JsonEqual for FixedSizeBinaryArray {
- fn equals_json(&self, json: &[&Value]) -> bool {
- if self.len() != json.len() {
- return false;
- }
-
- (0..self.len()).all(|i| match json[i] {
- JString(s) => {
- // binary data is sometimes hex encoded, this checks if bytes are equal,
- // and if not converting to hex is attempted
- self.is_valid(i)
- && (s.as_str().as_bytes() == self.value(i)
- || Vec::from_hex(s.as_str()) == Ok(self.value(i).to_vec()))
- }
- JNull => self.is_null(i),
- _ => false,
- })
- }
-}
-
-impl PartialEq<Value> for FixedSizeBinaryArray {
- fn eq(&self, json: &Value) -> bool {
- match json {
- Value::Array(json_array) => self.equals_json_values(&json_array),
- _ => false,
- }
- }
-}
-
-impl PartialEq<FixedSizeBinaryArray> for Value {
- fn eq(&self, arrow: &FixedSizeBinaryArray) -> bool {
- match self {
- Value::Array(json_array) => arrow.equals_json_values(&json_array),
- _ => false,
- }
- }
-}
-
-impl JsonEqual for DecimalArray {
- fn equals_json(&self, json: &[&Value]) -> bool {
- if self.len() != json.len() {
- return false;
- }
-
- (0..self.len()).all(|i| match json[i] {
- JString(s) => {
- self.is_valid(i)
- && (s
- .parse::<i128>()
- .map_or_else(|_| false, |v| v == self.value(i)))
- }
- JNull => self.is_null(i),
- _ => false,
- })
- }
-}
-
-impl PartialEq<Value> for DecimalArray {
- fn eq(&self, json: &Value) -> bool {
- match json {
- Value::Array(json_array) => self.equals_json_values(&json_array),
- _ => false,
- }
- }
-}
-
-impl PartialEq<DecimalArray> for Value {
- fn eq(&self, arrow: &DecimalArray) -> bool {
- match self {
- Value::Array(json_array) => arrow.equals_json_values(&json_array),
- _ => false,
- }
- }
-}
-
-impl JsonEqual for UnionArray {
- fn equals_json(&self, _json: &[&Value]) -> bool {
- unimplemented!(
- "Added to allow UnionArray to implement the Array trait: see ARROW-8547"
- )
- }
-}
-
-impl JsonEqual for NullArray {
- fn equals_json(&self, json: &[&Value]) -> bool {
- if self.len() != json.len() {
- return false;
- }
-
- // all JSON values must be nulls
- json.iter().all(|&v| v == &JNull)
- }
-}
-
-impl PartialEq<NullArray> for Value {
- fn eq(&self, arrow: &NullArray) -> bool {
- match self {
- Value::Array(json_array) => arrow.equals_json_values(&json_array),
- _ => false,
- }
- }
-}
-
-impl PartialEq<Value> for NullArray {
- fn eq(&self, json: &Value) -> bool {
- match json {
- Value::Array(json_array) => self.equals_json_values(&json_array),
- _ => false,
- }
- }
-}
-
-#[cfg(test)]
-mod tests {
- use super::*;
-
- use crate::error::Result;
- use std::{convert::TryFrom, sync::Arc};
-
- fn create_list_array<U: AsRef<[i32]>, T: AsRef<[Option<U>]>>(
- builder: &mut ListBuilder<Int32Builder>,
- data: T,
- ) -> Result<ListArray> {
- for d in data.as_ref() {
- if let Some(v) = d {
- builder.values().append_slice(v.as_ref())?;
- builder.append(true)?
- } else {
- builder.append(false)?
- }
- }
- Ok(builder.finish())
- }
-
- /// Create a fixed size list of 2 value lengths
- fn create_fixed_size_list_array<U: AsRef<[i32]>, T: AsRef<[Option<U>]>>(
- builder: &mut FixedSizeListBuilder<Int32Builder>,
- data: T,
- ) -> Result<FixedSizeListArray> {
- for d in data.as_ref() {
- if let Some(v) = d {
- builder.values().append_slice(v.as_ref())?;
- builder.append(true)?
- } else {
- for _ in 0..builder.value_length() {
- builder.values().append_null()?;
- }
- builder.append(false)?
- }
- }
- Ok(builder.finish())
- }
-
- #[test]
- fn test_primitive_json_equal() {
- // Test equaled array
- let arrow_array = Int32Array::from(vec![Some(1), None, Some(2), Some(3)]);
- let json_array: Value = serde_json::from_str(
- r#"
- [
- 1, null, 2, 3
- ]
- "#,
- )
- .unwrap();
- assert!(arrow_array.eq(&json_array));
- assert!(json_array.eq(&arrow_array));
-
- // Test unequaled array
- let arrow_array = Int32Array::from(vec![Some(1), None, Some(2), Some(3)]);
- let json_array: Value = serde_json::from_str(
- r#"
- [
- 1, 1, 2, 3
- ]
- "#,
- )
- .unwrap();
- assert!(arrow_array.ne(&json_array));
- assert!(json_array.ne(&arrow_array));
-
- // Test unequal length case
- let arrow_array = Int32Array::from(vec![Some(1), None, Some(2), Some(3)]);
- let json_array: Value = serde_json::from_str(
- r#"
- [
- 1, 1
- ]
- "#,
- )
- .unwrap();
- assert!(arrow_array.ne(&json_array));
- assert!(json_array.ne(&arrow_array));
-
- // Test not json array type case
- let arrow_array = Int32Array::from(vec![Some(1), None, Some(2), Some(3)]);
- let json_array: Value = serde_json::from_str(
- r#"
- {
- "a": 1
- }
- "#,
- )
- .unwrap();
- assert!(arrow_array.ne(&json_array));
- assert!(json_array.ne(&arrow_array));
- }
-
- #[test]
- fn test_list_json_equal() {
- // Test equal case
- let arrow_array = create_list_array(
- &mut ListBuilder::new(Int32Builder::new(10)),
- &[Some(&[1, 2, 3]), None, Some(&[4, 5, 6])],
- )
- .unwrap();
- let json_array: Value = serde_json::from_str(
- r#"
- [
- [1, 2, 3],
- null,
- [4, 5, 6]
- ]
- "#,
- )
- .unwrap();
- assert!(arrow_array.eq(&json_array));
- assert!(json_array.eq(&arrow_array));
-
- // Test unequal case
- let arrow_array = create_list_array(
- &mut ListBuilder::new(Int32Builder::new(10)),
- &[Some(&[1, 2, 3]), None, Some(&[4, 5, 6])],
- )
- .unwrap();
- let json_array: Value = serde_json::from_str(
- r#"
- [
- [1, 2, 3],
- [7, 8],
- [4, 5, 6]
- ]
- "#,
- )
- .unwrap();
- assert!(arrow_array.ne(&json_array));
- assert!(json_array.ne(&arrow_array));
-
- // Test incorrect type case
- let arrow_array = create_list_array(
- &mut ListBuilder::new(Int32Builder::new(10)),
- &[Some(&[1, 2, 3]), None, Some(&[4, 5, 6])],
- )
- .unwrap();
- let json_array: Value = serde_json::from_str(
- r#"
- {
- "a": 1
- }
- "#,
- )
- .unwrap();
- assert!(arrow_array.ne(&json_array));
- assert!(json_array.ne(&arrow_array));
- }
-
- #[test]
- fn test_fixed_size_list_json_equal() {
- // Test equal case
- let arrow_array = create_fixed_size_list_array(
- &mut FixedSizeListBuilder::new(Int32Builder::new(10), 3),
- &[Some(&[1, 2, 3]), None, Some(&[4, 5, 6])],
- )
- .unwrap();
- let json_array: Value = serde_json::from_str(
- r#"
- [
- [1, 2, 3],
- null,
- [4, 5, 6]
- ]
- "#,
- )
- .unwrap();
- println!("{:?}", arrow_array);
- println!("{:?}", json_array);
- assert!(arrow_array.eq(&json_array));
- assert!(json_array.eq(&arrow_array));
-
- // Test unequal case
- let arrow_array = create_fixed_size_list_array(
- &mut FixedSizeListBuilder::new(Int32Builder::new(10), 3),
- &[Some(&[1, 2, 3]), None, Some(&[4, 5, 6])],
- )
- .unwrap();
- let json_array: Value = serde_json::from_str(
- r#"
- [
- [1, 2, 3],
- [7, 8, 9],
- [4, 5, 6]
- ]
- "#,
- )
- .unwrap();
- assert!(arrow_array.ne(&json_array));
- assert!(json_array.ne(&arrow_array));
-
- // Test incorrect type case
- let arrow_array = create_fixed_size_list_array(
- &mut FixedSizeListBuilder::new(Int32Builder::new(10), 3),
- &[Some(&[1, 2, 3]), None, Some(&[4, 5, 6])],
- )
- .unwrap();
- let json_array: Value = serde_json::from_str(
- r#"
- {
- "a": 1
- }
- "#,
- )
- .unwrap();
- assert!(arrow_array.ne(&json_array));
- assert!(json_array.ne(&arrow_array));
- }
-
- #[test]
- fn test_string_json_equal() {
- // Test the equal case
- let arrow_array =
- StringArray::from(vec![Some("hello"), None, None, Some("world"), None, None]);
- let json_array: Value = serde_json::from_str(
- r#"
- [
- "hello",
- null,
- null,
- "world",
- null,
- null
- ]
- "#,
- )
- .unwrap();
- assert!(arrow_array.eq(&json_array));
- assert!(json_array.eq(&arrow_array));
-
- // Test unequal case
- let arrow_array =
- StringArray::from(vec![Some("hello"), None, None, Some("world"), None, None]);
- let json_array: Value = serde_json::from_str(
- r#"
- [
- "hello",
- null,
- null,
- "arrow",
- null,
- null
- ]
- "#,
- )
- .unwrap();
- assert!(arrow_array.ne(&json_array));
- assert!(json_array.ne(&arrow_array));
-
- // Test unequal length case
- let arrow_array =
- StringArray::from(vec![Some("hello"), None, None, Some("world"), None]);
- let json_array: Value = serde_json::from_str(
- r#"
- [
- "hello",
- null,
- null,
- "arrow",
- null,
- null
- ]
- "#,
- )
- .unwrap();
- assert!(arrow_array.ne(&json_array));
- assert!(json_array.ne(&arrow_array));
-
- // Test incorrect type case
- let arrow_array =
- StringArray::from(vec![Some("hello"), None, None, Some("world"), None]);
- let json_array: Value = serde_json::from_str(
- r#"
- {
- "a": 1
- }
- "#,
- )
- .unwrap();
- assert!(arrow_array.ne(&json_array));
- assert!(json_array.ne(&arrow_array));
-
- // Test incorrect value type case
- let arrow_array =
- StringArray::from(vec![Some("hello"), None, None, Some("world"), None]);
- let json_array: Value = serde_json::from_str(
- r#"
- [
- "hello",
- null,
- null,
- 1,
- null,
- null
- ]
- "#,
- )
- .unwrap();
- assert!(arrow_array.ne(&json_array));
- assert!(json_array.ne(&arrow_array));
- }
-
- #[test]
- fn test_binary_json_equal() {
- // Test the equal case
- let mut builder = BinaryBuilder::new(6);
- builder.append_value(b"hello").unwrap();
- builder.append_null().unwrap();
- builder.append_null().unwrap();
- builder.append_value(b"world").unwrap();
- builder.append_null().unwrap();
- builder.append_null().unwrap();
- let arrow_array = builder.finish();
- let json_array: Value = serde_json::from_str(
- r#"
- [
- "hello",
- null,
- null,
- "world",
- null,
- null
- ]
- "#,
- )
- .unwrap();
- assert!(arrow_array.eq(&json_array));
- assert!(json_array.eq(&arrow_array));
-
- // Test unequal case
- let arrow_array =
- StringArray::from(vec![Some("hello"), None, None, Some("world"), None, None]);
- let json_array: Value = serde_json::from_str(
- r#"
- [
- "hello",
- null,
- null,
- "arrow",
- null,
- null
- ]
- "#,
- )
- .unwrap();
- assert!(arrow_array.ne(&json_array));
- assert!(json_array.ne(&arrow_array));
-
- // Test unequal length case
- let arrow_array =
- StringArray::from(vec![Some("hello"), None, None, Some("world"), None]);
- let json_array: Value = serde_json::from_str(
- r#"
- [
- "hello",
- null,
- null,
- "arrow",
- null,
- null
- ]
- "#,
- )
- .unwrap();
- assert!(arrow_array.ne(&json_array));
- assert!(json_array.ne(&arrow_array));
-
- // Test incorrect type case
- let arrow_array =
- StringArray::from(vec![Some("hello"), None, None, Some("world"), None]);
- let json_array: Value = serde_json::from_str(
- r#"
- {
- "a": 1
- }
- "#,
- )
- .unwrap();
- assert!(arrow_array.ne(&json_array));
- assert!(json_array.ne(&arrow_array));
-
- // Test incorrect value type case
- let arrow_array =
- StringArray::from(vec![Some("hello"), None, None, Some("world"), None]);
- let json_array: Value = serde_json::from_str(
- r#"
- [
- "hello",
- null,
- null,
- 1,
- null,
- null
- ]
- "#,
- )
- .unwrap();
- assert!(arrow_array.ne(&json_array));
- assert!(json_array.ne(&arrow_array));
- }
-
- #[test]
- fn test_fixed_size_binary_json_equal() {
- // Test the equal case
- let mut builder = FixedSizeBinaryBuilder::new(15, 5);
- builder.append_value(b"hello").unwrap();
- builder.append_null().unwrap();
- builder.append_value(b"world").unwrap();
- let arrow_array: FixedSizeBinaryArray = builder.finish();
- let json_array: Value = serde_json::from_str(
- r#"
- [
- "hello",
- null,
- "world"
- ]
- "#,
- )
- .unwrap();
- assert!(arrow_array.eq(&json_array));
- assert!(json_array.eq(&arrow_array));
-
- // Test unequal case
- builder.append_value(b"hello").unwrap();
- builder.append_null().unwrap();
- builder.append_value(b"world").unwrap();
- let arrow_array: FixedSizeBinaryArray = builder.finish();
- let json_array: Value = serde_json::from_str(
- r#"
- [
- "hello",
- null,
- "arrow"
- ]
- "#,
- )
- .unwrap();
- assert!(arrow_array.ne(&json_array));
- assert!(json_array.ne(&arrow_array));
-
- // Test unequal length case
- let json_array: Value = serde_json::from_str(
- r#"
- [
- "hello",
- null,
- null,
- "world"
- ]
- "#,
- )
- .unwrap();
- assert!(arrow_array.ne(&json_array));
- assert!(json_array.ne(&arrow_array));
-
- // Test incorrect type case
- let json_array: Value = serde_json::from_str(
- r#"
- {
- "a": 1
- }
- "#,
- )
- .unwrap();
- assert!(arrow_array.ne(&json_array));
- assert!(json_array.ne(&arrow_array));
-
- // Test incorrect value type case
- let json_array: Value = serde_json::from_str(
- r#"
- [
- "hello",
- null,
- 1
- ]
- "#,
- )
- .unwrap();
- assert!(arrow_array.ne(&json_array));
- assert!(json_array.ne(&arrow_array));
- }
-
- #[test]
- fn test_decimal_json_equal() {
- // Test the equal case
- let mut builder = DecimalBuilder::new(30, 23, 6);
- builder.append_value(1_000).unwrap();
- builder.append_null().unwrap();
- builder.append_value(-250).unwrap();
- let arrow_array: DecimalArray = builder.finish();
- let json_array: Value = serde_json::from_str(
- r#"
- [
- "1000",
- null,
- "-250"
- ]
- "#,
- )
- .unwrap();
- assert!(arrow_array.eq(&json_array));
- assert!(json_array.eq(&arrow_array));
-
- // Test unequal case
- builder.append_value(1_000).unwrap();
- builder.append_null().unwrap();
- builder.append_value(55).unwrap();
- let arrow_array: DecimalArray = builder.finish();
- let json_array: Value = serde_json::from_str(
- r#"
- [
- "1000",
- null,
- "-250"
- ]
- "#,
- )
- .unwrap();
- assert!(arrow_array.ne(&json_array));
- assert!(json_array.ne(&arrow_array));
-
- // Test unequal length case
- let json_array: Value = serde_json::from_str(
- r#"
- [
- "1000",
- null,
- null,
- "55"
- ]
- "#,
- )
- .unwrap();
- assert!(arrow_array.ne(&json_array));
- assert!(json_array.ne(&arrow_array));
-
- // Test incorrect type case
- let json_array: Value = serde_json::from_str(
- r#"
- {
- "a": 1
- }
- "#,
- )
- .unwrap();
- assert!(arrow_array.ne(&json_array));
- assert!(json_array.ne(&arrow_array));
-
- // Test incorrect value type case
- let json_array: Value = serde_json::from_str(
- r#"
- [
- "hello",
- null,
- 1
- ]
- "#,
- )
- .unwrap();
- assert!(arrow_array.ne(&json_array));
- assert!(json_array.ne(&arrow_array));
- }
-
- #[test]
- fn test_struct_json_equal() {
- let strings: ArrayRef = Arc::new(StringArray::from(vec![
- Some("joe"),
- None,
- None,
- Some("mark"),
- Some("doe"),
- ]));
- let ints: ArrayRef = Arc::new(Int32Array::from(vec![
- Some(1),
- Some(2),
- None,
- Some(4),
- Some(5),
- ]));
-
- let arrow_array =
- StructArray::try_from(vec![("f1", strings.clone()), ("f2", ints.clone())])
- .unwrap();
-
- let json_array: Value = serde_json::from_str(
- r#"
- [
- {
- "f1": "joe",
- "f2": 1
- },
- {
- "f2": 2
- },
- null,
- {
- "f1": "mark",
- "f2": 4
- },
- {
- "f1": "doe",
- "f2": 5
- }
- ]
- "#,
- )
- .unwrap();
- assert!(arrow_array.eq(&json_array));
- assert!(json_array.eq(&arrow_array));
-
- // Test unequal length case
- let json_array: Value = serde_json::from_str(
- r#"
- [
- {
- "f1": "joe",
- "f2": 1
- },
- {
- "f2": 2
- },
- null,
- {
- "f1": "mark",
- "f2": 4
- }
- ]
- "#,
- )
- .unwrap();
- assert!(arrow_array.ne(&json_array));
- assert!(json_array.ne(&arrow_array));
-
- // Test incorrect type case
- let json_array: Value = serde_json::from_str(
- r#"
- {
- "f1": "joe",
- "f2": 1
- }
- "#,
- )
- .unwrap();
- assert!(arrow_array.ne(&json_array));
- assert!(json_array.ne(&arrow_array));
-
- // Test not all object case
- let json_array: Value = serde_json::from_str(
- r#"
- [
- {
- "f1": "joe",
- "f2": 1
- },
- 2,
- null,
- {
- "f1": "mark",
- "f2": 4
- }
- ]
- "#,
- )
- .unwrap();
- assert!(arrow_array.ne(&json_array));
- assert!(json_array.ne(&arrow_array));
- }
-
- #[test]
- fn test_null_json_equal() {
- // Test equaled array
- let arrow_array = NullArray::new(4);
- let json_array: Value = serde_json::from_str(
- r#"
- [
- null, null, null, null
- ]
- "#,
- )
- .unwrap();
- assert!(arrow_array.eq(&json_array));
- assert!(json_array.eq(&arrow_array));
-
- // Test unequaled array
- let arrow_array = NullArray::new(2);
- let json_array: Value = serde_json::from_str(
- r#"
- [
- null, null, null
- ]
- "#,
- )
- .unwrap();
- assert!(arrow_array.ne(&json_array));
- assert!(json_array.ne(&arrow_array));
- }
-}
diff --git a/arrow/src/array/ffi.rs b/arrow/src/array/ffi.rs
deleted file mode 100644
index 847649c..0000000
--- a/arrow/src/array/ffi.rs
+++ /dev/null
@@ -1,130 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-//! Contains functionality to load an ArrayData from the C Data Interface
-
-use std::convert::TryFrom;
-
-use crate::{
- error::{ArrowError, Result},
- ffi,
- ffi::ArrowArrayRef,
-};
-
-use super::ArrayData;
-
-impl TryFrom<ffi::ArrowArray> for ArrayData {
- type Error = ArrowError;
-
- fn try_from(value: ffi::ArrowArray) -> Result<Self> {
- value.to_data()
- }
-}
-
-impl TryFrom<ArrayData> for ffi::ArrowArray {
- type Error = ArrowError;
-
- fn try_from(value: ArrayData) -> Result<Self> {
- unsafe { ffi::ArrowArray::try_new(value) }
- }
-}
-
-#[cfg(test)]
-mod tests {
- use crate::error::Result;
- use crate::{
- array::{
- Array, ArrayData, BooleanArray, Int64Array, StructArray, UInt32Array,
- UInt64Array,
- },
- datatypes::{DataType, Field},
- ffi::ArrowArray,
- };
- use std::convert::TryFrom;
- use std::sync::Arc;
-
- fn test_round_trip(expected: &ArrayData) -> Result<()> {
- // create a `ArrowArray` from the data.
- let d1 = ArrowArray::try_from(expected.clone())?;
-
- // here we export the array as 2 pointers. We would have no control over ownership if it was not for
- // the release mechanism.
- let (array, schema) = ArrowArray::into_raw(d1);
-
- // simulate an external consumer by being the consumer
- let d1 = unsafe { ArrowArray::try_from_raw(array, schema) }?;
-
- let result = &ArrayData::try_from(d1)?;
-
- assert_eq!(result, expected);
- Ok(())
- }
-
- #[test]
- fn test_u32() -> Result<()> {
- let array = UInt32Array::from(vec![Some(2), None, Some(1), None]);
- let data = array.data();
- test_round_trip(data)
- }
-
- #[test]
- fn test_u64() -> Result<()> {
- let array = UInt64Array::from(vec![Some(2), None, Some(1), None]);
- let data = array.data();
- test_round_trip(data)
- }
-
- #[test]
- fn test_i64() -> Result<()> {
- let array = Int64Array::from(vec![Some(2), None, Some(1), None]);
- let data = array.data();
- test_round_trip(data)
- }
-
- #[test]
- fn test_struct() -> Result<()> {
- let inner = StructArray::from(vec![
- (
- Field::new("a1", DataType::Boolean, false),
- Arc::new(BooleanArray::from(vec![true, true, false, false]))
- as Arc<dyn Array>,
- ),
- (
- Field::new("a2", DataType::UInt32, false),
- Arc::new(UInt32Array::from(vec![1, 2, 3, 4])),
- ),
- ]);
-
- let array = StructArray::from(vec![
- (
- Field::new("a", inner.data_type().clone(), false),
- Arc::new(inner) as Arc<dyn Array>,
- ),
- (
- Field::new("b", DataType::Boolean, false),
- Arc::new(BooleanArray::from(vec![false, false, true, true]))
- as Arc<dyn Array>,
- ),
- (
- Field::new("c", DataType::UInt32, false),
- Arc::new(UInt32Array::from(vec![42, 28, 19, 31])),
- ),
- ]);
- let data = array.data();
- test_round_trip(data)
- }
-}
diff --git a/arrow/src/array/iterator.rs b/arrow/src/array/iterator.rs
deleted file mode 100644
index d97aa16..0000000
--- a/arrow/src/array/iterator.rs
+++ /dev/null
@@ -1,527 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-use crate::datatypes::ArrowPrimitiveType;
-
-use super::{
- Array, ArrayRef, BinaryOffsetSizeTrait, BooleanArray, GenericBinaryArray,
- GenericListArray, GenericStringArray, OffsetSizeTrait, PrimitiveArray,
- StringOffsetSizeTrait,
-};
-
-/// an iterator that returns Some(T) or None, that can be used on any PrimitiveArray
-// Note: This implementation is based on std's [Vec]s' [IntoIter].
-#[derive(Debug)]
-pub struct PrimitiveIter<'a, T: ArrowPrimitiveType> {
- array: &'a PrimitiveArray<T>,
- current: usize,
- current_end: usize,
-}
-
-impl<'a, T: ArrowPrimitiveType> PrimitiveIter<'a, T> {
- /// create a new iterator
- pub fn new(array: &'a PrimitiveArray<T>) -> Self {
- PrimitiveIter::<T> {
- array,
- current: 0,
- current_end: array.len(),
- }
- }
-}
-
-impl<'a, T: ArrowPrimitiveType> std::iter::Iterator for PrimitiveIter<'a, T> {
- type Item = Option<T::Native>;
-
- #[inline]
- fn next(&mut self) -> Option<Self::Item> {
- if self.current == self.current_end {
- None
- } else if self.array.is_null(self.current) {
- self.current += 1;
- Some(None)
- } else {
- let old = self.current;
- self.current += 1;
- // Safety:
- // we just checked bounds in `self.current_end == self.current`
- // this is safe on the premise that this struct is initialized with
- // current = array.len()
- // and that current_end is ever only decremented
- unsafe { Some(Some(self.array.value_unchecked(old))) }
- }
- }
-
- fn size_hint(&self) -> (usize, Option<usize>) {
- (
- self.array.len() - self.current,
- Some(self.array.len() - self.current),
- )
- }
-}
-
-impl<'a, T: ArrowPrimitiveType> std::iter::DoubleEndedIterator for PrimitiveIter<'a, T> {
- fn next_back(&mut self) -> Option<Self::Item> {
- if self.current_end == self.current {
- None
- } else {
- self.current_end -= 1;
- Some(if self.array.is_null(self.current_end) {
- None
- } else {
- // Safety:
- // we just checked bounds in `self.current_end == self.current`
- // this is safe on the premise that this struct is initialized with
- // current = array.len()
- // and that current_end is ever only decremented
- unsafe { Some(self.array.value_unchecked(self.current_end)) }
- })
- }
- }
-}
-
-/// all arrays have known size.
-impl<'a, T: ArrowPrimitiveType> std::iter::ExactSizeIterator for PrimitiveIter<'a, T> {}
-
-/// an iterator that returns Some(bool) or None.
-// Note: This implementation is based on std's [Vec]s' [IntoIter].
-#[derive(Debug)]
-pub struct BooleanIter<'a> {
- array: &'a BooleanArray,
- current: usize,
- current_end: usize,
-}
-
-impl<'a> BooleanIter<'a> {
- /// create a new iterator
- pub fn new(array: &'a BooleanArray) -> Self {
- BooleanIter {
- array,
- current: 0,
- current_end: array.len(),
- }
- }
-}
-
-impl<'a> std::iter::Iterator for BooleanIter<'a> {
- type Item = Option<bool>;
-
- fn next(&mut self) -> Option<Self::Item> {
- if self.current == self.current_end {
- None
- } else if self.array.is_null(self.current) {
- self.current += 1;
- Some(None)
- } else {
- let old = self.current;
- self.current += 1;
- // Safety:
- // we just checked bounds in `self.current_end == self.current`
- // this is safe on the premise that this struct is initialized with
- // current = array.len()
- // and that current_end is ever only decremented
- unsafe { Some(Some(self.array.value_unchecked(old))) }
- }
- }
-
- fn size_hint(&self) -> (usize, Option<usize>) {
- (
- self.array.len() - self.current,
- Some(self.array.len() - self.current),
- )
- }
-}
-
-impl<'a> std::iter::DoubleEndedIterator for BooleanIter<'a> {
- fn next_back(&mut self) -> Option<Self::Item> {
- if self.current_end == self.current {
- None
- } else {
- self.current_end -= 1;
- Some(if self.array.is_null(self.current_end) {
- None
- } else {
- // Safety:
- // we just checked bounds in `self.current_end == self.current`
- // this is safe on the premise that this struct is initialized with
- // current = array.len()
- // and that current_end is ever only decremented
- unsafe { Some(self.array.value_unchecked(self.current_end)) }
- })
- }
- }
-}
-
-/// all arrays have known size.
-impl<'a> std::iter::ExactSizeIterator for BooleanIter<'a> {}
-
-/// an iterator that returns `Some(&str)` or `None`, for string arrays
-#[derive(Debug)]
-pub struct GenericStringIter<'a, T>
-where
- T: StringOffsetSizeTrait,
-{
- array: &'a GenericStringArray<T>,
- current: usize,
- current_end: usize,
-}
-
-impl<'a, T: StringOffsetSizeTrait> GenericStringIter<'a, T> {
- /// create a new iterator
- pub fn new(array: &'a GenericStringArray<T>) -> Self {
- GenericStringIter::<T> {
- array,
- current: 0,
- current_end: array.len(),
- }
- }
-}
-
-impl<'a, T: StringOffsetSizeTrait> std::iter::Iterator for GenericStringIter<'a, T> {
- type Item = Option<&'a str>;
-
- fn next(&mut self) -> Option<Self::Item> {
- let i = self.current;
- if i >= self.current_end {
- None
- } else if self.array.is_null(i) {
- self.current += 1;
- Some(None)
- } else {
- self.current += 1;
- // Safety:
- // we just checked bounds in `self.current_end == self.current`
- // this is safe on the premise that this struct is initialized with
- // current = array.len()
- // and that current_end is ever only decremented
- unsafe { Some(Some(self.array.value_unchecked(i))) }
- }
- }
-
- fn size_hint(&self) -> (usize, Option<usize>) {
- (
- self.current_end - self.current,
- Some(self.current_end - self.current),
- )
- }
-}
-
-impl<'a, T: StringOffsetSizeTrait> std::iter::DoubleEndedIterator
- for GenericStringIter<'a, T>
-{
- fn next_back(&mut self) -> Option<Self::Item> {
- if self.current_end == self.current {
- None
- } else {
- self.current_end -= 1;
- Some(if self.array.is_null(self.current_end) {
- None
- } else {
- // Safety:
- // we just checked bounds in `self.current_end == self.current`
- // this is safe on the premise that this struct is initialized with
- // current = array.len()
- // and that current_end is ever only decremented
- unsafe { Some(self.array.value_unchecked(self.current_end)) }
- })
- }
- }
-}
-
-/// all arrays have known size.
-impl<'a, T: StringOffsetSizeTrait> std::iter::ExactSizeIterator
- for GenericStringIter<'a, T>
-{
-}
-
-/// an iterator that returns `Some(&[u8])` or `None`, for binary arrays
-#[derive(Debug)]
-pub struct GenericBinaryIter<'a, T>
-where
- T: BinaryOffsetSizeTrait,
-{
- array: &'a GenericBinaryArray<T>,
- current: usize,
- current_end: usize,
-}
-
-impl<'a, T: BinaryOffsetSizeTrait> GenericBinaryIter<'a, T> {
- /// create a new iterator
- pub fn new(array: &'a GenericBinaryArray<T>) -> Self {
- GenericBinaryIter::<T> {
- array,
- current: 0,
- current_end: array.len(),
- }
- }
-}
-
-impl<'a, T: BinaryOffsetSizeTrait> std::iter::Iterator for GenericBinaryIter<'a, T> {
- type Item = Option<&'a [u8]>;
-
- fn next(&mut self) -> Option<Self::Item> {
- let i = self.current;
- if i >= self.current_end {
- None
- } else if self.array.is_null(i) {
- self.current += 1;
- Some(None)
- } else {
- self.current += 1;
- // Safety:
- // we just checked bounds in `self.current_end == self.current`
- // this is safe on the premise that this struct is initialized with
- // current = array.len()
- // and that current_end is ever only decremented
- unsafe { Some(Some(self.array.value_unchecked(i))) }
- }
- }
-
- fn size_hint(&self) -> (usize, Option<usize>) {
- (
- self.current_end - self.current,
- Some(self.current_end - self.current),
- )
- }
-}
-
-impl<'a, T: BinaryOffsetSizeTrait> std::iter::DoubleEndedIterator
- for GenericBinaryIter<'a, T>
-{
- fn next_back(&mut self) -> Option<Self::Item> {
- if self.current_end == self.current {
- None
- } else {
- self.current_end -= 1;
- Some(if self.array.is_null(self.current_end) {
- None
- } else {
- // Safety:
- // we just checked bounds in `self.current_end == self.current`
- // this is safe on the premise that this struct is initialized with
- // current = array.len()
- // and that current_end is ever only decremented
- unsafe { Some(self.array.value_unchecked(self.current_end)) }
- })
- }
- }
-}
-
-/// all arrays have known size.
-impl<'a, T: BinaryOffsetSizeTrait> std::iter::ExactSizeIterator
- for GenericBinaryIter<'a, T>
-{
-}
-
-#[derive(Debug)]
-pub struct GenericListArrayIter<'a, S>
-where
- S: OffsetSizeTrait,
-{
- array: &'a GenericListArray<S>,
- current: usize,
- current_end: usize,
-}
-
-impl<'a, S: OffsetSizeTrait> GenericListArrayIter<'a, S> {
- pub fn new(array: &'a GenericListArray<S>) -> Self {
- GenericListArrayIter::<S> {
- array,
- current: 0,
- current_end: array.len(),
- }
- }
-}
-
-impl<'a, S: OffsetSizeTrait> std::iter::Iterator for GenericListArrayIter<'a, S> {
- type Item = Option<ArrayRef>;
-
- fn next(&mut self) -> Option<Self::Item> {
- let i = self.current;
- if i >= self.current_end {
- None
- } else if self.array.is_null(i) {
- self.current += 1;
- Some(None)
- } else {
- self.current += 1;
- // Safety:
- // we just checked bounds in `self.current_end == self.current`
- // this is safe on the premise that this struct is initialized with
- // current = array.len()
- // and that current_end is ever only decremented
- unsafe { Some(Some(self.array.value_unchecked(i))) }
- }
- }
-
- fn size_hint(&self) -> (usize, Option<usize>) {
- (
- self.current_end - self.current,
- Some(self.current_end - self.current),
- )
- }
-}
-
-impl<'a, S: OffsetSizeTrait> std::iter::DoubleEndedIterator
- for GenericListArrayIter<'a, S>
-{
- fn next_back(&mut self) -> Option<Self::Item> {
- if self.current_end == self.current {
- None
- } else {
- self.current_end -= 1;
- Some(if self.array.is_null(self.current_end) {
- None
- } else {
- // Safety:
- // we just checked bounds in `self.current_end == self.current`
- // this is safe on the premise that this struct is initialized with
- // current = array.len()
- // and that current_end is ever only decremented
- unsafe { Some(self.array.value_unchecked(self.current_end)) }
- })
- }
- }
-}
-
-/// all arrays have known size.
-impl<'a, S: OffsetSizeTrait> std::iter::ExactSizeIterator
- for GenericListArrayIter<'a, S>
-{
-}
-
-#[cfg(test)]
-mod tests {
- use std::sync::Arc;
-
- use crate::array::{ArrayRef, BinaryArray, BooleanArray, Int32Array, StringArray};
-
- #[test]
- fn test_primitive_array_iter_round_trip() {
- let array = Int32Array::from(vec![Some(0), None, Some(2), None, Some(4)]);
- let array = Arc::new(array) as ArrayRef;
-
- let array = array.as_any().downcast_ref::<Int32Array>().unwrap();
-
- // to and from iter, with a +1
- let result: Int32Array = array.iter().map(|e| e.map(|e| e + 1)).collect();
-
- let expected = Int32Array::from(vec![Some(1), None, Some(3), None, Some(5)]);
- assert_eq!(result, expected);
-
- // check if DoubleEndedIterator is implemented
- let result: Int32Array = array.iter().rev().collect();
- let rev_array = Int32Array::from(vec![Some(4), None, Some(2), None, Some(0)]);
- assert_eq!(result, rev_array);
- // check if ExactSizeIterator is implemented
- let _ = array.iter().rposition(|opt_b| opt_b == Some(1));
- }
-
- #[test]
- fn test_double_ended() {
- let array = Int32Array::from(vec![Some(0), None, Some(2), None, Some(4)]);
- let mut a = array.iter();
- assert_eq!(a.next(), Some(Some(0)));
- assert_eq!(a.next(), Some(None));
- assert_eq!(a.next_back(), Some(Some(4)));
- assert_eq!(a.next_back(), Some(None));
- assert_eq!(a.next_back(), Some(Some(2)));
- // the two sides have met: None is returned by both
- assert_eq!(a.next_back(), None);
- assert_eq!(a.next(), None);
- }
-
- #[test]
- fn test_string_array_iter_round_trip() {
- let array =
- StringArray::from(vec![Some("a"), None, Some("aaa"), None, Some("aaaaa")]);
- let array = Arc::new(array) as ArrayRef;
-
- let array = array.as_any().downcast_ref::<StringArray>().unwrap();
-
- // to and from iter, with a +1
- let result: StringArray = array
- .iter()
- .map(|e| {
- e.map(|e| {
- let mut a = e.to_string();
- a.push('b');
- a
- })
- })
- .collect();
-
- let expected =
- StringArray::from(vec![Some("ab"), None, Some("aaab"), None, Some("aaaaab")]);
- assert_eq!(result, expected);
-
- // check if DoubleEndedIterator is implemented
- let result: StringArray = array.iter().rev().collect();
- let rev_array =
- StringArray::from(vec![Some("aaaaa"), None, Some("aaa"), None, Some("a")]);
- assert_eq!(result, rev_array);
- // check if ExactSizeIterator is implemented
- let _ = array.iter().rposition(|opt_b| opt_b == Some("a"));
- }
-
- #[test]
- fn test_binary_array_iter_round_trip() {
- let array = BinaryArray::from(vec![
- Some(b"a" as &[u8]),
- None,
- Some(b"aaa"),
- None,
- Some(b"aaaaa"),
- ]);
-
- // to and from iter
- let result: BinaryArray = array.iter().collect();
-
- assert_eq!(result, array);
-
- // check if DoubleEndedIterator is implemented
- let result: BinaryArray = array.iter().rev().collect();
- let rev_array = BinaryArray::from(vec![
- Some(b"aaaaa" as &[u8]),
- None,
- Some(b"aaa"),
- None,
- Some(b"a"),
- ]);
- assert_eq!(result, rev_array);
-
- // check if ExactSizeIterator is implemented
- let _ = array.iter().rposition(|opt_b| opt_b == Some(&[9]));
- }
-
- #[test]
- fn test_boolean_array_iter_round_trip() {
- let array = BooleanArray::from(vec![Some(true), None, Some(false)]);
-
- // to and from iter
- let result: BooleanArray = array.iter().collect();
-
- assert_eq!(result, array);
-
- // check if DoubleEndedIterator is implemented
- let result: BooleanArray = array.iter().rev().collect();
- let rev_array = BooleanArray::from(vec![Some(false), None, Some(true)]);
- assert_eq!(result, rev_array);
-
- // check if ExactSizeIterator is implemented
- let _ = array.iter().rposition(|opt_b| opt_b == Some(true));
- }
-}
diff --git a/arrow/src/array/mod.rs b/arrow/src/array/mod.rs
deleted file mode 100644
index 6a0b94a..0000000
--- a/arrow/src/array/mod.rs
+++ /dev/null
@@ -1,283 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-//! The central type in Apache Arrow are arrays, represented
-//! by the [`Array` trait](crate::array::Array).
-//! An array represents a known-length sequence of values all
-//! having the same type.
-//!
-//! Internally, those values are represented by one or several
-//! [buffers](crate::buffer::Buffer), the number and meaning
-//! of which depend on the array’s data type, as documented in
-//! [the Arrow data layout specification](https://arrow.apache.org/docs/format/Columnar.html).
-//! For example, the type `Int16Array` represents an Apache
-//! Arrow array of 16-bit integers.
-//!
-//! Those buffers consist of the value data itself and an
-//! optional [bitmap buffer](crate::bitmap::Bitmap) that
-//! indicates which array entries are null values.
-//! The bitmap buffer can be entirely omitted if the array is
-//! known to have zero null values.
-//!
-//! There are concrete implementations of this trait for each
-//! data type, that help you access individual values of the
-//! array.
-//!
-//! # Building an Array
-//!
-//! Arrow's `Arrays` are immutable, but there is the trait
-//! [`ArrayBuilder`](crate::array::ArrayBuilder)
-//! that helps you with constructing new `Arrays`. As with the
-//! `Array` trait, there are builder implementations for all
-//! concrete array types.
-//!
-//! # Example
-//! ```
-//! extern crate arrow;
-//!
-//! use arrow::array::Int16Array;
-//!
-//! // Create a new builder with a capacity of 100
-//! let mut builder = Int16Array::builder(100);
-//!
-//! // Append a single primitive value
-//! builder.append_value(1).unwrap();
-//!
-//! // Append a null value
-//! builder.append_null().unwrap();
-//!
-//! // Append a slice of primitive values
-//! builder.append_slice(&[2, 3, 4]).unwrap();
-//!
-//! // Build the array
-//! let array = builder.finish();
-//!
-//! assert_eq!(
-//! 5,
-//! array.len(),
-//! "The array has 5 values, counting the null value"
-//! );
-//!
-//! assert_eq!(2, array.value(2), "Get the value with index 2");
-//!
-//! assert_eq!(
-//! &array.values()[3..5],
-//! &[3, 4],
-//! "Get slice of len 2 starting at idx 3"
-//! )
-//! ```
-
-#[allow(clippy::module_inception)]
-mod array;
-mod array_binary;
-mod array_boolean;
-mod array_dictionary;
-mod array_list;
-mod array_primitive;
-mod array_string;
-mod array_struct;
-mod array_union;
-mod builder;
-mod cast;
-mod data;
-mod equal;
-mod equal_json;
-mod ffi;
-mod iterator;
-mod null;
-mod ord;
-mod raw_pointer;
-mod transform;
-
-use crate::datatypes::*;
-
-// --------------------- Array & ArrayData ---------------------
-
-pub use self::array::Array;
-pub use self::array::ArrayRef;
-pub use self::data::ArrayData;
-pub use self::data::ArrayDataBuilder;
-pub use self::data::ArrayDataRef;
-
-pub use self::array_binary::BinaryArray;
-pub use self::array_binary::DecimalArray;
-pub use self::array_binary::FixedSizeBinaryArray;
-pub use self::array_binary::LargeBinaryArray;
-pub use self::array_boolean::BooleanArray;
-pub use self::array_dictionary::DictionaryArray;
-pub use self::array_list::FixedSizeListArray;
-pub use self::array_list::LargeListArray;
-pub use self::array_list::ListArray;
-pub use self::array_primitive::PrimitiveArray;
-pub use self::array_string::LargeStringArray;
-pub use self::array_string::StringArray;
-pub use self::array_struct::StructArray;
-pub use self::array_union::UnionArray;
-pub use self::null::NullArray;
-
-pub use self::array::make_array;
-pub use self::array::new_empty_array;
-pub use self::array::new_null_array;
-
-pub type Int8Array = PrimitiveArray<Int8Type>;
-pub type Int16Array = PrimitiveArray<Int16Type>;
-pub type Int32Array = PrimitiveArray<Int32Type>;
-pub type Int64Array = PrimitiveArray<Int64Type>;
-pub type UInt8Array = PrimitiveArray<UInt8Type>;
-pub type UInt16Array = PrimitiveArray<UInt16Type>;
-pub type UInt32Array = PrimitiveArray<UInt32Type>;
-pub type UInt64Array = PrimitiveArray<UInt64Type>;
-pub type Float32Array = PrimitiveArray<Float32Type>;
-pub type Float64Array = PrimitiveArray<Float64Type>;
-
-pub type Int8DictionaryArray = DictionaryArray<Int8Type>;
-pub type Int16DictionaryArray = DictionaryArray<Int16Type>;
-pub type Int32DictionaryArray = DictionaryArray<Int32Type>;
-pub type Int64DictionaryArray = DictionaryArray<Int64Type>;
-pub type UInt8DictionaryArray = DictionaryArray<UInt8Type>;
-pub type UInt16DictionaryArray = DictionaryArray<UInt16Type>;
-pub type UInt32DictionaryArray = DictionaryArray<UInt32Type>;
-pub type UInt64DictionaryArray = DictionaryArray<UInt64Type>;
-
-pub type TimestampSecondArray = PrimitiveArray<TimestampSecondType>;
-pub type TimestampMillisecondArray = PrimitiveArray<TimestampMillisecondType>;
-pub type TimestampMicrosecondArray = PrimitiveArray<TimestampMicrosecondType>;
-pub type TimestampNanosecondArray = PrimitiveArray<TimestampNanosecondType>;
-pub type Date32Array = PrimitiveArray<Date32Type>;
-pub type Date64Array = PrimitiveArray<Date64Type>;
-pub type Time32SecondArray = PrimitiveArray<Time32SecondType>;
-pub type Time32MillisecondArray = PrimitiveArray<Time32MillisecondType>;
-pub type Time64MicrosecondArray = PrimitiveArray<Time64MicrosecondType>;
-pub type Time64NanosecondArray = PrimitiveArray<Time64NanosecondType>;
-pub type IntervalYearMonthArray = PrimitiveArray<IntervalYearMonthType>;
-pub type IntervalDayTimeArray = PrimitiveArray<IntervalDayTimeType>;
-pub type DurationSecondArray = PrimitiveArray<DurationSecondType>;
-pub type DurationMillisecondArray = PrimitiveArray<DurationMillisecondType>;
-pub type DurationMicrosecondArray = PrimitiveArray<DurationMicrosecondType>;
-pub type DurationNanosecondArray = PrimitiveArray<DurationNanosecondType>;
-
-pub use self::array_binary::BinaryOffsetSizeTrait;
-pub use self::array_binary::GenericBinaryArray;
-pub use self::array_list::GenericListArray;
-pub use self::array_list::OffsetSizeTrait;
-pub use self::array_string::GenericStringArray;
-pub use self::array_string::StringOffsetSizeTrait;
-
-// --------------------- Array Builder ---------------------
-
-pub use self::builder::BooleanBufferBuilder;
-pub use self::builder::BufferBuilder;
-
-pub type Int8BufferBuilder = BufferBuilder<i8>;
-pub type Int16BufferBuilder = BufferBuilder<i16>;
-pub type Int32BufferBuilder = BufferBuilder<i32>;
-pub type Int64BufferBuilder = BufferBuilder<i64>;
-pub type UInt8BufferBuilder = BufferBuilder<u8>;
-pub type UInt16BufferBuilder = BufferBuilder<u16>;
-pub type UInt32BufferBuilder = BufferBuilder<u32>;
-pub type UInt64BufferBuilder = BufferBuilder<u64>;
-pub type Float32BufferBuilder = BufferBuilder<f32>;
-pub type Float64BufferBuilder = BufferBuilder<f64>;
-
-pub type TimestampSecondBufferBuilder = BufferBuilder<TimestampSecondType>;
-pub type TimestampMillisecondBufferBuilder = BufferBuilder<TimestampMillisecondType>;
-pub type TimestampMicrosecondBufferBuilder = BufferBuilder<TimestampMicrosecondType>;
-pub type TimestampNanosecondBufferBuilder = BufferBuilder<TimestampNanosecondType>;
-pub type Date32BufferBuilder = BufferBuilder<Date32Type>;
-pub type Date64BufferBuilder = BufferBuilder<Date64Type>;
-pub type Time32SecondBufferBuilder = BufferBuilder<Time32SecondType>;
-pub type Time32MillisecondBufferBuilder = BufferBuilder<Time32MillisecondType>;
-pub type Time64MicrosecondBufferBuilder = BufferBuilder<Time64MicrosecondType>;
-pub type Time64NanosecondBufferBuilder = BufferBuilder<Time64NanosecondType>;
-pub type IntervalYearMonthBufferBuilder = BufferBuilder<IntervalYearMonthType>;
-pub type IntervalDayTimeBufferBuilder = BufferBuilder<IntervalDayTimeType>;
-pub type DurationSecondBufferBuilder = BufferBuilder<DurationSecondType>;
-pub type DurationMillisecondBufferBuilder = BufferBuilder<DurationMillisecondType>;
-pub type DurationMicrosecondBufferBuilder = BufferBuilder<DurationMicrosecondType>;
-pub type DurationNanosecondBufferBuilder = BufferBuilder<DurationNanosecondType>;
-
-pub use self::builder::ArrayBuilder;
-pub use self::builder::BinaryBuilder;
-pub use self::builder::BooleanBuilder;
-pub use self::builder::DecimalBuilder;
-pub use self::builder::FixedSizeBinaryBuilder;
-pub use self::builder::FixedSizeListBuilder;
-pub use self::builder::GenericStringBuilder;
-pub use self::builder::LargeBinaryBuilder;
-pub use self::builder::LargeListBuilder;
-pub use self::builder::LargeStringBuilder;
-pub use self::builder::ListBuilder;
-pub use self::builder::PrimitiveBuilder;
-pub use self::builder::PrimitiveDictionaryBuilder;
-pub use self::builder::StringBuilder;
-pub use self::builder::StringDictionaryBuilder;
-pub use self::builder::StructBuilder;
-pub use self::builder::UnionBuilder;
-
-pub type Int8Builder = PrimitiveBuilder<Int8Type>;
-pub type Int16Builder = PrimitiveBuilder<Int16Type>;
-pub type Int32Builder = PrimitiveBuilder<Int32Type>;
-pub type Int64Builder = PrimitiveBuilder<Int64Type>;
-pub type UInt8Builder = PrimitiveBuilder<UInt8Type>;
-pub type UInt16Builder = PrimitiveBuilder<UInt16Type>;
-pub type UInt32Builder = PrimitiveBuilder<UInt32Type>;
-pub type UInt64Builder = PrimitiveBuilder<UInt64Type>;
-pub type Float32Builder = PrimitiveBuilder<Float32Type>;
-pub type Float64Builder = PrimitiveBuilder<Float64Type>;
-
-pub type TimestampSecondBuilder = PrimitiveBuilder<TimestampSecondType>;
-pub type TimestampMillisecondBuilder = PrimitiveBuilder<TimestampMillisecondType>;
-pub type TimestampMicrosecondBuilder = PrimitiveBuilder<TimestampMicrosecondType>;
-pub type TimestampNanosecondBuilder = PrimitiveBuilder<TimestampNanosecondType>;
-pub type Date32Builder = PrimitiveBuilder<Date32Type>;
-pub type Date64Builder = PrimitiveBuilder<Date64Type>;
-pub type Time32SecondBuilder = PrimitiveBuilder<Time32SecondType>;
-pub type Time32MillisecondBuilder = PrimitiveBuilder<Time32MillisecondType>;
-pub type Time64MicrosecondBuilder = PrimitiveBuilder<Time64MicrosecondType>;
-pub type Time64NanosecondBuilder = PrimitiveBuilder<Time64NanosecondType>;
-pub type IntervalYearMonthBuilder = PrimitiveBuilder<IntervalYearMonthType>;
-pub type IntervalDayTimeBuilder = PrimitiveBuilder<IntervalDayTimeType>;
-pub type DurationSecondBuilder = PrimitiveBuilder<DurationSecondType>;
-pub type DurationMillisecondBuilder = PrimitiveBuilder<DurationMillisecondType>;
-pub type DurationMicrosecondBuilder = PrimitiveBuilder<DurationMicrosecondType>;
-pub type DurationNanosecondBuilder = PrimitiveBuilder<DurationNanosecondType>;
-
-pub use self::transform::{Capacities, MutableArrayData};
-
-// --------------------- Array Iterator ---------------------
-
-pub use self::iterator::*;
-
-// --------------------- Array Equality ---------------------
-
-pub use self::equal_json::JsonEqual;
-
-// --------------------- Array's values comparison ---------------------
-
-pub use self::ord::{build_compare, DynComparator};
-
-// --------------------- Array downcast helper functions ---------------------
-
-pub use self::cast::{
- as_boolean_array, as_dictionary_array, as_generic_list_array, as_large_list_array,
- as_largestring_array, as_list_array, as_null_array, as_primitive_array,
- as_string_array, as_struct_array,
-};
-
-// ------------------------------ C Data Interface ---------------------------
-
-pub use self::array::make_array_from_raw;
diff --git a/arrow/src/array/null.rs b/arrow/src/array/null.rs
deleted file mode 100644
index 40513af..0000000
--- a/arrow/src/array/null.rs
+++ /dev/null
@@ -1,159 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-//! Contains the `NullArray` type.
-
-use std::any::Any;
-use std::fmt;
-use std::mem;
-
-use crate::array::{Array, ArrayData};
-use crate::datatypes::*;
-
-/// An Array where all elements are nulls
-///
-/// A `NullArray` is a simplified array where all values are null.
-///
-/// # Example: Create an array
-///
-/// ```
-/// use arrow::array::{Array, NullArray};
-///
-/// # fn main() -> arrow::error::Result<()> {
-/// let array = NullArray::new(10);
-///
-/// assert_eq!(array.len(), 10);
-/// assert_eq!(array.null_count(), 10);
-///
-/// # Ok(())
-/// # }
-/// ```
-pub struct NullArray {
- data: ArrayData,
-}
-
-impl NullArray {
- /// Create a new [`NullArray`] of the specified length
- ///
- /// *Note*: Use [`crate::array::new_null_array`] if you need an array of some
- /// other [`DataType`].
- ///
- pub fn new(length: usize) -> Self {
- let array_data = ArrayData::builder(DataType::Null).len(length).build();
- NullArray::from(array_data)
- }
-}
-
-impl Array for NullArray {
- fn as_any(&self) -> &Any {
- self
- }
-
- fn data(&self) -> &ArrayData {
- &self.data
- }
-
- /// Returns whether the element at `index` is null.
- /// All elements of a `NullArray` are always null.
- fn is_null(&self, _index: usize) -> bool {
- true
- }
-
- /// Returns whether the element at `index` is valid.
- /// All elements of a `NullArray` are always invalid.
- fn is_valid(&self, _index: usize) -> bool {
- false
- }
-
- /// Returns the total number of null values in this array.
- /// The null count of a `NullArray` always equals its length.
- fn null_count(&self) -> usize {
- self.data_ref().len()
- }
-
- /// Returns the total number of bytes of memory occupied by the buffers owned by this [NullArray].
- fn get_buffer_memory_size(&self) -> usize {
- self.data.get_buffer_memory_size()
- }
-
- /// Returns the total number of bytes of memory occupied physically by this [NullArray].
- fn get_array_memory_size(&self) -> usize {
- mem::size_of_val(self)
- }
-}
-
-impl From<ArrayData> for NullArray {
- fn from(data: ArrayData) -> Self {
- assert_eq!(
- data.data_type(),
- &DataType::Null,
- "NullArray data type should be Null"
- );
- assert_eq!(
- data.buffers().len(),
- 0,
- "NullArray data should contain 0 buffers"
- );
- assert!(
- data.null_buffer().is_none(),
- "NullArray data should not contain a null buffer, as no buffers are required"
- );
- Self { data }
- }
-}
-
-impl fmt::Debug for NullArray {
- fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
- write!(f, "NullArray({})", self.len())
- }
-}
-
-#[cfg(test)]
-mod tests {
- use super::*;
-
- #[test]
- fn test_null_array() {
- let null_arr = NullArray::new(32);
-
- assert_eq!(null_arr.len(), 32);
- assert_eq!(null_arr.null_count(), 32);
- assert_eq!(null_arr.is_valid(0), false);
-
- assert_eq!(0, null_arr.get_buffer_memory_size());
- assert_eq!(
- null_arr.get_buffer_memory_size() + std::mem::size_of::<NullArray>(),
- null_arr.get_array_memory_size()
- );
- }
-
- #[test]
- fn test_null_array_slice() {
- let array1 = NullArray::new(32);
-
- let array2 = array1.slice(8, 16);
- assert_eq!(array2.len(), 16);
- assert_eq!(array2.null_count(), 16);
- assert_eq!(array2.offset(), 8);
- }
-
- #[test]
- fn test_debug_null_array() {
- let array = NullArray::new(1024 * 1024);
- assert_eq!(format!("{:?}", array), "NullArray(1048576)");
- }
-}
diff --git a/arrow/src/array/ord.rs b/arrow/src/array/ord.rs
deleted file mode 100644
index d7b1829..0000000
--- a/arrow/src/array/ord.rs
+++ /dev/null
@@ -1,325 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-//! Contains functions and function factories to compare arrays.
-
-use std::cmp::Ordering;
-
-use crate::array::*;
-use crate::datatypes::TimeUnit;
-use crate::datatypes::*;
-use crate::error::{ArrowError, Result};
-
-use num::Float;
-
-/// Compare the values at two arbitrary indices in two arrays.
-pub type DynComparator<'a> = Box<dyn Fn(usize, usize) -> Ordering + 'a>;
-
-/// compares two floats, placing NaNs at last
-fn cmp_nans_last<T: Float>(a: &T, b: &T) -> Ordering {
- match (a.is_nan(), b.is_nan()) {
- (true, true) => Ordering::Equal,
- (true, false) => Ordering::Greater,
- (false, true) => Ordering::Less,
- _ => a.partial_cmp(b).unwrap(),
- }
-}
-
-fn compare_primitives<'a, T: ArrowPrimitiveType>(
- left: &'a Array,
- right: &'a Array,
-) -> DynComparator<'a>
-where
- T::Native: Ord,
-{
- let left = left.as_any().downcast_ref::<PrimitiveArray<T>>().unwrap();
- let right = right.as_any().downcast_ref::<PrimitiveArray<T>>().unwrap();
- Box::new(move |i, j| left.value(i).cmp(&right.value(j)))
-}
-
-fn compare_boolean<'a>(left: &'a Array, right: &'a Array) -> DynComparator<'a> {
- let left = left.as_any().downcast_ref::<BooleanArray>().unwrap();
- let right = right.as_any().downcast_ref::<BooleanArray>().unwrap();
- Box::new(move |i, j| left.value(i).cmp(&right.value(j)))
-}
-
-fn compare_float<'a, T: ArrowPrimitiveType>(
- left: &'a Array,
- right: &'a Array,
-) -> DynComparator<'a>
-where
- T::Native: Float,
-{
- let left = left.as_any().downcast_ref::<PrimitiveArray<T>>().unwrap();
- let right = right.as_any().downcast_ref::<PrimitiveArray<T>>().unwrap();
- Box::new(move |i, j| cmp_nans_last(&left.value(i), &right.value(j)))
-}
-
-fn compare_string<'a, T>(left: &'a Array, right: &'a Array) -> DynComparator<'a>
-where
- T: StringOffsetSizeTrait,
-{
- let left = left
- .as_any()
- .downcast_ref::<GenericStringArray<T>>()
- .unwrap();
- let right = right
- .as_any()
- .downcast_ref::<GenericStringArray<T>>()
- .unwrap();
- Box::new(move |i, j| left.value(i).cmp(&right.value(j)))
-}
-
-fn compare_dict_string<'a, T>(left: &'a Array, right: &'a Array) -> DynComparator<'a>
-where
- T: ArrowDictionaryKeyType,
-{
- let left = left.as_any().downcast_ref::<DictionaryArray<T>>().unwrap();
- let right = right.as_any().downcast_ref::<DictionaryArray<T>>().unwrap();
- let left_keys = left.keys();
- let right_keys = right.keys();
-
- let left_values = StringArray::from(left.values().data().clone());
- let right_values = StringArray::from(right.values().data().clone());
-
- Box::new(move |i: usize, j: usize| {
- let key_left = left_keys.value(i).to_usize().unwrap();
- let key_right = right_keys.value(j).to_usize().unwrap();
- let left = left_values.value(key_left);
- let right = right_values.value(key_right);
- left.cmp(&right)
- })
-}
-
-/// returns a comparison function that compares two values at two different positions
-/// between the two arrays.
-/// The arrays' types must be equal.
-/// # Example
-/// ```
-/// use arrow::array::{build_compare, Int32Array};
-///
-/// # fn main() -> arrow::error::Result<()> {
-/// let array1 = Int32Array::from(vec![1, 2]);
-/// let array2 = Int32Array::from(vec![3, 4]);
-///
-/// let cmp = build_compare(&array1, &array2)?;
-///
-/// // 1 (index 0 of array1) is smaller than 4 (index 1 of array2)
-/// assert_eq!(std::cmp::Ordering::Less, (cmp)(0, 1));
-/// # Ok(())
-/// # }
-/// ```
-// This is a factory of comparisons.
-// The lifetime 'a enforces that we cannot use the closure beyond any of the array's lifetime.
-pub fn build_compare<'a>(left: &'a Array, right: &'a Array) -> Result<DynComparator<'a>> {
- use DataType::*;
- use IntervalUnit::*;
- use TimeUnit::*;
- Ok(match (left.data_type(), right.data_type()) {
- (a, b) if a != b => {
- return Err(ArrowError::InvalidArgumentError(
- "Can't compare arrays of different types".to_string(),
- ));
- }
- (Boolean, Boolean) => compare_boolean(left, right),
- (UInt8, UInt8) => compare_primitives::<UInt8Type>(left, right),
- (UInt16, UInt16) => compare_primitives::<UInt16Type>(left, right),
- (UInt32, UInt32) => compare_primitives::<UInt32Type>(left, right),
- (UInt64, UInt64) => compare_primitives::<UInt64Type>(left, right),
- (Int8, Int8) => compare_primitives::<Int8Type>(left, right),
- (Int16, Int16) => compare_primitives::<Int16Type>(left, right),
- (Int32, Int32) => compare_primitives::<Int32Type>(left, right),
- (Int64, Int64) => compare_primitives::<Int64Type>(left, right),
- (Float32, Float32) => compare_float::<Float32Type>(left, right),
- (Float64, Float64) => compare_float::<Float64Type>(left, right),
- (Date32, Date32) => compare_primitives::<Date32Type>(left, right),
- (Date64, Date64) => compare_primitives::<Date64Type>(left, right),
- (Time32(Second), Time32(Second)) => {
- compare_primitives::<Time32SecondType>(left, right)
- }
- (Time32(Millisecond), Time32(Millisecond)) => {
- compare_primitives::<Time32MillisecondType>(left, right)
- }
- (Time64(Microsecond), Time64(Microsecond)) => {
- compare_primitives::<Time64MicrosecondType>(left, right)
- }
- (Time64(Nanosecond), Time64(Nanosecond)) => {
- compare_primitives::<Time64NanosecondType>(left, right)
- }
- (Timestamp(Second, _), Timestamp(Second, _)) => {
- compare_primitives::<TimestampSecondType>(left, right)
- }
- (Timestamp(Millisecond, _), Timestamp(Millisecond, _)) => {
- compare_primitives::<TimestampMillisecondType>(left, right)
- }
- (Timestamp(Microsecond, _), Timestamp(Microsecond, _)) => {
- compare_primitives::<TimestampMicrosecondType>(left, right)
- }
- (Timestamp(Nanosecond, _), Timestamp(Nanosecond, _)) => {
- compare_primitives::<TimestampNanosecondType>(left, right)
- }
- (Interval(YearMonth), Interval(YearMonth)) => {
- compare_primitives::<IntervalYearMonthType>(left, right)
- }
- (Interval(DayTime), Interval(DayTime)) => {
- compare_primitives::<IntervalDayTimeType>(left, right)
- }
- (Duration(Second), Duration(Second)) => {
- compare_primitives::<DurationSecondType>(left, right)
- }
- (Duration(Millisecond), Duration(Millisecond)) => {
- compare_primitives::<DurationMillisecondType>(left, right)
- }
- (Duration(Microsecond), Duration(Microsecond)) => {
- compare_primitives::<DurationMicrosecondType>(left, right)
- }
- (Duration(Nanosecond), Duration(Nanosecond)) => {
- compare_primitives::<DurationNanosecondType>(left, right)
- }
- (Utf8, Utf8) => compare_string::<i32>(left, right),
- (LargeUtf8, LargeUtf8) => compare_string::<i64>(left, right),
- (
- Dictionary(key_type_lhs, value_type_lhs),
- Dictionary(key_type_rhs, value_type_rhs),
- ) => {
- if value_type_lhs.as_ref() != &DataType::Utf8
- || value_type_rhs.as_ref() != &DataType::Utf8
- {
- return Err(ArrowError::InvalidArgumentError(
- "Arrow still does not support comparisons of non-string dictionary arrays"
- .to_string(),
- ));
- }
- match (key_type_lhs.as_ref(), key_type_rhs.as_ref()) {
- (a, b) if a != b => {
- return Err(ArrowError::InvalidArgumentError(
- "Can't compare arrays of different types".to_string(),
- ));
- }
- (UInt8, UInt8) => compare_dict_string::<UInt8Type>(left, right),
- (UInt16, UInt16) => compare_dict_string::<UInt16Type>(left, right),
- (UInt32, UInt32) => compare_dict_string::<UInt32Type>(left, right),
- (UInt64, UInt64) => compare_dict_string::<UInt64Type>(left, right),
- (Int8, Int8) => compare_dict_string::<Int8Type>(left, right),
- (Int16, Int16) => compare_dict_string::<Int16Type>(left, right),
- (Int32, Int32) => compare_dict_string::<Int32Type>(left, right),
- (Int64, Int64) => compare_dict_string::<Int64Type>(left, right),
- (lhs, _) => {
- return Err(ArrowError::InvalidArgumentError(format!(
- "Dictionaries do not support keys of type {:?}",
- lhs
- )));
- }
- }
- }
- (lhs, _) => {
- return Err(ArrowError::InvalidArgumentError(format!(
- "The data type type {:?} has no natural order",
- lhs
- )));
- }
- })
-}
-
-#[cfg(test)]
-pub mod tests {
- use super::*;
- use crate::array::{Float64Array, Int32Array};
- use crate::error::Result;
- use std::cmp::Ordering;
- use std::iter::FromIterator;
-
- #[test]
- fn test_i32() -> Result<()> {
- let array = Int32Array::from(vec![1, 2]);
-
- let cmp = build_compare(&array, &array)?;
-
- assert_eq!(Ordering::Less, (cmp)(0, 1));
- Ok(())
- }
-
- #[test]
- fn test_i32_i32() -> Result<()> {
- let array1 = Int32Array::from(vec![1]);
- let array2 = Int32Array::from(vec![2]);
-
- let cmp = build_compare(&array1, &array2)?;
-
- assert_eq!(Ordering::Less, (cmp)(0, 0));
- Ok(())
- }
-
- #[test]
- fn test_f64() -> Result<()> {
- let array = Float64Array::from(vec![1.0, 2.0]);
-
- let cmp = build_compare(&array, &array)?;
-
- assert_eq!(Ordering::Less, (cmp)(0, 1));
- Ok(())
- }
-
- #[test]
- fn test_f64_nan() -> Result<()> {
- let array = Float64Array::from(vec![1.0, f64::NAN]);
-
- let cmp = build_compare(&array, &array)?;
-
- assert_eq!(Ordering::Less, (cmp)(0, 1));
- Ok(())
- }
-
- #[test]
- fn test_f64_zeros() -> Result<()> {
- let array = Float64Array::from(vec![-0.0, 0.0]);
-
- let cmp = build_compare(&array, &array)?;
-
- assert_eq!(Ordering::Equal, (cmp)(0, 1));
- assert_eq!(Ordering::Equal, (cmp)(1, 0));
- Ok(())
- }
-
- #[test]
- fn test_dict() -> Result<()> {
- let data = vec!["a", "b", "c", "a", "a", "c", "c"];
- let array = DictionaryArray::<Int16Type>::from_iter(data.into_iter());
-
- let cmp = build_compare(&array, &array)?;
-
- assert_eq!(Ordering::Less, (cmp)(0, 1));
- assert_eq!(Ordering::Equal, (cmp)(3, 4));
- assert_eq!(Ordering::Greater, (cmp)(2, 3));
- Ok(())
- }
-
- #[test]
- fn test_multiple_dict() -> Result<()> {
- let d1 = vec!["a", "b", "c", "d"];
- let a1 = DictionaryArray::<Int16Type>::from_iter(d1.into_iter());
- let d2 = vec!["e", "f", "g", "a"];
- let a2 = DictionaryArray::<Int16Type>::from_iter(d2.into_iter());
-
- let cmp = build_compare(&a1, &a2)?;
-
- assert_eq!(Ordering::Less, (cmp)(0, 0));
- assert_eq!(Ordering::Equal, (cmp)(0, 3));
- assert_eq!(Ordering::Greater, (cmp)(1, 3));
- Ok(())
- }
-}
diff --git a/arrow/src/array/raw_pointer.rs b/arrow/src/array/raw_pointer.rs
deleted file mode 100644
index cd1f802..0000000
--- a/arrow/src/array/raw_pointer.rs
+++ /dev/null
@@ -1,65 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-use std::ptr::NonNull;
-
-/// This struct is highly `unsafe` and offers the possibility to self-reference a [arrow::Buffer] from [arrow::array::ArrayData].
-/// as a pointer to the beginning of its contents.
-pub(super) struct RawPtrBox<T> {
- ptr: NonNull<T>,
-}
-
-impl<T> RawPtrBox<T> {
- /// # Safety
- /// The user must guarantee that:
- /// * the contents where `ptr` points to are never `moved`. This is guaranteed when they are Pinned.
- /// * the lifetime of this struct does not outlive the lifetime of `ptr`.
- /// Failure to fulfill any the above conditions results in undefined behavior.
- /// # Panic
- /// This function panics if:
- /// * `ptr` is null
- /// * `ptr` is not aligned to a slice of type `T`. This is guaranteed if it was built from a slice of type `T`.
- pub(super) unsafe fn new(ptr: *const u8) -> Self {
- let ptr = NonNull::new(ptr as *mut u8).expect("Pointer cannot be null");
- assert_eq!(
- ptr.as_ptr().align_offset(std::mem::align_of::<T>()),
- 0,
- "memory is not aligned"
- );
- Self { ptr: ptr.cast() }
- }
-
- pub(super) fn as_ptr(&self) -> *const T {
- self.ptr.as_ptr()
- }
-}
-
-unsafe impl<T> Send for RawPtrBox<T> {}
-unsafe impl<T> Sync for RawPtrBox<T> {}
-
-#[cfg(test)]
-mod tests {
- use super::*;
-
- #[test]
- #[should_panic(expected = "memory is not aligned")]
- #[cfg_attr(miri, ignore)] // sometimes does not panic as expected
- fn test_primitive_array_alignment() {
- let bytes = vec![0u8, 1u8];
- unsafe { RawPtrBox::<u64>::new(bytes.as_ptr().offset(1)) };
- }
-}
diff --git a/arrow/src/array/transform/boolean.rs b/arrow/src/array/transform/boolean.rs
deleted file mode 100644
index 1829149..0000000
--- a/arrow/src/array/transform/boolean.rs
+++ /dev/null
@@ -1,45 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-use crate::array::ArrayData;
-
-use super::{
- Extend, _MutableArrayData,
- utils::{resize_for_bits, set_bits},
-};
-
-pub(super) fn build_extend(array: &ArrayData) -> Extend {
- let values = array.buffers()[0].as_slice();
- Box::new(
- move |mutable: &mut _MutableArrayData, _, start: usize, len: usize| {
- let buffer = &mut mutable.buffer1;
- resize_for_bits(buffer, mutable.len + len);
- set_bits(
- &mut buffer.as_slice_mut(),
- values,
- mutable.len,
- array.offset() + start,
- len,
- );
- },
- )
-}
-
-pub(super) fn extend_nulls(mutable: &mut _MutableArrayData, len: usize) {
- let buffer = &mut mutable.buffer1;
- resize_for_bits(buffer, mutable.len + len);
-}
diff --git a/arrow/src/array/transform/fixed_binary.rs b/arrow/src/array/transform/fixed_binary.rs
deleted file mode 100644
index 36952d4..0000000
--- a/arrow/src/array/transform/fixed_binary.rs
+++ /dev/null
@@ -1,65 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-use crate::{array::ArrayData, datatypes::DataType};
-
-use super::{Extend, _MutableArrayData};
-
-pub(super) fn build_extend(array: &ArrayData) -> Extend {
- let size = match array.data_type() {
- DataType::FixedSizeBinary(i) => *i as usize,
- _ => unreachable!(),
- };
-
- let values = &array.buffers()[0].as_slice()[array.offset() * size..];
- if array.null_count() == 0 {
- // fast case where we can copy regions without null issues
- Box::new(
- move |mutable: &mut _MutableArrayData, _, start: usize, len: usize| {
- let buffer = &mut mutable.buffer1;
- buffer.extend_from_slice(&values[start * size..(start + len) * size]);
- },
- )
- } else {
- Box::new(
- move |mutable: &mut _MutableArrayData, _, start: usize, len: usize| {
- // nulls present: append item by item, ignoring null entries
- let values_buffer = &mut mutable.buffer1;
-
- (start..start + len).for_each(|i| {
- if array.is_valid(i) {
- // append value
- let bytes = &values[i * size..(i + 1) * size];
- values_buffer.extend_from_slice(bytes);
- } else {
- values_buffer.extend_zeros(size);
- }
- })
- },
- )
- }
-}
-
-pub(super) fn extend_nulls(mutable: &mut _MutableArrayData, len: usize) {
- let size = match mutable.data_type {
- DataType::FixedSizeBinary(i) => i as usize,
- _ => unreachable!(),
- };
-
- let values_buffer = &mut mutable.buffer1;
- values_buffer.extend_zeros(len * size);
-}
diff --git a/arrow/src/array/transform/list.rs b/arrow/src/array/transform/list.rs
deleted file mode 100644
index 8eb2bd1..0000000
--- a/arrow/src/array/transform/list.rs
+++ /dev/null
@@ -1,99 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-use crate::array::{ArrayData, OffsetSizeTrait};
-
-use super::{
- Extend, _MutableArrayData,
- utils::{extend_offsets, get_last_offset},
-};
-
-pub(super) fn build_extend<T: OffsetSizeTrait>(array: &ArrayData) -> Extend {
- let offsets = array.buffer::<T>(0);
- if array.null_count() == 0 {
- // fast case where we can copy regions without nullability checks
- Box::new(
- move |mutable: &mut _MutableArrayData,
- index: usize,
- start: usize,
- len: usize| {
- let offset_buffer = &mut mutable.buffer1;
-
- // this is safe due to how offset is built. See details on `get_last_offset`
- let last_offset: T = unsafe { get_last_offset(offset_buffer) };
-
- // offsets
- extend_offsets::<T>(
- offset_buffer,
- last_offset,
- &offsets[start..start + len + 1],
- );
-
- mutable.child_data[0].extend(
- index,
- offsets[start].to_usize().unwrap(),
- offsets[start + len].to_usize().unwrap(),
- )
- },
- )
- } else {
- // nulls present: append item by item, ignoring null entries
- Box::new(
- move |mutable: &mut _MutableArrayData,
- index: usize,
- start: usize,
- len: usize| {
- let offset_buffer = &mut mutable.buffer1;
-
- // this is safe due to how offset is built. See details on `get_last_offset`
- let mut last_offset: T = unsafe { get_last_offset(offset_buffer) };
-
- let delta_len = array.len() - array.null_count();
- offset_buffer.reserve(delta_len * std::mem::size_of::<T>());
-
- let child = &mut mutable.child_data[0];
- (start..start + len).for_each(|i| {
- if array.is_valid(i) {
- // compute the new offset
- last_offset += offsets[i + 1] - offsets[i];
-
- // append value
- child.extend(
- index,
- offsets[i].to_usize().unwrap(),
- offsets[i + 1].to_usize().unwrap(),
- );
- }
- // append offset
- offset_buffer.push(last_offset);
- })
- },
- )
- }
-}
-
-pub(super) fn extend_nulls<T: OffsetSizeTrait>(
- mutable: &mut _MutableArrayData,
- len: usize,
-) {
- let offset_buffer = &mut mutable.buffer1;
-
- // this is safe due to how offset is built. See details on `get_last_offset`
- let last_offset: T = unsafe { get_last_offset(offset_buffer) };
-
- (0..len).for_each(|_| offset_buffer.push(last_offset))
-}
diff --git a/arrow/src/array/transform/mod.rs b/arrow/src/array/transform/mod.rs
deleted file mode 100644
index 0194b93..0000000
--- a/arrow/src/array/transform/mod.rs
+++ /dev/null
@@ -1,1456 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-use crate::{
- buffer::MutableBuffer,
- datatypes::DataType,
- error::{ArrowError, Result},
- util::bit_util,
-};
-use std::mem;
-
-use super::{
- data::{into_buffers, new_buffers},
- ArrayData, ArrayDataBuilder,
-};
-use crate::array::StringOffsetSizeTrait;
-
-mod boolean;
-mod fixed_binary;
-mod list;
-mod null;
-mod primitive;
-mod structure;
-mod utils;
-mod variable_size;
-
-type ExtendNullBits<'a> = Box<Fn(&mut _MutableArrayData, usize, usize) + 'a>;
-// function that extends `[start..start+len]` to the mutable array.
-// this is dynamic because different data_types influence how buffers and childs are extended.
-type Extend<'a> = Box<Fn(&mut _MutableArrayData, usize, usize, usize) + 'a>;
-
-type ExtendNulls = Box<Fn(&mut _MutableArrayData, usize) -> ()>;
-
-/// A mutable [ArrayData] that knows how to freeze itself into an [ArrayData].
-/// This is just a data container.
-#[derive(Debug)]
-struct _MutableArrayData<'a> {
- pub data_type: DataType,
- pub null_count: usize,
-
- pub len: usize,
- pub null_buffer: MutableBuffer,
-
- // arrow specification only allows up to 3 buffers (2 ignoring the nulls above).
- // Thus, we place them in the stack to avoid bound checks and greater data locality.
- pub buffer1: MutableBuffer,
- pub buffer2: MutableBuffer,
- pub child_data: Vec<MutableArrayData<'a>>,
-}
-
-impl<'a> _MutableArrayData<'a> {
- fn freeze(self, dictionary: Option<ArrayData>) -> ArrayDataBuilder {
- let buffers = into_buffers(&self.data_type, self.buffer1, self.buffer2);
-
- let child_data = match self.data_type {
- DataType::Dictionary(_, _) => vec![dictionary.unwrap()],
- _ => {
- let mut child_data = Vec::with_capacity(self.child_data.len());
- for child in self.child_data {
- child_data.push(child.freeze());
- }
- child_data
- }
- };
-
- let mut array_data_builder = ArrayDataBuilder::new(self.data_type)
- .offset(0)
- .len(self.len)
- .null_count(self.null_count)
- .buffers(buffers)
- .child_data(child_data);
- if self.null_count > 0 {
- array_data_builder =
- array_data_builder.null_bit_buffer(self.null_buffer.into());
- }
-
- array_data_builder
- }
-}
-
-fn build_extend_null_bits(array: &ArrayData, use_nulls: bool) -> ExtendNullBits {
- if let Some(bitmap) = array.null_bitmap() {
- let bytes = bitmap.bits.as_slice();
- Box::new(move |mutable, start, len| {
- utils::resize_for_bits(&mut mutable.null_buffer, mutable.len + len);
- mutable.null_count += utils::set_bits(
- mutable.null_buffer.as_slice_mut(),
- bytes,
- mutable.len,
- array.offset() + start,
- len,
- );
- })
- } else if use_nulls {
- Box::new(|mutable, _, len| {
- utils::resize_for_bits(&mut mutable.null_buffer, mutable.len + len);
- let write_data = mutable.null_buffer.as_slice_mut();
- let offset = mutable.len;
- (0..len).for_each(|i| {
- bit_util::set_bit(write_data, offset + i);
- });
- })
- } else {
- Box::new(|_, _, _| {})
- }
-}
-
-/// Struct to efficiently and interactively create an [ArrayData] from an existing [ArrayData] by
-/// copying chunks.
-/// The main use case of this struct is to perform unary operations to arrays of arbitrary types, such as `filter` and `take`.
-/// # Example:
-///
-/// ```
-/// use arrow::{array::{Int32Array, Array, MutableArrayData}};
-///
-/// let array = Int32Array::from(vec![1, 2, 3, 4, 5]);
-/// let array = array.data();
-/// // Create a new `MutableArrayData` from an array and with a capacity of 4.
-/// // Capacity here is equivalent to `Vec::with_capacity`
-/// let arrays = vec![array];
-/// let mut mutable = MutableArrayData::new(arrays, false, 4);
-/// mutable.extend(0, 1, 3); // extend from the slice [1..3], [2,3]
-/// mutable.extend(0, 0, 3); // extend from the slice [0..3], [1,2,3]
-/// // `.freeze()` to convert `MutableArrayData` into a `ArrayData`.
-/// let new_array = Int32Array::from(mutable.freeze());
-/// assert_eq!(Int32Array::from(vec![2, 3, 1, 2, 3]), new_array);
-/// ```
-pub struct MutableArrayData<'a> {
- arrays: Vec<&'a ArrayData>,
- // The attributes in [_MutableArrayData] cannot be in [MutableArrayData] due to
- // mutability invariants (interior mutability):
- // [MutableArrayData] contains a function that can only mutate [_MutableArrayData], not
- // [MutableArrayData] itself
- data: _MutableArrayData<'a>,
-
- // the child data of the `Array` in Dictionary arrays.
- // This is not stored in `MutableArrayData` because these values constant and only needed
- // at the end, when freezing [_MutableArrayData].
- dictionary: Option<ArrayData>,
-
- // function used to extend values from arrays. This function's lifetime is bound to the array
- // because it reads values from it.
- extend_values: Vec<Extend<'a>>,
- // function used to extend nulls from arrays. This function's lifetime is bound to the array
- // because it reads nulls from it.
- extend_null_bits: Vec<ExtendNullBits<'a>>,
-
- // function used to extend nulls.
- // this is independent of the arrays and therefore has no lifetime.
- extend_nulls: ExtendNulls,
-}
-
-impl<'a> std::fmt::Debug for MutableArrayData<'a> {
- fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
- // ignores the closures.
- f.debug_struct("MutableArrayData")
- .field("data", &self.data)
- .finish()
- }
-}
-
-/// Builds an extend that adds `offset` to the source primitive
-/// Additionally validates that `max` fits into the
-/// the underlying primitive returning None if not
-fn build_extend_dictionary(
- array: &ArrayData,
- offset: usize,
- max: usize,
-) -> Option<Extend> {
- use crate::datatypes::*;
- use std::convert::TryInto;
-
- match array.data_type() {
- DataType::Dictionary(child_data_type, _) => match child_data_type.as_ref() {
- DataType::UInt8 => {
- let _: u8 = max.try_into().ok()?;
- let offset: u8 = offset.try_into().ok()?;
- Some(primitive::build_extend_with_offset(array, offset))
- }
- DataType::UInt16 => {
- let _: u16 = max.try_into().ok()?;
- let offset: u16 = offset.try_into().ok()?;
- Some(primitive::build_extend_with_offset(array, offset))
- }
- DataType::UInt32 => {
- let _: u32 = max.try_into().ok()?;
- let offset: u32 = offset.try_into().ok()?;
- Some(primitive::build_extend_with_offset(array, offset))
- }
- DataType::UInt64 => {
- let _: u64 = max.try_into().ok()?;
- let offset: u64 = offset.try_into().ok()?;
- Some(primitive::build_extend_with_offset(array, offset))
- }
- DataType::Int8 => {
- let _: i8 = max.try_into().ok()?;
- let offset: i8 = offset.try_into().ok()?;
- Some(primitive::build_extend_with_offset(array, offset))
- }
- DataType::Int16 => {
- let _: i16 = max.try_into().ok()?;
- let offset: i16 = offset.try_into().ok()?;
- Some(primitive::build_extend_with_offset(array, offset))
- }
- DataType::Int32 => {
- let _: i32 = max.try_into().ok()?;
- let offset: i32 = offset.try_into().ok()?;
- Some(primitive::build_extend_with_offset(array, offset))
- }
- DataType::Int64 => {
- let _: i64 = max.try_into().ok()?;
- let offset: i64 = offset.try_into().ok()?;
- Some(primitive::build_extend_with_offset(array, offset))
- }
- _ => unreachable!(),
- },
- _ => None,
- }
-}
-
-fn build_extend(array: &ArrayData) -> Extend {
- use crate::datatypes::*;
- match array.data_type() {
- DataType::Null => null::build_extend(array),
- DataType::Boolean => boolean::build_extend(array),
- DataType::UInt8 => primitive::build_extend::<u8>(array),
- DataType::UInt16 => primitive::build_extend::<u16>(array),
- DataType::UInt32 => primitive::build_extend::<u32>(array),
- DataType::UInt64 => primitive::build_extend::<u64>(array),
- DataType::Int8 => primitive::build_extend::<i8>(array),
- DataType::Int16 => primitive::build_extend::<i16>(array),
- DataType::Int32 => primitive::build_extend::<i32>(array),
- DataType::Int64 => primitive::build_extend::<i64>(array),
- DataType::Float32 => primitive::build_extend::<f32>(array),
- DataType::Float64 => primitive::build_extend::<f64>(array),
- DataType::Date32
- | DataType::Time32(_)
- | DataType::Interval(IntervalUnit::YearMonth) => {
- primitive::build_extend::<i32>(array)
- }
- DataType::Date64
- | DataType::Time64(_)
- | DataType::Timestamp(_, _)
- | DataType::Duration(_)
- | DataType::Interval(IntervalUnit::DayTime) => {
- primitive::build_extend::<i64>(array)
- }
- DataType::Utf8 | DataType::Binary => variable_size::build_extend::<i32>(array),
- DataType::LargeUtf8 | DataType::LargeBinary => {
- variable_size::build_extend::<i64>(array)
- }
- DataType::List(_) => list::build_extend::<i32>(array),
- DataType::LargeList(_) => list::build_extend::<i64>(array),
- DataType::Dictionary(_, _) => unreachable!("should use build_extend_dictionary"),
- DataType::Struct(_) => structure::build_extend(array),
- DataType::FixedSizeBinary(_) => fixed_binary::build_extend(array),
- DataType::Float16 => unreachable!(),
- /*
- DataType::FixedSizeList(_, _) => {}
- DataType::Union(_) => {}
- */
- _ => todo!("Take and filter operations still not supported for this datatype"),
- }
-}
-
-fn build_extend_nulls(data_type: &DataType) -> ExtendNulls {
- use crate::datatypes::*;
- Box::new(match data_type {
- DataType::Null => null::extend_nulls,
- DataType::Boolean => boolean::extend_nulls,
- DataType::UInt8 => primitive::extend_nulls::<u8>,
- DataType::UInt16 => primitive::extend_nulls::<u16>,
- DataType::UInt32 => primitive::extend_nulls::<u32>,
- DataType::UInt64 => primitive::extend_nulls::<u64>,
- DataType::Int8 => primitive::extend_nulls::<i8>,
- DataType::Int16 => primitive::extend_nulls::<i16>,
- DataType::Int32 => primitive::extend_nulls::<i32>,
- DataType::Int64 => primitive::extend_nulls::<i64>,
- DataType::Float32 => primitive::extend_nulls::<f32>,
- DataType::Float64 => primitive::extend_nulls::<f64>,
- DataType::Date32
- | DataType::Time32(_)
- | DataType::Interval(IntervalUnit::YearMonth) => primitive::extend_nulls::<i32>,
- DataType::Date64
- | DataType::Time64(_)
- | DataType::Timestamp(_, _)
- | DataType::Duration(_)
- | DataType::Interval(IntervalUnit::DayTime) => primitive::extend_nulls::<i64>,
- DataType::Utf8 | DataType::Binary => variable_size::extend_nulls::<i32>,
- DataType::LargeUtf8 | DataType::LargeBinary => variable_size::extend_nulls::<i64>,
- DataType::List(_) => list::extend_nulls::<i32>,
- DataType::LargeList(_) => list::extend_nulls::<i64>,
- DataType::Dictionary(child_data_type, _) => match child_data_type.as_ref() {
- DataType::UInt8 => primitive::extend_nulls::<u8>,
- DataType::UInt16 => primitive::extend_nulls::<u16>,
- DataType::UInt32 => primitive::extend_nulls::<u32>,
- DataType::UInt64 => primitive::extend_nulls::<u64>,
- DataType::Int8 => primitive::extend_nulls::<i8>,
- DataType::Int16 => primitive::extend_nulls::<i16>,
- DataType::Int32 => primitive::extend_nulls::<i32>,
- DataType::Int64 => primitive::extend_nulls::<i64>,
- _ => unreachable!(),
- },
- DataType::Struct(_) => structure::extend_nulls,
- DataType::FixedSizeBinary(_) => fixed_binary::extend_nulls,
- DataType::Float16 => unreachable!(),
- /*
- DataType::FixedSizeList(_, _) => {}
- DataType::Union(_) => {}
- */
- _ => todo!("Take and filter operations still not supported for this datatype"),
- })
-}
-
-fn preallocate_offset_and_binary_buffer<Offset: StringOffsetSizeTrait>(
- capacity: usize,
- binary_size: usize,
-) -> [MutableBuffer; 2] {
- // offsets
- let mut buffer = MutableBuffer::new((1 + capacity) * mem::size_of::<Offset>());
- // safety: `unsafe` code assumes that this buffer is initialized with one element
- if Offset::is_large() {
- buffer.push(0i64);
- } else {
- buffer.push(0i32)
- }
-
- [
- buffer,
- MutableBuffer::new(binary_size * mem::size_of::<u8>()),
- ]
-}
-
-/// Define capacities of child data or data buffers.
-#[derive(Debug, Clone)]
-pub enum Capacities {
- /// Binary, Utf8 and LargeUtf8 data types
- /// Define
- /// * the capacity of the array offsets
- /// * the capacity of the binary/ str buffer
- Binary(usize, Option<usize>),
- /// List and LargeList data types
- /// Define
- /// * the capacity of the array offsets
- /// * the capacity of the child data
- List(usize, Option<Box<Capacities>>),
- /// Struct type
- /// * the capacity of the array
- /// * the capacities of the fields
- Struct(usize, Option<Vec<Capacities>>),
- /// Dictionary type
- /// * the capacity of the array/keys
- /// * the capacity of the values
- Dictionary(usize, Option<Box<Capacities>>),
- /// Don't preallocate inner buffers and rely on array growth strategy
- Array(usize),
-}
-impl<'a> MutableArrayData<'a> {
- /// returns a new [MutableArrayData] with capacity to `capacity` slots and specialized to create an
- /// [ArrayData] from multiple `arrays`.
- ///
- /// `use_nulls` is a flag used to optimize insertions. It should be `false` if the only source of nulls
- /// are the arrays themselves and `true` if the user plans to call [MutableArrayData::extend_nulls].
- /// In other words, if `use_nulls` is `false`, calling [MutableArrayData::extend_nulls] should not be used.
- pub fn new(arrays: Vec<&'a ArrayData>, use_nulls: bool, capacity: usize) -> Self {
- Self::with_capacities(arrays, use_nulls, Capacities::Array(capacity))
- }
-
- /// Similar to [MutableArray::new], but lets users define the preallocated capacities of the array.
- /// See also [MutableArray::new] for more information on the arguments.
- ///
- /// # Panic
- /// This function panics if the given `capacities` don't match the data type of `arrays`. Or when
- /// a [Capacities] variant is not yet supported.
- pub fn with_capacities(
- arrays: Vec<&'a ArrayData>,
- mut use_nulls: bool,
- capacities: Capacities,
- ) -> Self {
- let data_type = arrays[0].data_type();
- use crate::datatypes::*;
-
- // if any of the arrays has nulls, insertions from any array requires setting bits
- // as there is at least one array with nulls.
- if arrays.iter().any(|array| array.null_count() > 0) {
- use_nulls = true;
- };
-
- let mut array_capacity;
-
- let [buffer1, buffer2] = match (data_type, &capacities) {
- (DataType::LargeUtf8, Capacities::Binary(capacity, Some(value_cap)))
- | (DataType::LargeBinary, Capacities::Binary(capacity, Some(value_cap))) => {
- array_capacity = *capacity;
- preallocate_offset_and_binary_buffer::<i64>(*capacity, *value_cap)
- }
- (DataType::Utf8, Capacities::Binary(capacity, Some(value_cap)))
- | (DataType::Binary, Capacities::Binary(capacity, Some(value_cap))) => {
- array_capacity = *capacity;
- preallocate_offset_and_binary_buffer::<i32>(*capacity, *value_cap)
- }
- (_, Capacities::Array(capacity)) => {
- array_capacity = *capacity;
- new_buffers(data_type, *capacity)
- }
- _ => panic!("Capacities: {:?} not yet supported", capacities),
- };
-
- let child_data = match &data_type {
- DataType::Null
- | DataType::Boolean
- | DataType::UInt8
- | DataType::UInt16
- | DataType::UInt32
- | DataType::UInt64
- | DataType::Int8
- | DataType::Int16
- | DataType::Int32
- | DataType::Int64
- | DataType::Float32
- | DataType::Float64
- | DataType::Date32
- | DataType::Date64
- | DataType::Time32(_)
- | DataType::Time64(_)
- | DataType::Duration(_)
- | DataType::Timestamp(_, _)
- | DataType::Utf8
- | DataType::Binary
- | DataType::LargeUtf8
- | DataType::LargeBinary
- | DataType::Interval(_)
- | DataType::FixedSizeBinary(_) => vec![],
- DataType::List(_) | DataType::LargeList(_) => {
- let childs = arrays
- .iter()
- .map(|array| &array.child_data()[0])
- .collect::<Vec<_>>();
-
- let capacities = if let Capacities::List(capacity, ref child_capacities) =
- capacities
- {
- array_capacity = capacity;
- child_capacities
- .clone()
- .map(|c| *c)
- .unwrap_or(Capacities::Array(array_capacity))
- } else {
- Capacities::Array(array_capacity)
- };
-
- vec![MutableArrayData::with_capacities(
- childs, use_nulls, capacities,
- )]
- }
- // the dictionary type just appends keys and clones the values.
- DataType::Dictionary(_, _) => vec![],
- DataType::Float16 => unreachable!(),
- DataType::Struct(fields) => match capacities {
- Capacities::Struct(capacity, Some(ref child_capacities)) => {
- array_capacity = capacity;
- (0..fields.len())
- .zip(child_capacities)
- .map(|(i, child_cap)| {
- let child_arrays = arrays
- .iter()
- .map(|array| &array.child_data()[i])
- .collect::<Vec<_>>();
- MutableArrayData::with_capacities(
- child_arrays,
- use_nulls,
- child_cap.clone(),
- )
- })
- .collect::<Vec<_>>()
- }
- Capacities::Struct(capacity, None) => {
- array_capacity = capacity;
- (0..fields.len())
- .map(|i| {
- let child_arrays = arrays
- .iter()
- .map(|array| &array.child_data()[i])
- .collect::<Vec<_>>();
- MutableArrayData::new(child_arrays, use_nulls, capacity)
- })
- .collect::<Vec<_>>()
- }
- _ => (0..fields.len())
- .map(|i| {
- let child_arrays = arrays
- .iter()
- .map(|array| &array.child_data()[i])
- .collect::<Vec<_>>();
- MutableArrayData::new(child_arrays, use_nulls, array_capacity)
- })
- .collect::<Vec<_>>(),
- },
- _ => {
- todo!("Take and filter operations still not supported for this datatype")
- }
- };
-
- let dictionary = match &data_type {
- DataType::Dictionary(_, _) => match arrays.len() {
- 0 => unreachable!(),
- 1 => Some(arrays[0].child_data()[0].clone()),
- _ => {
- if let Capacities::Dictionary(_, _) = capacities {
- panic!("dictionary capacity not yet supported")
- }
- // Concat dictionaries together
- let dictionaries: Vec<_> =
- arrays.iter().map(|array| &array.child_data()[0]).collect();
- let lengths: Vec<_> = dictionaries
- .iter()
- .map(|dictionary| dictionary.len())
- .collect();
- let capacity = lengths.iter().sum();
-
- let mut mutable =
- MutableArrayData::new(dictionaries, false, capacity);
-
- for (i, len) in lengths.iter().enumerate() {
- mutable.extend(i, 0, *len)
- }
-
- Some(mutable.freeze())
- }
- },
- _ => None,
- };
-
- let extend_nulls = build_extend_nulls(data_type);
-
- let extend_null_bits = arrays
- .iter()
- .map(|array| build_extend_null_bits(array, use_nulls))
- .collect();
-
- let null_buffer = if use_nulls {
- let null_bytes = bit_util::ceil(array_capacity, 8);
- MutableBuffer::from_len_zeroed(null_bytes)
- } else {
- // create 0 capacity mutable buffer with the intention that it won't be used
- MutableBuffer::with_capacity(0)
- };
-
- let extend_values = match &data_type {
- DataType::Dictionary(_, _) => {
- let mut next_offset = 0;
- let extend_values: Result<Vec<_>> = arrays
- .iter()
- .map(|array| {
- let offset = next_offset;
- next_offset += array.child_data()[0].len();
- build_extend_dictionary(array, offset, next_offset)
- .ok_or(ArrowError::DictionaryKeyOverflowError)
- })
- .collect();
-
- extend_values.expect("MutableArrayData::new is infallible")
- }
- _ => arrays.iter().map(|array| build_extend(array)).collect(),
- };
-
- let data = _MutableArrayData {
- data_type: data_type.clone(),
- len: 0,
- null_count: 0,
- null_buffer,
- buffer1,
- buffer2,
- child_data,
- };
- Self {
- arrays,
- data,
- dictionary,
- extend_values,
- extend_null_bits,
- extend_nulls,
- }
- }
-
- /// Extends this [MutableArrayData] with elements from the bounded [ArrayData] at `start`
- /// and for a size of `len`.
- /// # Panic
- /// This function panics if the range is out of bounds, i.e. if `start + len >= array.len()`.
- pub fn extend(&mut self, index: usize, start: usize, end: usize) {
- let len = end - start;
- (self.extend_null_bits[index])(&mut self.data, start, len);
- (self.extend_values[index])(&mut self.data, index, start, len);
- self.data.len += len;
- }
-
- /// Extends this [MutableArrayData] with null elements, disregarding the bound arrays
- pub fn extend_nulls(&mut self, len: usize) {
- // TODO: null_buffer should probably be extended here as well
- // otherwise is_valid() could later panic
- // add test to confirm
- self.data.null_count += len;
- (self.extend_nulls)(&mut self.data, len);
- self.data.len += len;
- }
-
- /// Returns the current length
- #[inline]
- pub fn len(&self) -> usize {
- self.data.len
- }
-
- /// Returns true if len is 0
- #[inline]
- pub fn is_empty(&self) -> bool {
- self.data.len == 0
- }
-
- /// Returns the current null count
- #[inline]
- pub fn null_count(&self) -> usize {
- self.data.null_count
- }
-
- /// Creates a [ArrayData] from the pushed regions up to this point, consuming `self`.
- pub fn freeze(self) -> ArrayData {
- self.data.freeze(self.dictionary).build()
- }
-
- /// Creates a [ArrayDataBuilder] from the pushed regions up to this point, consuming `self`.
- /// This is useful for extending the default behavior of MutableArrayData.
- pub fn into_builder(self) -> ArrayDataBuilder {
- self.data.freeze(self.dictionary)
- }
-}
-
-#[cfg(test)]
-mod tests {
- use std::{convert::TryFrom, sync::Arc};
-
- use super::*;
-
- use crate::{
- array::{
- Array, ArrayData, ArrayRef, BooleanArray, DictionaryArray,
- FixedSizeBinaryArray, Int16Array, Int16Type, Int32Array, Int64Array,
- Int64Builder, ListBuilder, NullArray, PrimitiveBuilder, StringArray,
- StringDictionaryBuilder, StructArray, UInt8Array,
- },
- buffer::Buffer,
- datatypes::Field,
- };
- use crate::{
- array::{ListArray, StringBuilder},
- error::Result,
- };
-
- /// tests extending from a primitive array w/ offset nor nulls
- #[test]
- fn test_primitive() {
- let b = UInt8Array::from(vec![Some(1), Some(2), Some(3)]);
- let arrays = vec![b.data()];
- let mut a = MutableArrayData::new(arrays, false, 3);
- a.extend(0, 0, 2);
- let result = a.freeze();
- let array = UInt8Array::from(result);
- let expected = UInt8Array::from(vec![Some(1), Some(2)]);
- assert_eq!(array, expected);
- }
-
- /// tests extending from a primitive array with offset w/ nulls
- #[test]
- fn test_primitive_offset() {
- let b = UInt8Array::from(vec![Some(1), Some(2), Some(3)]);
- let b = b.slice(1, 2);
- let arrays = vec![b.data()];
- let mut a = MutableArrayData::new(arrays, false, 2);
- a.extend(0, 0, 2);
- let result = a.freeze();
- let array = UInt8Array::from(result);
- let expected = UInt8Array::from(vec![Some(2), Some(3)]);
- assert_eq!(array, expected);
- }
-
- /// tests extending from a primitive array with offset and nulls
- #[test]
- fn test_primitive_null_offset() {
- let b = UInt8Array::from(vec![Some(1), None, Some(3)]);
- let b = b.slice(1, 2);
- let arrays = vec![b.data()];
- let mut a = MutableArrayData::new(arrays, false, 2);
- a.extend(0, 0, 2);
- let result = a.freeze();
- let array = UInt8Array::from(result);
- let expected = UInt8Array::from(vec![None, Some(3)]);
- assert_eq!(array, expected);
- }
-
- #[test]
- fn test_primitive_null_offset_nulls() {
- let b = UInt8Array::from(vec![Some(1), Some(2), Some(3)]);
- let b = b.slice(1, 2);
- let arrays = vec![b.data()];
- let mut a = MutableArrayData::new(arrays, true, 2);
- a.extend(0, 0, 2);
- a.extend_nulls(3);
- a.extend(0, 1, 2);
- let result = a.freeze();
- let array = UInt8Array::from(result);
- let expected =
- UInt8Array::from(vec![Some(2), Some(3), None, None, None, Some(3)]);
- assert_eq!(array, expected);
- }
-
- #[test]
- fn test_list_null_offset() -> Result<()> {
- let int_builder = Int64Builder::new(24);
- let mut builder = ListBuilder::<Int64Builder>::new(int_builder);
- builder.values().append_slice(&[1, 2, 3])?;
- builder.append(true)?;
- builder.values().append_slice(&[4, 5])?;
- builder.append(true)?;
- builder.values().append_slice(&[6, 7, 8])?;
- builder.append(true)?;
- let array = builder.finish();
- let arrays = vec![array.data()];
-
- let mut mutable = MutableArrayData::new(arrays, false, 0);
- mutable.extend(0, 0, 1);
-
- let result = mutable.freeze();
- let array = ListArray::from(result);
-
- let int_builder = Int64Builder::new(24);
- let mut builder = ListBuilder::<Int64Builder>::new(int_builder);
- builder.values().append_slice(&[1, 2, 3])?;
- builder.append(true)?;
- let expected = builder.finish();
-
- assert_eq!(array, expected);
-
- Ok(())
- }
-
- /// tests extending from a variable-sized (strings and binary) array w/ offset with nulls
- #[test]
- fn test_variable_sized_nulls() {
- let array = StringArray::from(vec![Some("a"), Some("bc"), None, Some("defh")]);
- let arrays = vec![array.data()];
-
- let mut mutable = MutableArrayData::new(arrays, false, 0);
-
- mutable.extend(0, 1, 3);
-
- let result = mutable.freeze();
- let result = StringArray::from(result);
-
- let expected = StringArray::from(vec![Some("bc"), None]);
- assert_eq!(result, expected);
- }
-
- /// tests extending from a variable-sized (strings and binary) array
- /// with an offset and nulls
- #[test]
- fn test_variable_sized_offsets() {
- let array = StringArray::from(vec![Some("a"), Some("bc"), None, Some("defh")]);
- let array = array.slice(1, 3);
-
- let arrays = vec![array.data()];
-
- let mut mutable = MutableArrayData::new(arrays, false, 0);
-
- mutable.extend(0, 0, 3);
-
- let result = mutable.freeze();
- let result = StringArray::from(result);
-
- let expected = StringArray::from(vec![Some("bc"), None, Some("defh")]);
- assert_eq!(result, expected);
- }
-
- #[test]
- fn test_string_offsets() {
- let array = StringArray::from(vec![Some("a"), Some("bc"), None, Some("defh")]);
- let array = array.slice(1, 3);
-
- let arrays = vec![array.data()];
-
- let mut mutable = MutableArrayData::new(arrays, false, 0);
-
- mutable.extend(0, 0, 3);
-
- let result = mutable.freeze();
- let result = StringArray::from(result);
-
- let expected = StringArray::from(vec![Some("bc"), None, Some("defh")]);
- assert_eq!(result, expected);
- }
-
- #[test]
- fn test_multiple_with_nulls() {
- let array1 = StringArray::from(vec!["hello", "world"]);
- let array2 = StringArray::from(vec![Some("1"), None]);
-
- let arrays = vec![array1.data(), array2.data()];
-
- let mut mutable = MutableArrayData::new(arrays, false, 5);
-
- mutable.extend(0, 0, 2);
- mutable.extend(1, 0, 2);
-
- let result = mutable.freeze();
- let result = StringArray::from(result);
-
- let expected =
- StringArray::from(vec![Some("hello"), Some("world"), Some("1"), None]);
- assert_eq!(result, expected);
- }
-
- #[test]
- fn test_string_null_offset_nulls() {
- let array = StringArray::from(vec![Some("a"), Some("bc"), None, Some("defh")]);
- let array = array.slice(1, 3);
-
- let arrays = vec![array.data()];
-
- let mut mutable = MutableArrayData::new(arrays, true, 0);
-
- mutable.extend(0, 1, 3);
- mutable.extend_nulls(1);
-
- let result = mutable.freeze();
- let result = StringArray::from(result);
-
- let expected = StringArray::from(vec![None, Some("defh"), None]);
- assert_eq!(result, expected);
- }
-
- #[test]
- fn test_bool() {
- let array = BooleanArray::from(vec![Some(false), Some(true), None, Some(false)]);
- let arrays = vec![array.data()];
-
- let mut mutable = MutableArrayData::new(arrays, false, 0);
-
- mutable.extend(0, 1, 3);
-
- let result = mutable.freeze();
- let result = BooleanArray::from(result);
-
- let expected = BooleanArray::from(vec![Some(true), None]);
- assert_eq!(result, expected);
- }
-
- #[test]
- fn test_null() {
- let array1 = NullArray::new(10);
- let array2 = NullArray::new(5);
- let arrays = vec![array1.data(), array2.data()];
-
- let mut mutable = MutableArrayData::new(arrays, false, 0);
-
- mutable.extend(0, 1, 3);
- mutable.extend(1, 0, 1);
-
- let result = mutable.freeze();
- let result = NullArray::from(result);
-
- let expected = NullArray::new(3);
- assert_eq!(result, expected);
- }
-
- fn create_dictionary_array(values: &[&str], keys: &[Option<&str>]) -> ArrayData {
- let values = StringArray::from(values.to_vec());
- let mut builder = StringDictionaryBuilder::new_with_dictionary(
- PrimitiveBuilder::<Int16Type>::new(3),
- &values,
- )
- .unwrap();
- for key in keys {
- if let Some(v) = key {
- builder.append(v).unwrap();
- } else {
- builder.append_null().unwrap()
- }
- }
- builder.finish().data().clone()
- }
-
- #[test]
- fn test_dictionary() {
- // (a, b, c), (0, 1, 0, 2) => (a, b, a, c)
- let array = create_dictionary_array(
- &["a", "b", "c"],
- &[Some("a"), Some("b"), None, Some("c")],
- );
- let arrays = vec![&array];
-
- let mut mutable = MutableArrayData::new(arrays, false, 0);
-
- mutable.extend(0, 1, 3);
-
- let result = mutable.freeze();
- let result = DictionaryArray::from(result);
-
- let expected = Int16Array::from(vec![Some(1), None]);
- assert_eq!(result.keys(), &expected);
- }
-
- #[test]
- fn test_struct() {
- let strings: ArrayRef = Arc::new(StringArray::from(vec![
- Some("joe"),
- None,
- None,
- Some("mark"),
- Some("doe"),
- ]));
- let ints: ArrayRef = Arc::new(Int32Array::from(vec![
- Some(1),
- Some(2),
- Some(3),
- Some(4),
- Some(5),
- ]));
-
- let array =
- StructArray::try_from(vec![("f1", strings.clone()), ("f2", ints.clone())])
- .unwrap();
- let arrays = vec![array.data()];
- let mut mutable = MutableArrayData::new(arrays, false, 0);
-
- mutable.extend(0, 1, 3);
- let data = mutable.freeze();
- let array = StructArray::from(data);
-
- let expected = StructArray::try_from(vec![
- ("f1", strings.slice(1, 2)),
- ("f2", ints.slice(1, 2)),
- ])
- .unwrap();
- assert_eq!(array, expected)
- }
-
- #[test]
- fn test_struct_offset() {
- let strings: ArrayRef = Arc::new(StringArray::from(vec![
- Some("joe"),
- None,
- None,
- Some("mark"),
- Some("doe"),
- ]));
- let ints: ArrayRef = Arc::new(Int32Array::from(vec![
- Some(1),
- Some(2),
- Some(3),
- Some(4),
- Some(5),
- ]));
-
- let array =
- StructArray::try_from(vec![("f1", strings.clone()), ("f2", ints.clone())])
- .unwrap()
- .slice(1, 3);
- let arrays = vec![array.data()];
- let mut mutable = MutableArrayData::new(arrays, false, 0);
-
- mutable.extend(0, 1, 3);
- let data = mutable.freeze();
- let array = StructArray::from(data);
-
- let expected_strings: ArrayRef =
- Arc::new(StringArray::from(vec![None, Some("mark")]));
- let expected = StructArray::try_from(vec![
- ("f1", expected_strings),
- ("f2", ints.slice(2, 2)),
- ])
- .unwrap();
-
- assert_eq!(array, expected);
- }
-
- #[test]
- fn test_struct_nulls() {
- let strings: ArrayRef = Arc::new(StringArray::from(vec![
- Some("joe"),
- None,
- None,
- Some("mark"),
- Some("doe"),
- ]));
- let ints: ArrayRef = Arc::new(Int32Array::from(vec![
- Some(1),
- Some(2),
- None,
- Some(4),
- Some(5),
- ]));
-
- let array =
- StructArray::try_from(vec![("f1", strings.clone()), ("f2", ints.clone())])
- .unwrap();
- let arrays = vec![array.data()];
-
- let mut mutable = MutableArrayData::new(arrays, false, 0);
-
- mutable.extend(0, 1, 3);
- let data = mutable.freeze();
- let array = StructArray::from(data);
-
- let expected_string = Arc::new(StringArray::from(vec![None, None])) as ArrayRef;
- let expected_int = Arc::new(Int32Array::from(vec![Some(2), None])) as ArrayRef;
-
- let expected =
- StructArray::try_from(vec![("f1", expected_string), ("f2", expected_int)])
- .unwrap();
- assert_eq!(array, expected)
- }
-
- #[test]
- fn test_struct_many() {
- let strings: ArrayRef = Arc::new(StringArray::from(vec![
- Some("joe"),
- None,
- None,
- Some("mark"),
- Some("doe"),
- ]));
- let ints: ArrayRef = Arc::new(Int32Array::from(vec![
- Some(1),
- Some(2),
- None,
- Some(4),
- Some(5),
- ]));
-
- let array =
- StructArray::try_from(vec![("f1", strings.clone()), ("f2", ints.clone())])
- .unwrap();
- let arrays = vec![array.data(), array.data()];
- let mut mutable = MutableArrayData::new(arrays, false, 0);
-
- mutable.extend(0, 1, 3);
- mutable.extend(1, 0, 2);
- let data = mutable.freeze();
- let array = StructArray::from(data);
-
- let expected_string =
- Arc::new(StringArray::from(vec![None, None, Some("joe"), None])) as ArrayRef;
- let expected_int =
- Arc::new(Int32Array::from(vec![Some(2), None, Some(1), Some(2)])) as ArrayRef;
-
- let expected =
- StructArray::try_from(vec![("f1", expected_string), ("f2", expected_int)])
- .unwrap();
- assert_eq!(array, expected)
- }
-
- #[test]
- fn test_binary_fixed_sized_offsets() {
- let array = FixedSizeBinaryArray::try_from_iter(
- vec![vec![0, 0], vec![0, 1], vec![0, 2]].into_iter(),
- )
- .expect("Failed to create FixedSizeBinaryArray from iterable");
- let array = array.slice(1, 2);
- // = [[0, 1], [0, 2]] due to the offset = 1
-
- let arrays = vec![array.data()];
-
- let mut mutable = MutableArrayData::new(arrays, false, 0);
-
- mutable.extend(0, 1, 2);
- mutable.extend(0, 0, 1);
-
- let result = mutable.freeze();
- let result = FixedSizeBinaryArray::from(result);
-
- let expected =
- FixedSizeBinaryArray::try_from_iter(vec![vec![0, 2], vec![0, 1]].into_iter())
- .expect("Failed to create FixedSizeBinaryArray from iterable");
- assert_eq!(result, expected);
- }
-
- #[test]
- fn test_list_append() -> Result<()> {
- let mut builder = ListBuilder::<Int64Builder>::new(Int64Builder::new(24));
- builder.values().append_slice(&[1, 2, 3])?;
- builder.append(true)?;
- builder.values().append_slice(&[4, 5])?;
- builder.append(true)?;
- builder.values().append_slice(&[6, 7, 8])?;
- builder.values().append_slice(&[9, 10, 11])?;
- builder.append(true)?;
- let a = builder.finish();
-
- let a_builder = Int64Builder::new(24);
- let mut a_builder = ListBuilder::<Int64Builder>::new(a_builder);
- a_builder.values().append_slice(&[12, 13])?;
- a_builder.append(true)?;
- a_builder.append(true)?;
- a_builder.values().append_slice(&[14, 15])?;
- a_builder.append(true)?;
- let b = a_builder.finish();
-
- let c = b.slice(1, 2);
-
- let mut mutable =
- MutableArrayData::new(vec![a.data(), b.data(), c.data()], false, 1);
- mutable.extend(0, 0, a.len());
- mutable.extend(1, 0, b.len());
- mutable.extend(2, 0, c.len());
-
- let finished = mutable.freeze();
-
- let expected_int_array = Int64Array::from(vec![
- Some(1),
- Some(2),
- Some(3),
- Some(4),
- Some(5),
- Some(6),
- Some(7),
- Some(8),
- Some(9),
- Some(10),
- Some(11),
- // append first array
- Some(12),
- Some(13),
- Some(14),
- Some(15),
- // append second array
- Some(14),
- Some(15),
- ]);
- let list_value_offsets =
- Buffer::from_slice_ref(&[0i32, 3, 5, 11, 13, 13, 15, 15, 17]);
- let expected_list_data = ArrayData::new(
- DataType::List(Box::new(Field::new("item", DataType::Int64, true))),
- 8,
- None,
- None,
- 0,
- vec![list_value_offsets],
- vec![expected_int_array.data().clone()],
- );
- assert_eq!(finished, expected_list_data);
-
- Ok(())
- }
-
- #[test]
- fn test_list_nulls_append() -> Result<()> {
- let mut builder = ListBuilder::<Int64Builder>::new(Int64Builder::new(32));
- builder.values().append_slice(&[1, 2, 3])?;
- builder.append(true)?;
- builder.values().append_slice(&[4, 5])?;
- builder.append(true)?;
- builder.append(false)?;
- builder.values().append_slice(&[6, 7, 8])?;
- builder.values().append_null()?;
- builder.values().append_null()?;
- builder.values().append_slice(&[9, 10, 11])?;
- builder.append(true)?;
- let a = builder.finish();
- let a = a.data();
-
- let mut builder = ListBuilder::<Int64Builder>::new(Int64Builder::new(32));
- builder.values().append_slice(&[12, 13])?;
- builder.append(true)?;
- builder.append(false)?;
- builder.append(true)?;
- builder.values().append_null()?;
- builder.values().append_null()?;
- builder.values().append_slice(&[14, 15])?;
- builder.append(true)?;
- let b = builder.finish();
- let b = b.data();
- let c = b.slice(1, 2);
- let d = b.slice(2, 2);
-
- let mut mutable = MutableArrayData::new(vec![a, b, &c, &d], false, 10);
-
- mutable.extend(0, 0, a.len());
- mutable.extend(1, 0, b.len());
- mutable.extend(2, 0, c.len());
- mutable.extend(3, 0, d.len());
- let result = mutable.freeze();
-
- let expected_int_array = Int64Array::from(vec![
- Some(1),
- Some(2),
- Some(3),
- Some(4),
- Some(5),
- Some(6),
- Some(7),
- Some(8),
- None,
- None,
- Some(9),
- Some(10),
- Some(11),
- // second array
- Some(12),
- Some(13),
- None,
- None,
- Some(14),
- Some(15),
- // slice(1, 2) results in no values added
- None,
- None,
- Some(14),
- Some(15),
- ]);
- let list_value_offsets =
- Buffer::from_slice_ref(&[0, 3, 5, 5, 13, 15, 15, 15, 19, 19, 19, 19, 23]);
- let expected_list_data = ArrayData::new(
- DataType::List(Box::new(Field::new("item", DataType::Int64, true))),
- 12,
- None,
- Some(Buffer::from(&[0b11011011, 0b1110])),
- 0,
- vec![list_value_offsets],
- vec![expected_int_array.data().clone()],
- );
- assert_eq!(result, expected_list_data);
-
- Ok(())
- }
-
- #[test]
- fn test_list_of_strings_append() -> Result<()> {
- // [["alpha", "beta", None]]
- let mut builder = ListBuilder::new(StringBuilder::new(32));
- builder.values().append_value("Hello")?;
- builder.values().append_value("Arrow")?;
- builder.values().append_null()?;
- builder.append(true)?;
- let a = builder.finish();
-
- // [["alpha", "beta"], [None], ["gamma", "delta", None]]
- let mut builder = ListBuilder::new(StringBuilder::new(32));
- builder.values().append_value("alpha")?;
- builder.values().append_value("beta")?;
- builder.append(true)?;
- builder.values().append_null()?;
- builder.append(true)?;
- builder.values().append_value("gamma")?;
- builder.values().append_value("delta")?;
- builder.values().append_null()?;
- builder.append(true)?;
- let b = builder.finish();
-
- let mut mutable = MutableArrayData::new(vec![a.data(), b.data()], false, 10);
-
- mutable.extend(0, 0, a.len());
- mutable.extend(1, 0, b.len());
- mutable.extend(1, 1, 3);
- mutable.extend(1, 0, 0);
- let result = mutable.freeze();
-
- let expected_string_array = StringArray::from(vec![
- // extend a[0..a.len()]
- // a[0]
- Some("Hello"),
- Some("Arrow"),
- None,
- // extend b[0..b.len()]
- // b[0]
- Some("alpha"),
- Some("beta"),
- // b[1]
- None,
- // b[2]
- Some("gamma"),
- Some("delta"),
- None,
- // extend b[1..3]
- // b[1]
- None,
- // b[2]
- Some("gamma"),
- Some("delta"),
- None,
- // extend b[0..0]
- ]);
- let list_value_offsets = Buffer::from_slice_ref(&[0, 3, 5, 6, 9, 10, 13]);
- let expected_list_data = ArrayData::new(
- DataType::List(Box::new(Field::new("item", DataType::Utf8, true))),
- 6,
- None,
- None,
- 0,
- vec![list_value_offsets],
- vec![expected_string_array.data().clone()],
- );
- assert_eq!(result, expected_list_data);
- Ok(())
- }
-
- #[test]
- fn test_fixed_size_binary_append() {
- let a = vec![Some(vec![1, 2]), Some(vec![3, 4]), Some(vec![5, 6])];
- let a = FixedSizeBinaryArray::try_from_sparse_iter(a.into_iter())
- .expect("Failed to create FixedSizeBinaryArray from iterable");
-
- let b = vec![
- None,
- Some(vec![7, 8]),
- Some(vec![9, 10]),
- None,
- Some(vec![13, 14]),
- None,
- ];
- let b = FixedSizeBinaryArray::try_from_sparse_iter(b.into_iter())
- .expect("Failed to create FixedSizeBinaryArray from iterable");
-
- let mut mutable = MutableArrayData::new(vec![a.data(), b.data()], false, 10);
-
- mutable.extend(0, 0, a.len());
- mutable.extend(1, 0, b.len());
- mutable.extend(1, 1, 4);
- mutable.extend(1, 2, 3);
- mutable.extend(1, 5, 5);
- let result = mutable.freeze();
-
- let expected = vec![
- // a
- Some(vec![1, 2]),
- Some(vec![3, 4]),
- Some(vec![5, 6]),
- // b
- None,
- Some(vec![7, 8]),
- Some(vec![9, 10]),
- None,
- Some(vec![13, 14]),
- None,
- // b[1..4]
- Some(vec![7, 8]),
- Some(vec![9, 10]),
- None,
- // b[2..3]
- Some(vec![9, 10]),
- // b[4..4]
- ];
- let expected = FixedSizeBinaryArray::try_from_sparse_iter(expected.into_iter())
- .expect("Failed to create FixedSizeBinaryArray from iterable");
- assert_eq!(&result, expected.data());
- }
-
- /*
- // this is an old test used on a meanwhile removed dead code
- // that is still useful when `MutableArrayData` supports fixed-size lists.
- #[test]
- fn test_fixed_size_list_append() -> Result<()> {
- let int_builder = UInt16Builder::new(64);
- let mut builder = FixedSizeListBuilder::<UInt16Builder>::new(int_builder, 2);
- builder.values().append_slice(&[1, 2])?;
- builder.append(true)?;
- builder.values().append_slice(&[3, 4])?;
- builder.append(false)?;
- builder.values().append_slice(&[5, 6])?;
- builder.append(true)?;
-
- let a_builder = UInt16Builder::new(64);
- let mut a_builder = FixedSizeListBuilder::<UInt16Builder>::new(a_builder, 2);
- a_builder.values().append_slice(&[7, 8])?;
- a_builder.append(true)?;
- a_builder.values().append_slice(&[9, 10])?;
- a_builder.append(true)?;
- a_builder.values().append_slice(&[11, 12])?;
- a_builder.append(false)?;
- a_builder.values().append_slice(&[13, 14])?;
- a_builder.append(true)?;
- a_builder.values().append_null()?;
- a_builder.values().append_null()?;
- a_builder.append(true)?;
- let a = a_builder.finish();
-
- // append array
- builder.append_data(&[
- a.data(),
- a.slice(1, 3).data(),
- a.slice(2, 1).data(),
- a.slice(5, 0).data(),
- ])?;
- let finished = builder.finish();
-
- let expected_int_array = UInt16Array::from(vec![
- Some(1),
- Some(2),
- Some(3),
- Some(4),
- Some(5),
- Some(6),
- // append first array
- Some(7),
- Some(8),
- Some(9),
- Some(10),
- Some(11),
- Some(12),
- Some(13),
- Some(14),
- None,
- None,
- // append slice(1, 3)
- Some(9),
- Some(10),
- Some(11),
- Some(12),
- Some(13),
- Some(14),
- // append slice(2, 1)
- Some(11),
- Some(12),
- ]);
- let expected_list_data = ArrayData::new(
- DataType::FixedSizeList(
- Box::new(Field::new("item", DataType::UInt16, true)),
- 2,
- ),
- 12,
- None,
- None,
- 0,
- vec![],
- vec![expected_int_array.data()],
- );
- let expected_list =
- FixedSizeListArray::from(Arc::new(expected_list_data) as ArrayData);
- assert_eq!(&expected_list.values(), &finished.values());
- assert_eq!(expected_list.len(), finished.len());
-
- Ok(())
- }
- */
-}
diff --git a/arrow/src/array/transform/null.rs b/arrow/src/array/transform/null.rs
deleted file mode 100644
index e1335e1..0000000
--- a/arrow/src/array/transform/null.rs
+++ /dev/null
@@ -1,26 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-use crate::array::ArrayData;
-
-use super::{Extend, _MutableArrayData};
-
-pub(super) fn build_extend(_: &ArrayData) -> Extend {
- Box::new(move |_, _, _, _| {})
-}
-
-pub(super) fn extend_nulls(_: &mut _MutableArrayData, _: usize) {}
diff --git a/arrow/src/array/transform/primitive.rs b/arrow/src/array/transform/primitive.rs
deleted file mode 100644
index 4c765c0..0000000
--- a/arrow/src/array/transform/primitive.rs
+++ /dev/null
@@ -1,55 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-use std::mem::size_of;
-use std::ops::Add;
-
-use crate::{array::ArrayData, datatypes::ArrowNativeType};
-
-use super::{Extend, _MutableArrayData};
-
-pub(super) fn build_extend<T: ArrowNativeType>(array: &ArrayData) -> Extend {
- let values = array.buffer::<T>(0);
- Box::new(
- move |mutable: &mut _MutableArrayData, _, start: usize, len: usize| {
- mutable
- .buffer1
- .extend_from_slice(&values[start..start + len]);
- },
- )
-}
-
-pub(super) fn build_extend_with_offset<T>(array: &ArrayData, offset: T) -> Extend
-where
- T: ArrowNativeType + Add<Output = T>,
-{
- let values = array.buffer::<T>(0);
- Box::new(
- move |mutable: &mut _MutableArrayData, _, start: usize, len: usize| {
- mutable
- .buffer1
- .extend(values[start..start + len].iter().map(|x| *x + offset));
- },
- )
-}
-
-pub(super) fn extend_nulls<T: ArrowNativeType>(
- mutable: &mut _MutableArrayData,
- len: usize,
-) {
- mutable.buffer1.extend_zeros(len * size_of::<T>());
-}
diff --git a/arrow/src/array/transform/structure.rs b/arrow/src/array/transform/structure.rs
deleted file mode 100644
index c019f5a..0000000
--- a/arrow/src/array/transform/structure.rs
+++ /dev/null
@@ -1,67 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-use crate::array::ArrayData;
-
-use super::{Extend, _MutableArrayData};
-
-pub(super) fn build_extend(array: &ArrayData) -> Extend {
- if array.null_count() == 0 {
- Box::new(
- move |mutable: &mut _MutableArrayData,
- index: usize,
- start: usize,
- len: usize| {
- mutable.child_data.iter_mut().for_each(|child| {
- child.extend(
- index,
- array.offset() + start,
- array.offset() + start + len,
- )
- })
- },
- )
- } else {
- Box::new(
- move |mutable: &mut _MutableArrayData,
- index: usize,
- start: usize,
- len: usize| {
- (array.offset() + start..array.offset() + start + len).for_each(|i| {
- if array.is_valid(i) {
- mutable
- .child_data
- .iter_mut()
- .for_each(|child| child.extend(index, i, i + 1))
- } else {
- mutable
- .child_data
- .iter_mut()
- .for_each(|child| child.extend_nulls(1))
- }
- })
- },
- )
- }
-}
-
-pub(super) fn extend_nulls(mutable: &mut _MutableArrayData, len: usize) {
- mutable
- .child_data
- .iter_mut()
- .for_each(|child| child.extend_nulls(len))
-}
diff --git a/arrow/src/array/transform/utils.rs b/arrow/src/array/transform/utils.rs
deleted file mode 100644
index 8c718c7..0000000
--- a/arrow/src/array/transform/utils.rs
+++ /dev/null
@@ -1,76 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-use crate::{array::OffsetSizeTrait, buffer::MutableBuffer, util::bit_util};
-
-/// extends the `buffer` to be able to hold `len` bits, setting all bits of the new size to zero.
-#[inline]
-pub(super) fn resize_for_bits(buffer: &mut MutableBuffer, len: usize) {
- let needed_bytes = bit_util::ceil(len, 8);
- if buffer.len() < needed_bytes {
- buffer.resize(needed_bytes, 0);
- }
-}
-
-/// sets all bits on `write_data` on the range `[offset_write..offset_write+len]` to be equal to the
-/// bits on `data` on the range `[offset_read..offset_read+len]`
-pub(super) fn set_bits(
- write_data: &mut [u8],
- data: &[u8],
- offset_write: usize,
- offset_read: usize,
- len: usize,
-) -> usize {
- let mut count = 0;
- (0..len).for_each(|i| {
- if bit_util::get_bit(data, offset_read + i) {
- bit_util::set_bit(write_data, offset_write + i);
- } else {
- count += 1;
- }
- });
- count
-}
-
-pub(super) fn extend_offsets<T: OffsetSizeTrait>(
- buffer: &mut MutableBuffer,
- mut last_offset: T,
- offsets: &[T],
-) {
- buffer.reserve(offsets.len() * std::mem::size_of::<T>());
- offsets.windows(2).for_each(|offsets| {
- // compute the new offset
- let length = offsets[1] - offsets[0];
- last_offset += length;
- buffer.push(last_offset);
- });
-}
-
-#[inline]
-pub(super) unsafe fn get_last_offset<T: OffsetSizeTrait>(
- offset_buffer: &MutableBuffer,
-) -> T {
- // JUSTIFICATION
- // Benefit
- // 20% performance improvement extend of variable sized arrays (see bench `mutable_array`)
- // Soundness
- // * offset buffer is always extended in slices of T and aligned accordingly.
- // * Buffer[0] is initialized with one element, 0, and thus `mutable_offsets.len() - 1` is always valid.
- let (prefix, offsets, suffix) = offset_buffer.as_slice().align_to::<T>();
- debug_assert!(prefix.is_empty() && suffix.is_empty());
- *offsets.get_unchecked(offsets.len() - 1)
-}
diff --git a/arrow/src/array/transform/variable_size.rs b/arrow/src/array/transform/variable_size.rs
deleted file mode 100644
index c9304db..0000000
--- a/arrow/src/array/transform/variable_size.rs
+++ /dev/null
@@ -1,105 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-use crate::{
- array::{ArrayData, OffsetSizeTrait},
- buffer::MutableBuffer,
-};
-
-use super::{
- Extend, _MutableArrayData,
- utils::{extend_offsets, get_last_offset},
-};
-
-#[inline]
-fn extend_offset_values<T: OffsetSizeTrait>(
- buffer: &mut MutableBuffer,
- offsets: &[T],
- values: &[u8],
- start: usize,
- len: usize,
-) {
- let start_values = offsets[start].to_usize().unwrap();
- let end_values = offsets[start + len].to_usize().unwrap();
- let new_values = &values[start_values..end_values];
- buffer.extend_from_slice(new_values);
-}
-
-pub(super) fn build_extend<T: OffsetSizeTrait>(array: &ArrayData) -> Extend {
- let offsets = array.buffer::<T>(0);
- let values = array.buffers()[1].as_slice();
- if array.null_count() == 0 {
- // fast case where we can copy regions without null issues
- Box::new(
- move |mutable: &mut _MutableArrayData, _, start: usize, len: usize| {
- let offset_buffer = &mut mutable.buffer1;
- let values_buffer = &mut mutable.buffer2;
-
- // this is safe due to how offset is built. See details on `get_last_offset`
- let last_offset = unsafe { get_last_offset(offset_buffer) };
-
- extend_offsets::<T>(
- offset_buffer,
- last_offset,
- &offsets[start..start + len + 1],
- );
- // values
- extend_offset_values::<T>(values_buffer, offsets, values, start, len);
- },
- )
- } else {
- Box::new(
- move |mutable: &mut _MutableArrayData, _, start: usize, len: usize| {
- let offset_buffer = &mut mutable.buffer1;
- let values_buffer = &mut mutable.buffer2;
-
- // this is safe due to how offset is built. See details on `get_last_offset`
- let mut last_offset: T = unsafe { get_last_offset(offset_buffer) };
-
- // nulls present: append item by item, ignoring null entries
- offset_buffer.reserve(len * std::mem::size_of::<T>());
-
- (start..start + len).for_each(|i| {
- if array.is_valid(i) {
- // compute the new offset
- let length = offsets[i + 1] - offsets[i];
- last_offset += length;
-
- // append value
- let bytes = &values[offsets[i].to_usize().unwrap()
- ..offsets[i + 1].to_usize().unwrap()];
- values_buffer.extend_from_slice(bytes);
- }
- // offsets are always present
- offset_buffer.push(last_offset);
- })
- },
- )
- }
-}
-
-pub(super) fn extend_nulls<T: OffsetSizeTrait>(
- mutable: &mut _MutableArrayData,
- len: usize,
-) {
- let offset_buffer = &mut mutable.buffer1;
-
- // this is safe due to how offset is built. See details on `get_last_offset`
- let last_offset: T = unsafe { get_last_offset(offset_buffer) };
-
- (0..len).for_each(|_| offset_buffer.push(last_offset))
-}
diff --git a/arrow/src/bitmap.rs b/arrow/src/bitmap.rs
deleted file mode 100644
index b977f55..0000000
--- a/arrow/src/bitmap.rs
+++ /dev/null
@@ -1,157 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-//! Defines a bitmap, which is used to track which values in an Arrow array are null.
-//! This is called a "validity bitmap" in the Arrow documentation.
-
-use crate::buffer::Buffer;
-use crate::error::Result;
-use crate::util::bit_util;
-use std::mem;
-
-use std::ops::{BitAnd, BitOr};
-
-#[derive(Debug, Clone)]
-pub struct Bitmap {
- pub(crate) bits: Buffer,
-}
-
-impl Bitmap {
- pub fn new(num_bits: usize) -> Self {
- let num_bytes = num_bits / 8 + if num_bits % 8 > 0 { 1 } else { 0 };
- let r = num_bytes % 64;
- let len = if r == 0 {
- num_bytes
- } else {
- num_bytes + 64 - r
- };
- Bitmap {
- bits: Buffer::from(&vec![0xFF; len]),
- }
- }
-
- pub fn len(&self) -> usize {
- self.bits.len()
- }
-
- pub fn is_empty(&self) -> bool {
- self.bits.is_empty()
- }
-
- pub fn is_set(&self, i: usize) -> bool {
- assert!(i < (self.bits.len() << 3));
- unsafe { bit_util::get_bit_raw(self.bits.as_ptr(), i) }
- }
-
- pub fn buffer_ref(&self) -> &Buffer {
- &self.bits
- }
-
- pub fn into_buffer(self) -> Buffer {
- self.bits
- }
-
- /// Returns the total number of bytes of memory occupied by the buffers owned by this [Bitmap].
- pub fn get_buffer_memory_size(&self) -> usize {
- self.bits.capacity()
- }
-
- /// Returns the total number of bytes of memory occupied physically by this [Bitmap].
- pub fn get_array_memory_size(&self) -> usize {
- self.bits.capacity() + mem::size_of_val(self)
- }
-}
-
-impl<'a, 'b> BitAnd<&'b Bitmap> for &'a Bitmap {
- type Output = Result<Bitmap>;
-
- fn bitand(self, rhs: &'b Bitmap) -> Result<Bitmap> {
- Ok(Bitmap::from((&self.bits & &rhs.bits)?))
- }
-}
-
-impl<'a, 'b> BitOr<&'b Bitmap> for &'a Bitmap {
- type Output = Result<Bitmap>;
-
- fn bitor(self, rhs: &'b Bitmap) -> Result<Bitmap> {
- Ok(Bitmap::from((&self.bits | &rhs.bits)?))
- }
-}
-
-impl From<Buffer> for Bitmap {
- fn from(buf: Buffer) -> Self {
- Self { bits: buf }
- }
-}
-
-impl PartialEq for Bitmap {
- fn eq(&self, other: &Self) -> bool {
- // buffer equality considers capacity, but here we want to only compare
- // actual data contents
- let self_len = self.bits.len();
- let other_len = other.bits.len();
- if self_len != other_len {
- return false;
- }
- self.bits.as_slice()[..self_len] == other.bits.as_slice()[..self_len]
- }
-}
-
-#[cfg(test)]
-mod tests {
- use super::*;
-
- #[test]
- fn test_bitmap_length() {
- assert_eq!(64, Bitmap::new(63 * 8).len());
- assert_eq!(64, Bitmap::new(64 * 8).len());
- assert_eq!(128, Bitmap::new(65 * 8).len());
- }
-
- #[test]
- fn test_bitwise_and() {
- let bitmap1 = Bitmap::from(Buffer::from([0b01101010]));
- let bitmap2 = Bitmap::from(Buffer::from([0b01001110]));
- assert_eq!(
- Bitmap::from(Buffer::from([0b01001010])),
- (&bitmap1 & &bitmap2).unwrap()
- );
- }
-
- #[test]
- fn test_bitwise_or() {
- let bitmap1 = Bitmap::from(Buffer::from([0b01101010]));
- let bitmap2 = Bitmap::from(Buffer::from([0b01001110]));
- assert_eq!(
- Bitmap::from(Buffer::from([0b01101110])),
- (&bitmap1 | &bitmap2).unwrap()
- );
- }
-
- #[test]
- fn test_bitmap_is_set() {
- let bitmap = Bitmap::from(Buffer::from([0b01001010]));
- assert_eq!(false, bitmap.is_set(0));
- assert_eq!(true, bitmap.is_set(1));
- assert_eq!(false, bitmap.is_set(2));
- assert_eq!(true, bitmap.is_set(3));
- assert_eq!(false, bitmap.is_set(4));
- assert_eq!(false, bitmap.is_set(5));
- assert_eq!(true, bitmap.is_set(6));
- assert_eq!(false, bitmap.is_set(7));
- }
-}
diff --git a/arrow/src/buffer/immutable.rs b/arrow/src/buffer/immutable.rs
deleted file mode 100644
index f4aeae9..0000000
--- a/arrow/src/buffer/immutable.rs
+++ /dev/null
@@ -1,541 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-use std::fmt::Debug;
-use std::iter::FromIterator;
-use std::ptr::NonNull;
-use std::sync::Arc;
-use std::{convert::AsRef, usize};
-
-use crate::util::bit_chunk_iterator::BitChunks;
-use crate::{
- bytes::{Bytes, Deallocation},
- datatypes::ArrowNativeType,
- ffi,
-};
-
-use super::ops::bitwise_unary_op_helper;
-use super::MutableBuffer;
-
-/// Buffer represents a contiguous memory region that can be shared with other buffers and across
-/// thread boundaries.
-#[derive(Clone, PartialEq, Debug)]
-pub struct Buffer {
- /// the internal byte buffer.
- data: Arc<Bytes>,
-
- /// The offset into the buffer.
- offset: usize,
-}
-
-impl Buffer {
- /// Auxiliary method to create a new Buffer
- #[inline]
- pub fn from_bytes(bytes: Bytes) -> Self {
- Buffer {
- data: Arc::new(bytes),
- offset: 0,
- }
- }
-
- /// Initializes a [Buffer] from a slice of items.
- pub fn from_slice_ref<U: ArrowNativeType, T: AsRef<[U]>>(items: &T) -> Self {
- let slice = items.as_ref();
- let capacity = slice.len() * std::mem::size_of::<U>();
- let mut buffer = MutableBuffer::with_capacity(capacity);
- buffer.extend_from_slice(slice);
- buffer.into()
- }
-
- /// Creates a buffer from an existing memory region (must already be byte-aligned), this
- /// `Buffer` will free this piece of memory when dropped.
- ///
- /// # Arguments
- ///
- /// * `ptr` - Pointer to raw parts
- /// * `len` - Length of raw parts in **bytes**
- /// * `capacity` - Total allocated memory for the pointer `ptr`, in **bytes**
- ///
- /// # Safety
- ///
- /// This function is unsafe as there is no guarantee that the given pointer is valid for `len`
- /// bytes. If the `ptr` and `capacity` come from a `Buffer`, then this is guaranteed.
- pub unsafe fn from_raw_parts(ptr: NonNull<u8>, len: usize, capacity: usize) -> Self {
- assert!(len <= capacity);
- Buffer::build_with_arguments(ptr, len, Deallocation::Native(capacity))
- }
-
- /// Creates a buffer from an existing memory region (must already be byte-aligned), this
- /// `Buffer` **does not** free this piece of memory when dropped.
- ///
- /// # Arguments
- ///
- /// * `ptr` - Pointer to raw parts
- /// * `len` - Length of raw parts in **bytes**
- /// * `data` - An [ffi::FFI_ArrowArray] with the data
- ///
- /// # Safety
- ///
- /// This function is unsafe as there is no guarantee that the given pointer is valid for `len`
- /// bytes and that the foreign deallocator frees the region.
- pub unsafe fn from_unowned(
- ptr: NonNull<u8>,
- len: usize,
- data: Arc<ffi::FFI_ArrowArray>,
- ) -> Self {
- Buffer::build_with_arguments(ptr, len, Deallocation::Foreign(data))
- }
-
- /// Auxiliary method to create a new Buffer
- unsafe fn build_with_arguments(
- ptr: NonNull<u8>,
- len: usize,
- deallocation: Deallocation,
- ) -> Self {
- let bytes = Bytes::new(ptr, len, deallocation);
- Buffer {
- data: Arc::new(bytes),
- offset: 0,
- }
- }
-
- /// Returns the number of bytes in the buffer
- pub fn len(&self) -> usize {
- self.data.len() - self.offset
- }
-
- /// Returns the capacity of this buffer.
- /// For exernally owned buffers, this returns zero
- pub fn capacity(&self) -> usize {
- self.data.capacity()
- }
-
- /// Returns whether the buffer is empty.
- pub fn is_empty(&self) -> bool {
- self.data.len() - self.offset == 0
- }
-
- /// Returns the byte slice stored in this buffer
- pub fn as_slice(&self) -> &[u8] {
- &self.data[self.offset..]
- }
-
- /// Returns a new [Buffer] that is a slice of this buffer starting at `offset`.
- /// Doing so allows the same memory region to be shared between buffers.
- /// # Panics
- /// Panics iff `offset` is larger than `len`.
- pub fn slice(&self, offset: usize) -> Self {
- assert!(
- offset <= self.len(),
- "the offset of the new Buffer cannot exceed the existing length"
- );
- Self {
- data: self.data.clone(),
- offset: self.offset + offset,
- }
- }
-
- /// Returns a pointer to the start of this buffer.
- ///
- /// Note that this should be used cautiously, and the returned pointer should not be
- /// stored anywhere, to avoid dangling pointers.
- pub fn as_ptr(&self) -> *const u8 {
- unsafe { self.data.ptr().as_ptr().add(self.offset) }
- }
-
- /// View buffer as typed slice.
- ///
- /// # Safety
- ///
- /// `ArrowNativeType` is public so that it can be used as a trait bound for other public
- /// components, such as the `ToByteSlice` trait. However, this means that it can be
- /// implemented by user defined types, which it is not intended for.
- ///
- /// Also `typed_data::<bool>` is unsafe as `0x00` and `0x01` are the only valid values for
- /// `bool` in Rust. However, `bool` arrays in Arrow are bit-packed which breaks this condition.
- /// View buffer as typed slice.
- pub unsafe fn typed_data<T: ArrowNativeType + num::Num>(&self) -> &[T] {
- // JUSTIFICATION
- // Benefit
- // Many of the buffers represent specific types, and consumers of `Buffer` often need to re-interpret them.
- // Soundness
- // * The pointer is non-null by construction
- // * alignment asserted below.
- let (prefix, offsets, suffix) = self.as_slice().align_to::<T>();
- assert!(prefix.is_empty() && suffix.is_empty());
- offsets
- }
-
- /// Returns a slice of this buffer starting at a certain bit offset.
- /// If the offset is byte-aligned the returned buffer is a shallow clone,
- /// otherwise a new buffer is allocated and filled with a copy of the bits in the range.
- pub fn bit_slice(&self, offset: usize, len: usize) -> Self {
- if offset % 8 == 0 && len % 8 == 0 {
- return self.slice(offset / 8);
- }
-
- bitwise_unary_op_helper(&self, offset, len, |a| a)
- }
-
- /// Returns a `BitChunks` instance which can be used to iterate over this buffers bits
- /// in larger chunks and starting at arbitrary bit offsets.
- /// Note that both `offset` and `length` are measured in bits.
- pub fn bit_chunks(&self, offset: usize, len: usize) -> BitChunks {
- BitChunks::new(&self.as_slice(), offset, len)
- }
-
- /// Returns the number of 1-bits in this buffer.
- pub fn count_set_bits(&self) -> usize {
- let len_in_bits = self.len() * 8;
- // self.offset is already taken into consideration by the bit_chunks implementation
- self.count_set_bits_offset(0, len_in_bits)
- }
-
- /// Returns the number of 1-bits in this buffer, starting from `offset` with `length` bits
- /// inspected. Note that both `offset` and `length` are measured in bits.
- pub fn count_set_bits_offset(&self, offset: usize, len: usize) -> usize {
- let chunks = self.bit_chunks(offset, len);
- let mut count = chunks.iter().map(|c| c.count_ones() as usize).sum();
- count += chunks.remainder_bits().count_ones() as usize;
-
- count
- }
-}
-
-/// Creating a `Buffer` instance by copying the memory from a `AsRef<[u8]>` into a newly
-/// allocated memory region.
-impl<T: AsRef<[u8]>> From<T> for Buffer {
- fn from(p: T) -> Self {
- // allocate aligned memory buffer
- let slice = p.as_ref();
- let len = slice.len();
- let mut buffer = MutableBuffer::new(len);
- buffer.extend_from_slice(slice);
- buffer.into()
- }
-}
-
-/// Creating a `Buffer` instance by storing the boolean values into the buffer
-impl std::iter::FromIterator<bool> for Buffer {
- fn from_iter<I>(iter: I) -> Self
- where
- I: IntoIterator<Item = bool>,
- {
- MutableBuffer::from_iter(iter).into()
- }
-}
-
-impl std::ops::Deref for Buffer {
- type Target = [u8];
-
- fn deref(&self) -> &[u8] {
- unsafe { std::slice::from_raw_parts(self.as_ptr(), self.len()) }
- }
-}
-
-unsafe impl Sync for Buffer {}
-unsafe impl Send for Buffer {}
-
-impl From<MutableBuffer> for Buffer {
- #[inline]
- fn from(buffer: MutableBuffer) -> Self {
- buffer.into_buffer()
- }
-}
-
-impl Buffer {
- /// Creates a [`Buffer`] from an [`Iterator`] with a trusted (upper) length.
- /// Prefer this to `collect` whenever possible, as it is ~60% faster.
- /// # Example
- /// ```
- /// # use arrow::buffer::Buffer;
- /// let v = vec![1u32];
- /// let iter = v.iter().map(|x| x * 2);
- /// let buffer = unsafe { Buffer::from_trusted_len_iter(iter) };
- /// assert_eq!(buffer.len(), 4) // u32 has 4 bytes
- /// ```
- /// # Safety
- /// This method assumes that the iterator's size is correct and is undefined behavior
- /// to use it on an iterator that reports an incorrect length.
- // This implementation is required for two reasons:
- // 1. there is no trait `TrustedLen` in stable rust and therefore
- // we can't specialize `extend` for `TrustedLen` like `Vec` does.
- // 2. `from_trusted_len_iter` is faster.
- #[inline]
- pub unsafe fn from_trusted_len_iter<T: ArrowNativeType, I: Iterator<Item = T>>(
- iterator: I,
- ) -> Self {
- MutableBuffer::from_trusted_len_iter(iterator).into()
- }
-
- /// Creates a [`Buffer`] from an [`Iterator`] with a trusted (upper) length or errors
- /// if any of the items of the iterator is an error.
- /// Prefer this to `collect` whenever possible, as it is ~60% faster.
- /// # Safety
- /// This method assumes that the iterator's size is correct and is undefined behavior
- /// to use it on an iterator that reports an incorrect length.
- #[inline]
- pub unsafe fn try_from_trusted_len_iter<
- E,
- T: ArrowNativeType,
- I: Iterator<Item = std::result::Result<T, E>>,
- >(
- iterator: I,
- ) -> std::result::Result<Self, E> {
- Ok(MutableBuffer::try_from_trusted_len_iter(iterator)?.into())
- }
-}
-
-impl<T: ArrowNativeType> FromIterator<T> for Buffer {
- fn from_iter<I: IntoIterator<Item = T>>(iter: I) -> Self {
- let mut iterator = iter.into_iter();
- let size = std::mem::size_of::<T>();
-
- // first iteration, which will likely reserve sufficient space for the buffer.
- let mut buffer = match iterator.next() {
- None => MutableBuffer::new(0),
- Some(element) => {
- let (lower, _) = iterator.size_hint();
- let mut buffer = MutableBuffer::new(lower.saturating_add(1) * size);
- unsafe {
- std::ptr::write(buffer.as_mut_ptr() as *mut T, element);
- buffer.set_len(size);
- }
- buffer
- }
- };
-
- buffer.extend_from_iter(iterator);
- buffer.into()
- }
-}
-
-#[cfg(test)]
-mod tests {
- use std::thread;
-
- use super::*;
-
- #[test]
- fn test_buffer_data_equality() {
- let buf1 = Buffer::from(&[0, 1, 2, 3, 4]);
- let buf2 = Buffer::from(&[0, 1, 2, 3, 4]);
- assert_eq!(buf1, buf2);
-
- // slice with same offset should still preserve equality
- let buf3 = buf1.slice(2);
- assert_ne!(buf1, buf3);
- let buf4 = buf2.slice(2);
- assert_eq!(buf3, buf4);
-
- // Different capacities should still preserve equality
- let mut buf2 = MutableBuffer::new(65);
- buf2.extend_from_slice(&[0u8, 1, 2, 3, 4]);
-
- let buf2 = buf2.into();
- assert_eq!(buf1, buf2);
-
- // unequal because of different elements
- let buf2 = Buffer::from(&[0, 0, 2, 3, 4]);
- assert_ne!(buf1, buf2);
-
- // unequal because of different length
- let buf2 = Buffer::from(&[0, 1, 2, 3]);
- assert_ne!(buf1, buf2);
- }
-
- #[test]
- fn test_from_raw_parts() {
- let buf = Buffer::from(&[0, 1, 2, 3, 4]);
- assert_eq!(5, buf.len());
- assert!(!buf.as_ptr().is_null());
- assert_eq!([0, 1, 2, 3, 4], buf.as_slice());
- }
-
- #[test]
- fn test_from_vec() {
- let buf = Buffer::from(&[0, 1, 2, 3, 4]);
- assert_eq!(5, buf.len());
- assert!(!buf.as_ptr().is_null());
- assert_eq!([0, 1, 2, 3, 4], buf.as_slice());
- }
-
- #[test]
- fn test_copy() {
- let buf = Buffer::from(&[0, 1, 2, 3, 4]);
- let buf2 = buf;
- assert_eq!(5, buf2.len());
- assert_eq!(64, buf2.capacity());
- assert!(!buf2.as_ptr().is_null());
- assert_eq!([0, 1, 2, 3, 4], buf2.as_slice());
- }
-
- #[test]
- fn test_slice() {
- let buf = Buffer::from(&[2, 4, 6, 8, 10]);
- let buf2 = buf.slice(2);
-
- assert_eq!([6, 8, 10], buf2.as_slice());
- assert_eq!(3, buf2.len());
- assert_eq!(unsafe { buf.as_ptr().offset(2) }, buf2.as_ptr());
-
- let buf3 = buf2.slice(1);
- assert_eq!([8, 10], buf3.as_slice());
- assert_eq!(2, buf3.len());
- assert_eq!(unsafe { buf.as_ptr().offset(3) }, buf3.as_ptr());
-
- let buf4 = buf.slice(5);
- let empty_slice: [u8; 0] = [];
- assert_eq!(empty_slice, buf4.as_slice());
- assert_eq!(0, buf4.len());
- assert!(buf4.is_empty());
- assert_eq!(buf2.slice(2).as_slice(), &[10]);
- }
-
- #[test]
- #[should_panic(
- expected = "the offset of the new Buffer cannot exceed the existing length"
- )]
- fn test_slice_offset_out_of_bound() {
- let buf = Buffer::from(&[2, 4, 6, 8, 10]);
- buf.slice(6);
- }
-
- #[test]
- fn test_access_concurrently() {
- let buffer = Buffer::from(vec![1, 2, 3, 4, 5]);
- let buffer2 = buffer.clone();
- assert_eq!([1, 2, 3, 4, 5], buffer.as_slice());
-
- let buffer_copy = thread::spawn(move || {
- // access buffer in another thread.
- buffer
- })
- .join();
-
- assert!(buffer_copy.is_ok());
- assert_eq!(buffer2, buffer_copy.ok().unwrap());
- }
-
- macro_rules! check_as_typed_data {
- ($input: expr, $native_t: ty) => {{
- let buffer = Buffer::from_slice_ref($input);
- let slice: &[$native_t] = unsafe { buffer.typed_data::<$native_t>() };
- assert_eq!($input, slice);
- }};
- }
-
- #[test]
- #[allow(clippy::float_cmp)]
- fn test_as_typed_data() {
- check_as_typed_data!(&[1i8, 3i8, 6i8], i8);
- check_as_typed_data!(&[1u8, 3u8, 6u8], u8);
- check_as_typed_data!(&[1i16, 3i16, 6i16], i16);
- check_as_typed_data!(&[1i32, 3i32, 6i32], i32);
- check_as_typed_data!(&[1i64, 3i64, 6i64], i64);
- check_as_typed_data!(&[1u16, 3u16, 6u16], u16);
- check_as_typed_data!(&[1u32, 3u32, 6u32], u32);
- check_as_typed_data!(&[1u64, 3u64, 6u64], u64);
- check_as_typed_data!(&[1f32, 3f32, 6f32], f32);
- check_as_typed_data!(&[1f64, 3f64, 6f64], f64);
- }
-
- #[test]
- fn test_count_bits() {
- assert_eq!(0, Buffer::from(&[0b00000000]).count_set_bits());
- assert_eq!(8, Buffer::from(&[0b11111111]).count_set_bits());
- assert_eq!(3, Buffer::from(&[0b00001101]).count_set_bits());
- assert_eq!(6, Buffer::from(&[0b01001001, 0b01010010]).count_set_bits());
- assert_eq!(16, Buffer::from(&[0b11111111, 0b11111111]).count_set_bits());
- }
-
- #[test]
- fn test_count_bits_slice() {
- assert_eq!(
- 0,
- Buffer::from(&[0b11111111, 0b00000000])
- .slice(1)
- .count_set_bits()
- );
- assert_eq!(
- 8,
- Buffer::from(&[0b11111111, 0b11111111])
- .slice(1)
- .count_set_bits()
- );
- assert_eq!(
- 3,
- Buffer::from(&[0b11111111, 0b11111111, 0b00001101])
- .slice(2)
- .count_set_bits()
- );
- assert_eq!(
- 6,
- Buffer::from(&[0b11111111, 0b01001001, 0b01010010])
- .slice(1)
- .count_set_bits()
- );
- assert_eq!(
- 16,
- Buffer::from(&[0b11111111, 0b11111111, 0b11111111, 0b11111111])
- .slice(2)
- .count_set_bits()
- );
- }
-
- #[test]
- fn test_count_bits_offset_slice() {
- assert_eq!(8, Buffer::from(&[0b11111111]).count_set_bits_offset(0, 8));
- assert_eq!(3, Buffer::from(&[0b11111111]).count_set_bits_offset(0, 3));
- assert_eq!(5, Buffer::from(&[0b11111111]).count_set_bits_offset(3, 5));
- assert_eq!(1, Buffer::from(&[0b11111111]).count_set_bits_offset(3, 1));
- assert_eq!(0, Buffer::from(&[0b11111111]).count_set_bits_offset(8, 0));
- assert_eq!(2, Buffer::from(&[0b01010101]).count_set_bits_offset(0, 3));
- assert_eq!(
- 16,
- Buffer::from(&[0b11111111, 0b11111111]).count_set_bits_offset(0, 16)
- );
- assert_eq!(
- 10,
- Buffer::from(&[0b11111111, 0b11111111]).count_set_bits_offset(0, 10)
- );
- assert_eq!(
- 10,
- Buffer::from(&[0b11111111, 0b11111111]).count_set_bits_offset(3, 10)
- );
- assert_eq!(
- 8,
- Buffer::from(&[0b11111111, 0b11111111]).count_set_bits_offset(8, 8)
- );
- assert_eq!(
- 5,
- Buffer::from(&[0b11111111, 0b11111111]).count_set_bits_offset(11, 5)
- );
- assert_eq!(
- 0,
- Buffer::from(&[0b11111111, 0b11111111]).count_set_bits_offset(16, 0)
- );
- assert_eq!(
- 2,
- Buffer::from(&[0b01101101, 0b10101010]).count_set_bits_offset(7, 5)
- );
- assert_eq!(
- 4,
- Buffer::from(&[0b01101101, 0b10101010]).count_set_bits_offset(7, 9)
- );
- }
-}
diff --git a/arrow/src/buffer/mod.rs b/arrow/src/buffer/mod.rs
deleted file mode 100644
index cc5c63b..0000000
--- a/arrow/src/buffer/mod.rs
+++ /dev/null
@@ -1,69 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-//! This module contains two main structs: [Buffer] and [MutableBuffer]. A buffer represents
-//! a contiguous memory region that can be shared via `offsets`.
-
-mod immutable;
-pub use immutable::*;
-mod mutable;
-pub use mutable::*;
-mod ops;
-pub(super) use ops::*;
-
-use crate::error::{ArrowError, Result};
-use std::ops::{BitAnd, BitOr, Not};
-
-impl<'a, 'b> BitAnd<&'b Buffer> for &'a Buffer {
- type Output = Result<Buffer>;
-
- fn bitand(self, rhs: &'b Buffer) -> Result<Buffer> {
- if self.len() != rhs.len() {
- return Err(ArrowError::ComputeError(
- "Buffers must be the same size to apply Bitwise AND.".to_string(),
- ));
- }
-
- let len_in_bits = self.len() * 8;
- Ok(buffer_bin_and(&self, 0, &rhs, 0, len_in_bits))
- }
-}
-
-impl<'a, 'b> BitOr<&'b Buffer> for &'a Buffer {
- type Output = Result<Buffer>;
-
- fn bitor(self, rhs: &'b Buffer) -> Result<Buffer> {
- if self.len() != rhs.len() {
- return Err(ArrowError::ComputeError(
- "Buffers must be the same size to apply Bitwise OR.".to_string(),
- ));
- }
-
- let len_in_bits = self.len() * 8;
-
- Ok(buffer_bin_or(&self, 0, &rhs, 0, len_in_bits))
- }
-}
-
-impl Not for &Buffer {
- type Output = Buffer;
-
- fn not(self) -> Buffer {
- let len_in_bits = self.len() * 8;
- buffer_unary_not(&self, 0, len_in_bits)
- }
-}
diff --git a/arrow/src/buffer/mutable.rs b/arrow/src/buffer/mutable.rs
deleted file mode 100644
index 7d336e0..0000000
--- a/arrow/src/buffer/mutable.rs
+++ /dev/null
@@ -1,791 +0,0 @@
-use std::ptr::NonNull;
-
-use crate::{
- alloc,
- bytes::{Bytes, Deallocation},
- datatypes::{ArrowNativeType, ToByteSlice},
- util::bit_util,
-};
-
-use super::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.
-
-/// A [`MutableBuffer`] is Arrow's interface to build a [`Buffer`] out of items or slices of items.
-/// [`Buffer`]s created from [`MutableBuffer`] (via `into`) are guaranteed to have its pointer aligned
-/// along cache lines and in multiple of 64 bytes.
-/// Use [MutableBuffer::push] to insert an item, [MutableBuffer::extend_from_slice]
-/// to insert many items, and `into` to convert it to [`Buffer`].
-/// # Example
-/// ```
-/// # use arrow::buffer::{Buffer, MutableBuffer};
-/// let mut buffer = MutableBuffer::new(0);
-/// buffer.push(256u32);
-/// buffer.extend_from_slice(&[1u32]);
-/// let buffer: Buffer = buffer.into();
-/// assert_eq!(buffer.as_slice(), &[0u8, 1, 0, 0, 1, 0, 0, 0])
-/// ```
-#[derive(Debug)]
-pub struct MutableBuffer {
- // dangling iff capacity = 0
- data: NonNull<u8>,
- // invariant: len <= capacity
- len: usize,
- capacity: usize,
-}
-
-impl MutableBuffer {
- /// Allocate a new [MutableBuffer] with initial capacity to be at least `capacity`.
- #[inline]
- pub fn new(capacity: usize) -> Self {
- Self::with_capacity(capacity)
- }
-
- /// Allocate a new [MutableBuffer] with initial capacity to be at least `capacity`.
- #[inline]
- pub fn with_capacity(capacity: usize) -> Self {
- let capacity = bit_util::round_upto_multiple_of_64(capacity);
- let ptr = alloc::allocate_aligned(capacity);
- Self {
- data: ptr,
- len: 0,
- capacity,
- }
- }
-
- /// Allocates a new [MutableBuffer] with `len` and capacity to be at least `len` where
- /// all bytes are guaranteed to be `0u8`.
- /// # Example
- /// ```
- /// # use arrow::buffer::{Buffer, MutableBuffer};
- /// let mut buffer = MutableBuffer::from_len_zeroed(127);
- /// assert_eq!(buffer.len(), 127);
- /// assert!(buffer.capacity() >= 127);
- /// let data = buffer.as_slice_mut();
- /// assert_eq!(data[126], 0u8);
- /// ```
- pub fn from_len_zeroed(len: usize) -> Self {
- let new_capacity = bit_util::round_upto_multiple_of_64(len);
- let ptr = alloc::allocate_aligned_zeroed(new_capacity);
- Self {
- data: ptr,
- len,
- capacity: new_capacity,
- }
- }
-
- /// creates a new [MutableBuffer] with capacity and length capable of holding `len` bits.
- /// This is useful to create a buffer for packed bitmaps.
- pub fn new_null(len: usize) -> Self {
- let num_bytes = bit_util::ceil(len, 8);
- MutableBuffer::from_len_zeroed(num_bytes)
- }
-
- /// Set the bits in the range of `[0, end)` to 0 (if `val` is false), or 1 (if `val`
- /// is true). Also extend the length of this buffer to be `end`.
- ///
- /// This is useful when one wants to clear (or set) the bits and then manipulate
- /// the buffer directly (e.g., modifying the buffer by holding a mutable reference
- /// from `data_mut()`).
- pub fn with_bitset(mut self, end: usize, val: bool) -> Self {
- assert!(end <= self.capacity);
- let v = if val { 255 } else { 0 };
- unsafe {
- std::ptr::write_bytes(self.data.as_ptr(), v, end);
- self.len = end;
- }
- self
- }
-
- /// Ensure that `count` bytes from `start` contain zero bits
- ///
- /// This is used to initialize the bits in a buffer, however, it has no impact on the
- /// `len` of the buffer and so can be used to initialize the memory region from
- /// `len` to `capacity`.
- pub fn set_null_bits(&mut self, start: usize, count: usize) {
- assert!(start + count <= self.capacity);
- unsafe {
- std::ptr::write_bytes(self.data.as_ptr().add(start), 0, count);
- }
- }
-
- /// Ensures that this buffer has at least `self.len + additional` bytes. This re-allocates iff
- /// `self.len + additional > capacity`.
- /// # Example
- /// ```
- /// # use arrow::buffer::{Buffer, MutableBuffer};
- /// let mut buffer = MutableBuffer::new(0);
- /// buffer.reserve(253); // allocates for the first time
- /// (0..253u8).for_each(|i| buffer.push(i)); // no reallocation
- /// let buffer: Buffer = buffer.into();
- /// assert_eq!(buffer.len(), 253);
- /// ```
- // For performance reasons, this must be inlined so that the `if` is executed inside the caller, and not as an extra call that just
- // exits.
- #[inline(always)]
- pub fn reserve(&mut self, additional: usize) {
- let required_cap = self.len + additional;
- if required_cap > self.capacity {
- // JUSTIFICATION
- // Benefit
- // necessity
- // Soundness
- // `self.data` is valid for `self.capacity`.
- let (ptr, new_capacity) =
- unsafe { reallocate(self.data, self.capacity, required_cap) };
- self.data = ptr;
- self.capacity = new_capacity;
- }
- }
-
- /// Resizes the buffer, either truncating its contents (with no change in capacity), or
- /// growing it (potentially reallocating it) and writing `value` in the newly available bytes.
- /// # Example
- /// ```
- /// # use arrow::buffer::{Buffer, MutableBuffer};
- /// let mut buffer = MutableBuffer::new(0);
- /// buffer.resize(253, 2); // allocates for the first time
- /// assert_eq!(buffer.as_slice()[252], 2u8);
- /// ```
- // For performance reasons, this must be inlined so that the `if` is executed inside the caller, and not as an extra call that just
- // exits.
- #[inline(always)]
- pub fn resize(&mut self, new_len: usize, value: u8) {
- if new_len > self.len {
- let diff = new_len - self.len;
- self.reserve(diff);
- // write the value
- unsafe { self.data.as_ptr().add(self.len).write_bytes(value, diff) };
- }
- // this truncates the buffer when new_len < self.len
- self.len = new_len;
- }
-
- /// Shrinks the capacity of the buffer as much as possible.
- /// The new capacity will aligned to the nearest 64 bit alignment.
- ///
- /// # Example
- /// ```
- /// # use arrow::buffer::{Buffer, MutableBuffer};
- /// // 2 cache lines
- /// let mut buffer = MutableBuffer::new(128);
- /// assert_eq!(buffer.capacity(), 128);
- /// buffer.push(1);
- /// buffer.push(2);
- ///
- /// buffer.shrink_to_fit();
- /// assert!(buffer.capacity() >= 64 && buffer.capacity() < 128);
- /// ```
- pub fn shrink_to_fit(&mut self) {
- let new_capacity = bit_util::round_upto_multiple_of_64(self.len);
- if new_capacity < self.capacity {
- // JUSTIFICATION
- // Benefit
- // necessity
- // Soundness
- // `self.data` is valid for `self.capacity`.
- let ptr =
- unsafe { alloc::reallocate(self.data, self.capacity, new_capacity) };
-
- self.data = ptr;
- self.capacity = new_capacity;
- }
- }
-
- /// Returns whether this buffer is empty or not.
- #[inline]
- pub const fn is_empty(&self) -> bool {
- self.len == 0
- }
-
- /// Returns the length (the number of bytes written) in this buffer.
- /// The invariant `buffer.len() <= buffer.capacity()` is always upheld.
- #[inline]
- pub const fn len(&self) -> usize {
- self.len
- }
-
- /// Returns the total capacity in this buffer.
- /// The invariant `buffer.len() <= buffer.capacity()` is always upheld.
- #[inline]
- pub const fn capacity(&self) -> usize {
- self.capacity
- }
-
- /// Clear all existing data from this buffer.
- pub fn clear(&mut self) {
- self.len = 0
- }
-
- /// Returns the data stored in this buffer as a slice.
- pub fn as_slice(&self) -> &[u8] {
- self
- }
-
- /// Returns the data stored in this buffer as a mutable slice.
- pub fn as_slice_mut(&mut self) -> &mut [u8] {
- self
- }
-
- /// Returns a raw pointer to this buffer's internal memory
- /// This pointer is guaranteed to be aligned along cache-lines.
- #[inline]
- pub const fn as_ptr(&self) -> *const u8 {
- self.data.as_ptr()
- }
-
- /// Returns a mutable raw pointer to this buffer's internal memory
- /// This pointer is guaranteed to be aligned along cache-lines.
- #[inline]
- pub fn as_mut_ptr(&mut self) -> *mut u8 {
- self.data.as_ptr()
- }
-
- #[deprecated(
- since = "2.0.0",
- note = "This method is deprecated in favour of `into` from the trait `Into`."
- )]
- /// Freezes this buffer and return an immutable version of it.
- pub fn freeze(self) -> Buffer {
- self.into_buffer()
- }
-
- #[inline]
- pub(super) fn into_buffer(self) -> Buffer {
- let bytes = unsafe {
- Bytes::new(self.data, self.len, Deallocation::Native(self.capacity))
- };
- std::mem::forget(self);
- Buffer::from_bytes(bytes)
- }
-
- /// View this buffer asa slice of a specific type.
- /// # Safety
- /// This function must only be used when this buffer was extended with items of type `T`.
- /// Failure to do so results in undefined behavior.
- pub fn typed_data_mut<T: ArrowNativeType>(&mut self) -> &mut [T] {
- unsafe {
- let (prefix, offsets, suffix) = self.as_slice_mut().align_to_mut::<T>();
- assert!(prefix.is_empty() && suffix.is_empty());
- offsets
- }
- }
-
- /// Extends this buffer from a slice of items that can be represented in bytes, increasing its capacity if needed.
- /// # Example
- /// ```
- /// # use arrow::buffer::MutableBuffer;
- /// let mut buffer = MutableBuffer::new(0);
- /// buffer.extend_from_slice(&[2u32, 0]);
- /// assert_eq!(buffer.len(), 8) // u32 has 4 bytes
- /// ```
- #[inline]
- pub fn extend_from_slice<T: ToByteSlice>(&mut self, items: &[T]) {
- let len = items.len();
- let additional = len * std::mem::size_of::<T>();
- self.reserve(additional);
- unsafe {
- let dst = self.data.as_ptr().add(self.len);
- let src = items.as_ptr() as *const u8;
- std::ptr::copy_nonoverlapping(src, dst, additional)
- }
- self.len += additional;
- }
-
- /// Extends the buffer with a new item, increasing its capacity if needed.
- /// # Example
- /// ```
- /// # use arrow::buffer::MutableBuffer;
- /// let mut buffer = MutableBuffer::new(0);
- /// buffer.push(256u32);
- /// assert_eq!(buffer.len(), 4) // u32 has 4 bytes
- /// ```
- #[inline]
- pub fn push<T: ToByteSlice>(&mut self, item: T) {
- let additional = std::mem::size_of::<T>();
- self.reserve(additional);
- unsafe {
- let dst = self.data.as_ptr().add(self.len) as *mut T;
- std::ptr::write(dst, item);
- }
- self.len += additional;
- }
-
- /// Extends the buffer with a new item, without checking for sufficient capacity
- /// # Safety
- /// Caller must ensure that the capacity()-len()>=size_of<T>()
- #[inline]
- pub unsafe fn push_unchecked<T: ToByteSlice>(&mut self, item: T) {
- let additional = std::mem::size_of::<T>();
- let dst = self.data.as_ptr().add(self.len) as *mut T;
- std::ptr::write(dst, item);
- self.len += additional;
- }
-
- /// Extends the buffer by `additional` bytes equal to `0u8`, incrementing its capacity if needed.
- #[inline]
- pub fn extend_zeros(&mut self, additional: usize) {
- self.resize(self.len + additional, 0);
- }
-
- /// # Safety
- /// The caller must ensure that the buffer was properly initialized up to `len`.
- #[inline]
- pub(crate) unsafe fn set_len(&mut self, len: usize) {
- assert!(len <= self.capacity());
- self.len = len;
- }
-}
-
-/// # Safety
-/// `ptr` must be allocated for `old_capacity`.
-#[inline]
-unsafe fn reallocate(
- ptr: NonNull<u8>,
- old_capacity: usize,
- new_capacity: usize,
-) -> (NonNull<u8>, usize) {
- let new_capacity = bit_util::round_upto_multiple_of_64(new_capacity);
- let new_capacity = std::cmp::max(new_capacity, old_capacity * 2);
- let ptr = alloc::reallocate(ptr, old_capacity, new_capacity);
- (ptr, new_capacity)
-}
-
-impl<A: ArrowNativeType> Extend<A> for MutableBuffer {
- #[inline]
- fn extend<T: IntoIterator<Item = A>>(&mut self, iter: T) {
- let iterator = iter.into_iter();
- self.extend_from_iter(iterator)
- }
-}
-
-impl MutableBuffer {
- #[inline]
- pub(super) fn extend_from_iter<T: ArrowNativeType, I: Iterator<Item = T>>(
- &mut self,
- mut iterator: I,
- ) {
- let size = std::mem::size_of::<T>();
- let (lower, _) = iterator.size_hint();
- let additional = lower * size;
- self.reserve(additional);
-
- // this is necessary because of https://github.com/rust-lang/rust/issues/32155
- let mut len = SetLenOnDrop::new(&mut self.len);
- let mut dst = unsafe { self.data.as_ptr().add(len.local_len) as *mut T };
- let capacity = self.capacity;
-
- while len.local_len + size <= capacity {
- if let Some(item) = iterator.next() {
- unsafe {
- std::ptr::write(dst, item);
- dst = dst.add(1);
- }
- len.local_len += size;
- } else {
- break;
- }
- }
- drop(len);
-
- iterator.for_each(|item| self.push(item));
- }
-
- /// Creates a [`MutableBuffer`] from an [`Iterator`] with a trusted (upper) length.
- /// Prefer this to `collect` whenever possible, as it is faster ~60% faster.
- /// # Example
- /// ```
- /// # use arrow::buffer::MutableBuffer;
- /// let v = vec![1u32];
- /// let iter = v.iter().map(|x| x * 2);
- /// let buffer = unsafe { MutableBuffer::from_trusted_len_iter(iter) };
- /// assert_eq!(buffer.len(), 4) // u32 has 4 bytes
- /// ```
- /// # Safety
- /// This method assumes that the iterator's size is correct and is undefined behavior
- /// to use it on an iterator that reports an incorrect length.
- // This implementation is required for two reasons:
- // 1. there is no trait `TrustedLen` in stable rust and therefore
- // we can't specialize `extend` for `TrustedLen` like `Vec` does.
- // 2. `from_trusted_len_iter` is faster.
- #[inline]
- pub unsafe fn from_trusted_len_iter<T: ArrowNativeType, I: Iterator<Item = T>>(
- iterator: I,
- ) -> Self {
- let (_, upper) = iterator.size_hint();
- let upper = upper.expect("from_trusted_len_iter requires an upper limit");
- let len = upper * std::mem::size_of::<T>();
-
- let mut buffer = MutableBuffer::new(len);
-
- let mut dst = buffer.data.as_ptr() as *mut T;
- for item in iterator {
- // note how there is no reserve here (compared with `extend_from_iter`)
- std::ptr::write(dst, item);
- dst = dst.add(1);
- }
- assert_eq!(
- dst.offset_from(buffer.data.as_ptr() as *mut T) as usize,
- upper,
- "Trusted iterator length was not accurately reported"
- );
- buffer.len = len;
- buffer
- }
-
- /// Creates a [`MutableBuffer`] from a boolean [`Iterator`] with a trusted (upper) length.
- /// # use arrow::buffer::MutableBuffer;
- /// # Example
- /// ```
- /// # use arrow::buffer::MutableBuffer;
- /// let v = vec![false, true, false];
- /// let iter = v.iter().map(|x| *x || true);
- /// let buffer = unsafe { MutableBuffer::from_trusted_len_iter_bool(iter) };
- /// assert_eq!(buffer.len(), 1) // 3 booleans have 1 byte
- /// ```
- /// # Safety
- /// This method assumes that the iterator's size is correct and is undefined behavior
- /// to use it on an iterator that reports an incorrect length.
- // This implementation is required for two reasons:
- // 1. there is no trait `TrustedLen` in stable rust and therefore
- // we can't specialize `extend` for `TrustedLen` like `Vec` does.
- // 2. `from_trusted_len_iter_bool` is faster.
- #[inline]
- pub unsafe fn from_trusted_len_iter_bool<I: Iterator<Item = bool>>(
- mut iterator: I,
- ) -> Self {
- let (_, upper) = iterator.size_hint();
- let upper = upper.expect("from_trusted_len_iter requires an upper limit");
-
- let mut result = {
- let byte_capacity: usize = upper.saturating_add(7) / 8;
- MutableBuffer::new(byte_capacity)
- };
-
- 'a: loop {
- let mut byte_accum: u8 = 0;
- let mut mask: u8 = 1;
-
- //collect (up to) 8 bits into a byte
- while mask != 0 {
- if let Some(value) = iterator.next() {
- byte_accum |= match value {
- true => mask,
- false => 0,
- };
- mask <<= 1;
- } else {
- if mask != 1 {
- // Add last byte
- result.push_unchecked(byte_accum);
- }
- break 'a;
- }
- }
-
- // Soundness: from_trusted_len
- result.push_unchecked(byte_accum);
- }
- result
- }
-
- /// Creates a [`MutableBuffer`] from an [`Iterator`] with a trusted (upper) length or errors
- /// if any of the items of the iterator is an error.
- /// Prefer this to `collect` whenever possible, as it is faster ~60% faster.
- /// # Safety
- /// This method assumes that the iterator's size is correct and is undefined behavior
- /// to use it on an iterator that reports an incorrect length.
- #[inline]
- pub unsafe fn try_from_trusted_len_iter<
- E,
- T: ArrowNativeType,
- I: Iterator<Item = std::result::Result<T, E>>,
- >(
- iterator: I,
- ) -> std::result::Result<Self, E> {
- let (_, upper) = iterator.size_hint();
- let upper = upper.expect("try_from_trusted_len_iter requires an upper limit");
- let len = upper * std::mem::size_of::<T>();
-
- let mut buffer = MutableBuffer::new(len);
-
- let mut dst = buffer.data.as_ptr() as *mut T;
- for item in iterator {
- // note how there is no reserve here (compared with `extend_from_iter`)
- std::ptr::write(dst, item?);
- dst = dst.add(1);
- }
- assert_eq!(
- dst.offset_from(buffer.data.as_ptr() as *mut T) as usize,
- upper,
- "Trusted iterator length was not accurately reported"
- );
- buffer.len = len;
- Ok(buffer)
- }
-}
-
-impl std::ops::Deref for MutableBuffer {
- type Target = [u8];
-
- fn deref(&self) -> &[u8] {
- unsafe { std::slice::from_raw_parts(self.as_ptr(), self.len) }
- }
-}
-
-impl std::ops::DerefMut for MutableBuffer {
- fn deref_mut(&mut self) -> &mut [u8] {
- unsafe { std::slice::from_raw_parts_mut(self.as_mut_ptr(), self.len) }
- }
-}
-
-impl Drop for MutableBuffer {
- fn drop(&mut self) {
- unsafe { alloc::free_aligned(self.data, self.capacity) };
- }
-}
-
-impl PartialEq for MutableBuffer {
- fn eq(&self, other: &MutableBuffer) -> bool {
- if self.len != other.len {
- return false;
- }
- if self.capacity != other.capacity {
- return false;
- }
- self.as_slice() == other.as_slice()
- }
-}
-
-unsafe impl Sync for MutableBuffer {}
-unsafe impl Send for MutableBuffer {}
-
-struct SetLenOnDrop<'a> {
- len: &'a mut usize,
- local_len: usize,
-}
-
-impl<'a> SetLenOnDrop<'a> {
- #[inline]
- fn new(len: &'a mut usize) -> Self {
- SetLenOnDrop {
- local_len: *len,
- len,
- }
- }
-}
-
-impl Drop for SetLenOnDrop<'_> {
- #[inline]
- fn drop(&mut self) {
- *self.len = self.local_len;
- }
-}
-
-/// Creating a `MutableBuffer` instance by setting bits according to the boolean values
-impl std::iter::FromIterator<bool> for MutableBuffer {
- fn from_iter<I>(iter: I) -> Self
- where
- I: IntoIterator<Item = bool>,
- {
- let mut iterator = iter.into_iter();
- let mut result = {
- let byte_capacity: usize = iterator.size_hint().0.saturating_add(7) / 8;
- MutableBuffer::new(byte_capacity)
- };
-
- loop {
- let mut exhausted = false;
- let mut byte_accum: u8 = 0;
- let mut mask: u8 = 1;
-
- //collect (up to) 8 bits into a byte
- while mask != 0 {
- if let Some(value) = iterator.next() {
- byte_accum |= match value {
- true => mask,
- false => 0,
- };
- mask <<= 1;
- } else {
- exhausted = true;
- break;
- }
- }
-
- // break if the iterator was exhausted before it provided a bool for this byte
- if exhausted && mask == 1 {
- break;
- }
-
- //ensure we have capacity to write the byte
- if result.len() == result.capacity() {
- //no capacity for new byte, allocate 1 byte more (plus however many more the iterator advertises)
- let additional_byte_capacity = 1usize.saturating_add(
- iterator.size_hint().0.saturating_add(7) / 8, //convert bit count to byte count, rounding up
- );
- result.reserve(additional_byte_capacity)
- }
-
- // Soundness: capacity was allocated above
- unsafe { result.push_unchecked(byte_accum) };
- if exhausted {
- break;
- }
- }
- result
- }
-}
-
-#[cfg(test)]
-mod tests {
- use super::*;
-
- #[test]
- fn test_mutable_new() {
- let buf = MutableBuffer::new(63);
- assert_eq!(64, buf.capacity());
- assert_eq!(0, buf.len());
- assert!(buf.is_empty());
- }
-
- #[test]
- fn test_mutable_extend_from_slice() {
- let mut buf = MutableBuffer::new(100);
- buf.extend_from_slice(b"hello");
- assert_eq!(5, buf.len());
- assert_eq!(b"hello", buf.as_slice());
-
- buf.extend_from_slice(b" world");
- assert_eq!(11, buf.len());
- assert_eq!(b"hello world", buf.as_slice());
-
- buf.clear();
- assert_eq!(0, buf.len());
- buf.extend_from_slice(b"hello arrow");
- assert_eq!(11, buf.len());
- assert_eq!(b"hello arrow", buf.as_slice());
- }
-
- #[test]
- fn mutable_extend_from_iter() {
- let mut buf = MutableBuffer::new(0);
- buf.extend(vec![1u32, 2]);
- assert_eq!(8, buf.len());
- assert_eq!(&[1u8, 0, 0, 0, 2, 0, 0, 0], buf.as_slice());
-
- buf.extend(vec![3u32, 4]);
- assert_eq!(16, buf.len());
- assert_eq!(
- &[1u8, 0, 0, 0, 2, 0, 0, 0, 3, 0, 0, 0, 4, 0, 0, 0],
- buf.as_slice()
- );
- }
-
- #[test]
- fn test_from_trusted_len_iter() {
- let iter = vec![1u32, 2].into_iter();
- let buf = unsafe { Buffer::from_trusted_len_iter(iter) };
- assert_eq!(8, buf.len());
- assert_eq!(&[1u8, 0, 0, 0, 2, 0, 0, 0], buf.as_slice());
- }
-
- #[test]
- fn test_mutable_reserve() {
- let mut buf = MutableBuffer::new(1);
- assert_eq!(64, buf.capacity());
-
- // Reserving a smaller capacity should have no effect.
- buf.reserve(10);
- assert_eq!(64, buf.capacity());
-
- buf.reserve(80);
- assert_eq!(128, buf.capacity());
-
- buf.reserve(129);
- assert_eq!(256, buf.capacity());
- }
-
- #[test]
- fn test_mutable_resize() {
- let mut buf = MutableBuffer::new(1);
- assert_eq!(64, buf.capacity());
- assert_eq!(0, buf.len());
-
- buf.resize(20, 0);
- assert_eq!(64, buf.capacity());
- assert_eq!(20, buf.len());
-
- buf.resize(10, 0);
- assert_eq!(64, buf.capacity());
- assert_eq!(10, buf.len());
-
- buf.resize(100, 0);
- assert_eq!(128, buf.capacity());
- assert_eq!(100, buf.len());
-
- buf.resize(30, 0);
- assert_eq!(128, buf.capacity());
- assert_eq!(30, buf.len());
-
- buf.resize(0, 0);
- assert_eq!(128, buf.capacity());
- assert_eq!(0, buf.len());
- }
-
- #[test]
- fn test_mutable_into() {
- let mut buf = MutableBuffer::new(1);
- buf.extend_from_slice(b"aaaa bbbb cccc dddd");
- assert_eq!(19, buf.len());
- assert_eq!(64, buf.capacity());
- assert_eq!(b"aaaa bbbb cccc dddd", buf.as_slice());
-
- let immutable_buf: Buffer = buf.into();
- assert_eq!(19, immutable_buf.len());
- assert_eq!(64, immutable_buf.capacity());
- assert_eq!(b"aaaa bbbb cccc dddd", immutable_buf.as_slice());
- }
-
- #[test]
- fn test_mutable_equal() {
- let mut buf = MutableBuffer::new(1);
- let mut buf2 = MutableBuffer::new(1);
-
- buf.extend_from_slice(&[0xaa]);
- buf2.extend_from_slice(&[0xaa, 0xbb]);
- assert!(buf != buf2);
-
- buf.extend_from_slice(&[0xbb]);
- assert_eq!(buf, buf2);
-
- buf2.reserve(65);
- assert!(buf != buf2);
- }
-
- #[test]
- fn test_mutable_shrink_to_fit() {
- let mut buffer = MutableBuffer::new(128);
- assert_eq!(buffer.capacity(), 128);
- buffer.push(1);
- buffer.push(2);
-
- buffer.shrink_to_fit();
- assert!(buffer.capacity() >= 64 && buffer.capacity() < 128);
- }
-}
diff --git a/arrow/src/buffer/ops.rs b/arrow/src/buffer/ops.rs
deleted file mode 100644
index 9d88149..0000000
--- a/arrow/src/buffer/ops.rs
+++ /dev/null
@@ -1,429 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-#[cfg(feature = "simd")]
-use crate::util::bit_util;
-#[cfg(feature = "simd")]
-use packed_simd::u8x64;
-
-#[cfg(feature = "avx512")]
-use crate::arch::avx512::*;
-use crate::util::bit_util::ceil;
-#[cfg(any(feature = "simd", feature = "avx512"))]
-use std::borrow::BorrowMut;
-
-use super::{Buffer, MutableBuffer};
-
-/// Apply a bitwise operation `simd_op` / `scalar_op` to two inputs using simd instructions and return the result as a Buffer.
-/// The `simd_op` functions gets applied on chunks of 64 bytes (512 bits) at a time
-/// and the `scalar_op` gets applied to remaining bytes.
-/// Contrary to the non-simd version `bitwise_bin_op_helper`, the offset and length is specified in bytes
-/// and this version does not support operations starting at arbitrary bit offsets.
-#[cfg(feature = "simd")]
-pub fn bitwise_bin_op_simd_helper<F_SIMD, F_SCALAR>(
- left: &Buffer,
- left_offset: usize,
- right: &Buffer,
- right_offset: usize,
- len: usize,
- simd_op: F_SIMD,
- scalar_op: F_SCALAR,
-) -> Buffer
-where
- F_SIMD: Fn(u8x64, u8x64) -> u8x64,
- F_SCALAR: Fn(u8, u8) -> u8,
-{
- let mut result = MutableBuffer::new(len).with_bitset(len, false);
- let lanes = u8x64::lanes();
-
- let mut left_chunks = left.as_slice()[left_offset..].chunks_exact(lanes);
- let mut right_chunks = right.as_slice()[right_offset..].chunks_exact(lanes);
- let mut result_chunks = result.as_slice_mut().chunks_exact_mut(lanes);
-
- result_chunks
- .borrow_mut()
- .zip(left_chunks.borrow_mut().zip(right_chunks.borrow_mut()))
- .for_each(|(res, (left, right))| {
- unsafe { bit_util::bitwise_bin_op_simd(&left, &right, res, &simd_op) };
- });
-
- result_chunks
- .into_remainder()
- .iter_mut()
- .zip(
- left_chunks
- .remainder()
- .iter()
- .zip(right_chunks.remainder().iter()),
- )
- .for_each(|(res, (left, right))| {
- *res = scalar_op(*left, *right);
- });
-
- result.into()
-}
-
-/// Apply a bitwise operation `simd_op` / `scalar_op` to one input using simd instructions and return the result as a Buffer.
-/// The `simd_op` functions gets applied on chunks of 64 bytes (512 bits) at a time
-/// and the `scalar_op` gets applied to remaining bytes.
-/// Contrary to the non-simd version `bitwise_unary_op_helper`, the offset and length is specified in bytes
-/// and this version does not support operations starting at arbitrary bit offsets.
-#[cfg(feature = "simd")]
-pub fn bitwise_unary_op_simd_helper<F_SIMD, F_SCALAR>(
- left: &Buffer,
- left_offset: usize,
- len: usize,
- simd_op: F_SIMD,
- scalar_op: F_SCALAR,
-) -> Buffer
-where
- F_SIMD: Fn(u8x64) -> u8x64,
- F_SCALAR: Fn(u8) -> u8,
-{
- let mut result = MutableBuffer::new(len).with_bitset(len, false);
- let lanes = u8x64::lanes();
-
- let mut left_chunks = left.as_slice()[left_offset..].chunks_exact(lanes);
- let mut result_chunks = result.as_slice_mut().chunks_exact_mut(lanes);
-
- result_chunks
- .borrow_mut()
- .zip(left_chunks.borrow_mut())
- .for_each(|(res, left)| unsafe {
- let data_simd = u8x64::from_slice_unaligned_unchecked(left);
- let simd_result = simd_op(data_simd);
- simd_result.write_to_slice_unaligned_unchecked(res);
- });
-
- result_chunks
- .into_remainder()
- .iter_mut()
- .zip(left_chunks.remainder().iter())
- .for_each(|(res, left)| {
- *res = scalar_op(*left);
- });
-
- result.into()
-}
-
-/// Apply a bitwise operation `op` to two inputs and return the result as a Buffer.
-/// The inputs are treated as bitmaps, meaning that offsets and length are specified in number of bits.
-pub fn bitwise_bin_op_helper<F>(
- left: &Buffer,
- left_offset_in_bits: usize,
- right: &Buffer,
- right_offset_in_bits: usize,
- len_in_bits: usize,
- op: F,
-) -> Buffer
-where
- F: Fn(u64, u64) -> u64,
-{
- let left_chunks = left.bit_chunks(left_offset_in_bits, len_in_bits);
- let right_chunks = right.bit_chunks(right_offset_in_bits, len_in_bits);
-
- let chunks = left_chunks
- .iter()
- .zip(right_chunks.iter())
- .map(|(left, right)| op(left, right));
- // Soundness: `BitChunks` is a trusted len iterator
- let mut buffer = unsafe { MutableBuffer::from_trusted_len_iter(chunks) };
-
- let remainder_bytes = ceil(left_chunks.remainder_len(), 8);
- let rem = op(left_chunks.remainder_bits(), right_chunks.remainder_bits());
- // we are counting its starting from the least significant bit, to to_le_bytes should be correct
- let rem = &rem.to_le_bytes()[0..remainder_bytes];
- buffer.extend_from_slice(rem);
-
- buffer.into()
-}
-
-/// Apply a bitwise operation `op` to one input and return the result as a Buffer.
-/// The input is treated as a bitmap, meaning that offset and length are specified in number of bits.
-pub fn bitwise_unary_op_helper<F>(
- left: &Buffer,
- offset_in_bits: usize,
- len_in_bits: usize,
- op: F,
-) -> Buffer
-where
- F: Fn(u64) -> u64,
-{
- // reserve capacity and set length so we can get a typed view of u64 chunks
- let mut result =
- MutableBuffer::new(ceil(len_in_bits, 8)).with_bitset(len_in_bits / 64 * 8, false);
-
- let left_chunks = left.bit_chunks(offset_in_bits, len_in_bits);
- let result_chunks = result.typed_data_mut::<u64>().iter_mut();
-
- result_chunks
- .zip(left_chunks.iter())
- .for_each(|(res, left)| {
- *res = op(left);
- });
-
- let remainder_bytes = ceil(left_chunks.remainder_len(), 8);
- let rem = op(left_chunks.remainder_bits());
- // we are counting its starting from the least significant bit, to to_le_bytes should be correct
- let rem = &rem.to_le_bytes()[0..remainder_bytes];
- result.extend_from_slice(rem);
-
- result.into()
-}
-
-#[cfg(all(target_arch = "x86_64", feature = "avx512"))]
-pub fn buffer_bin_and(
- left: &Buffer,
- left_offset_in_bits: usize,
- right: &Buffer,
- right_offset_in_bits: usize,
- len_in_bits: usize,
-) -> Buffer {
- if left_offset_in_bits % 8 == 0
- && right_offset_in_bits % 8 == 0
- && len_in_bits % 8 == 0
- {
- let len = len_in_bits / 8;
- let left_offset = left_offset_in_bits / 8;
- let right_offset = right_offset_in_bits / 8;
-
- let mut result = MutableBuffer::new(len).with_bitset(len, false);
-
- let mut left_chunks =
- left.as_slice()[left_offset..].chunks_exact(AVX512_U8X64_LANES);
- let mut right_chunks =
- right.as_slice()[right_offset..].chunks_exact(AVX512_U8X64_LANES);
- let mut result_chunks =
- result.as_slice_mut().chunks_exact_mut(AVX512_U8X64_LANES);
-
- result_chunks
- .borrow_mut()
- .zip(left_chunks.borrow_mut().zip(right_chunks.borrow_mut()))
- .for_each(|(res, (left, right))| unsafe {
- avx512_bin_and(left, right, res);
- });
-
- result_chunks
- .into_remainder()
- .iter_mut()
- .zip(
- left_chunks
- .remainder()
- .iter()
- .zip(right_chunks.remainder().iter()),
- )
- .for_each(|(res, (left, right))| {
- *res = *left & *right;
- });
-
- result.into()
- } else {
- bitwise_bin_op_helper(
- &left,
- left_offset_in_bits,
- right,
- right_offset_in_bits,
- len_in_bits,
- |a, b| a & b,
- )
- }
-}
-
-#[cfg(all(feature = "simd", not(feature = "avx512")))]
-pub fn buffer_bin_and(
- left: &Buffer,
- left_offset_in_bits: usize,
- right: &Buffer,
- right_offset_in_bits: usize,
- len_in_bits: usize,
-) -> Buffer {
- if left_offset_in_bits % 8 == 0
- && right_offset_in_bits % 8 == 0
- && len_in_bits % 8 == 0
- {
- bitwise_bin_op_simd_helper(
- &left,
- left_offset_in_bits / 8,
- &right,
- right_offset_in_bits / 8,
- len_in_bits / 8,
- |a, b| a & b,
- |a, b| a & b,
- )
- } else {
- bitwise_bin_op_helper(
- &left,
- left_offset_in_bits,
- right,
- right_offset_in_bits,
- len_in_bits,
- |a, b| a & b,
- )
- }
-}
-
-// Note: do not target specific features like x86 without considering
-// other targets like wasm32, as those would fail to build
-#[cfg(all(not(any(feature = "simd", feature = "avx512"))))]
-pub fn buffer_bin_and(
- left: &Buffer,
- left_offset_in_bits: usize,
- right: &Buffer,
- right_offset_in_bits: usize,
- len_in_bits: usize,
-) -> Buffer {
- bitwise_bin_op_helper(
- &left,
- left_offset_in_bits,
- right,
- right_offset_in_bits,
- len_in_bits,
- |a, b| a & b,
- )
-}
-
-#[cfg(all(target_arch = "x86_64", feature = "avx512"))]
-pub fn buffer_bin_or(
- left: &Buffer,
- left_offset_in_bits: usize,
- right: &Buffer,
- right_offset_in_bits: usize,
- len_in_bits: usize,
-) -> Buffer {
- if left_offset_in_bits % 8 == 0
- && right_offset_in_bits % 8 == 0
- && len_in_bits % 8 == 0
- {
- let len = len_in_bits / 8;
- let left_offset = left_offset_in_bits / 8;
- let right_offset = right_offset_in_bits / 8;
-
- let mut result = MutableBuffer::new(len).with_bitset(len, false);
-
- let mut left_chunks =
- left.as_slice()[left_offset..].chunks_exact(AVX512_U8X64_LANES);
- let mut right_chunks =
- right.as_slice()[right_offset..].chunks_exact(AVX512_U8X64_LANES);
- let mut result_chunks =
- result.as_slice_mut().chunks_exact_mut(AVX512_U8X64_LANES);
-
- result_chunks
- .borrow_mut()
- .zip(left_chunks.borrow_mut().zip(right_chunks.borrow_mut()))
- .for_each(|(res, (left, right))| unsafe {
- avx512_bin_or(left, right, res);
- });
-
- result_chunks
- .into_remainder()
- .iter_mut()
- .zip(
- left_chunks
- .remainder()
- .iter()
- .zip(right_chunks.remainder().iter()),
- )
- .for_each(|(res, (left, right))| {
- *res = *left | *right;
- });
-
- result.into()
- } else {
- bitwise_bin_op_helper(
- &left,
- left_offset_in_bits,
- right,
- right_offset_in_bits,
- len_in_bits,
- |a, b| a | b,
- )
- }
-}
-
-#[cfg(all(feature = "simd", not(feature = "avx512")))]
-pub fn buffer_bin_or(
- left: &Buffer,
- left_offset_in_bits: usize,
- right: &Buffer,
- right_offset_in_bits: usize,
- len_in_bits: usize,
-) -> Buffer {
- if left_offset_in_bits % 8 == 0
- && right_offset_in_bits % 8 == 0
- && len_in_bits % 8 == 0
- {
- bitwise_bin_op_simd_helper(
- &left,
- left_offset_in_bits / 8,
- &right,
- right_offset_in_bits / 8,
- len_in_bits / 8,
- |a, b| a | b,
- |a, b| a | b,
- )
- } else {
- bitwise_bin_op_helper(
- &left,
- left_offset_in_bits,
- right,
- right_offset_in_bits,
- len_in_bits,
- |a, b| a | b,
- )
- }
-}
-
-#[cfg(all(not(any(feature = "simd", feature = "avx512"))))]
-pub fn buffer_bin_or(
- left: &Buffer,
- left_offset_in_bits: usize,
- right: &Buffer,
- right_offset_in_bits: usize,
- len_in_bits: usize,
-) -> Buffer {
- bitwise_bin_op_helper(
- &left,
- left_offset_in_bits,
- right,
- right_offset_in_bits,
- len_in_bits,
- |a, b| a | b,
- )
-}
-
-pub fn buffer_unary_not(
- left: &Buffer,
- offset_in_bits: usize,
- len_in_bits: usize,
-) -> Buffer {
- // SIMD implementation if available and byte-aligned
- #[cfg(feature = "simd")]
- if offset_in_bits % 8 == 0 && len_in_bits % 8 == 0 {
- return bitwise_unary_op_simd_helper(
- &left,
- offset_in_bits / 8,
- len_in_bits / 8,
- |a| !a,
- |a| !a,
- );
- }
- // Default implementation
- #[allow(unreachable_code)]
- {
- bitwise_unary_op_helper(&left, offset_in_bits, len_in_bits, |a| !a)
- }
-}
diff --git a/arrow/src/bytes.rs b/arrow/src/bytes.rs
deleted file mode 100644
index 38fa443..0000000
--- a/arrow/src/bytes.rs
+++ /dev/null
@@ -1,159 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-//! This module contains an implementation of a contiguous immutable memory region that knows
-//! how to de-allocate itself, [`Bytes`].
-//! Note that this is a low-level functionality of this crate.
-
-use core::slice;
-use std::ptr::NonNull;
-use std::sync::Arc;
-use std::{fmt::Debug, fmt::Formatter};
-
-use crate::{alloc, ffi};
-
-/// Mode of deallocating memory regions
-pub enum Deallocation {
- /// Native deallocation, using Rust deallocator with Arrow-specific memory aligment
- Native(usize),
- /// Foreign interface, via a callback
- Foreign(Arc<ffi::FFI_ArrowArray>),
-}
-
-impl Debug for Deallocation {
- fn fmt(&self, f: &mut Formatter) -> std::fmt::Result {
- match self {
- Deallocation::Native(capacity) => {
- write!(f, "Deallocation::Native {{ capacity: {} }}", capacity)
- }
- Deallocation::Foreign(_) => {
- write!(f, "Deallocation::Foreign {{ capacity: unknown }}")
- }
- }
- }
-}
-
-/// A continuous, fixed-size, immutable memory region that knows how to de-allocate itself.
-/// This structs' API is inspired by the `bytes::Bytes`, but it is not limited to using rust's
-/// global allocator nor u8 aligmnent.
-///
-/// In the most common case, this buffer is allocated using [`allocate_aligned`](memory::allocate_aligned)
-/// and deallocated accordingly [`free_aligned`](memory::free_aligned).
-/// When the region is allocated by an foreign allocator, [Deallocation::Foreign], this calls the
-/// foreign deallocator to deallocate the region when it is no longer needed.
-pub struct Bytes {
- /// The raw pointer to be begining of the region
- ptr: NonNull<u8>,
-
- /// The number of bytes visible to this region. This is always smaller than its capacity (when avaliable).
- len: usize,
-
- /// how to deallocate this region
- deallocation: Deallocation,
-}
-
-impl Bytes {
- /// Takes ownership of an allocated memory region,
- ///
- /// # Arguments
- ///
- /// * `ptr` - Pointer to raw parts
- /// * `len` - Length of raw parts in **bytes**
- /// * `capacity` - Total allocated memory for the pointer `ptr`, in **bytes**
- ///
- /// # Safety
- ///
- /// This function is unsafe as there is no guarantee that the given pointer is valid for `len`
- /// bytes. If the `ptr` and `capacity` come from a `Buffer`, then this is guaranteed.
- #[inline]
- pub unsafe fn new(
- ptr: std::ptr::NonNull<u8>,
- len: usize,
- deallocation: Deallocation,
- ) -> Bytes {
- Bytes {
- ptr,
- len,
- deallocation,
- }
- }
-
- fn as_slice(&self) -> &[u8] {
- self
- }
-
- #[inline]
- pub fn len(&self) -> usize {
- self.len
- }
-
- #[inline]
- pub fn is_empty(&self) -> bool {
- self.len == 0
- }
-
- #[inline]
- pub fn ptr(&self) -> NonNull<u8> {
- self.ptr
- }
-
- pub fn capacity(&self) -> usize {
- match self.deallocation {
- Deallocation::Native(capacity) => capacity,
- // we cannot determine this in general,
- // and thus we state that this is externally-owned memory
- Deallocation::Foreign(_) => 0,
- }
- }
-}
-
-impl Drop for Bytes {
- #[inline]
- fn drop(&mut self) {
- match &self.deallocation {
- Deallocation::Native(capacity) => {
- unsafe { alloc::free_aligned::<u8>(self.ptr, *capacity) };
- }
- // foreign interface knows how to deallocate itself.
- Deallocation::Foreign(_) => (),
- }
- }
-}
-
-impl std::ops::Deref for Bytes {
- type Target = [u8];
-
- fn deref(&self) -> &[u8] {
- unsafe { slice::from_raw_parts(self.ptr.as_ptr(), self.len) }
- }
-}
-
-impl PartialEq for Bytes {
- fn eq(&self, other: &Bytes) -> bool {
- self.as_slice() == other.as_slice()
- }
-}
-
-impl Debug for Bytes {
- fn fmt(&self, f: &mut Formatter) -> std::fmt::Result {
- write!(f, "Bytes {{ ptr: {:?}, len: {}, data: ", self.ptr, self.len,)?;
-
- f.debug_list().entries(self.iter()).finish()?;
-
- write!(f, " }}")
- }
-}
diff --git a/arrow/src/compute/kernels/aggregate.rs b/arrow/src/compute/kernels/aggregate.rs
deleted file mode 100644
index 2dd19c4..0000000
--- a/arrow/src/compute/kernels/aggregate.rs
+++ /dev/null
@@ -1,978 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-//! Defines aggregations over Arrow arrays.
-
-use multiversion::multiversion;
-use std::ops::Add;
-
-use crate::array::{
- Array, BooleanArray, GenericStringArray, PrimitiveArray, StringOffsetSizeTrait,
-};
-use crate::datatypes::{ArrowNativeType, ArrowNumericType};
-
-/// Generic test for NaN, the optimizer should be able to remove this for integer types.
-#[inline]
-fn is_nan<T: ArrowNativeType + PartialOrd + Copy>(a: T) -> bool {
- #[allow(clippy::eq_op)]
- !(a == a)
-}
-
-/// Helper macro to perform min/max of strings
-fn min_max_string<T: StringOffsetSizeTrait, F: Fn(&str, &str) -> bool>(
- array: &GenericStringArray<T>,
- cmp: F,
-) -> Option<&str> {
- let null_count = array.null_count();
-
- if null_count == array.len() {
- return None;
- }
- let data = array.data();
- let mut n;
- if null_count == 0 {
- n = array.value(0);
- for i in 1..data.len() {
- let item = array.value(i);
- if cmp(&n, item) {
- n = item;
- }
- }
- } else {
- n = "";
- let mut has_value = false;
-
- for i in 0..data.len() {
- let item = array.value(i);
- if data.is_valid(i) && (!has_value || cmp(&n, item)) {
- has_value = true;
- n = item;
- }
- }
- }
- Some(n)
-}
-
-/// Returns the minimum value in the array, according to the natural order.
-/// For floating point arrays any NaN values are considered to be greater than any other non-null value
-#[cfg(not(feature = "simd"))]
-pub fn min<T>(array: &PrimitiveArray<T>) -> Option<T::Native>
-where
- T: ArrowNumericType,
- T::Native: ArrowNativeType,
-{
- min_max_helper(array, |a, b| (is_nan(*a) & !is_nan(*b)) || a > b)
-}
-
-/// Returns the maximum value in the array, according to the natural order.
-/// For floating point arrays any NaN values are considered to be greater than any other non-null value
-#[cfg(not(feature = "simd"))]
-pub fn max<T>(array: &PrimitiveArray<T>) -> Option<T::Native>
-where
- T: ArrowNumericType,
- T::Native: ArrowNativeType,
-{
- min_max_helper(array, |a, b| (!is_nan(*a) & is_nan(*b)) || a < b)
-}
-
-/// Returns the maximum value in the string array, according to the natural order.
-pub fn max_string<T: StringOffsetSizeTrait>(
- array: &GenericStringArray<T>,
-) -> Option<&str> {
- min_max_string(array, |a, b| a < b)
-}
-
-/// Returns the minimum value in the string array, according to the natural order.
-pub fn min_string<T: StringOffsetSizeTrait>(
- array: &GenericStringArray<T>,
-) -> Option<&str> {
- min_max_string(array, |a, b| a > b)
-}
-
-/// Helper function to perform min/max lambda function on values from a numeric array.
-#[multiversion]
-#[clone(target = "x86_64+avx")]
-fn min_max_helper<T, F>(array: &PrimitiveArray<T>, cmp: F) -> Option<T::Native>
-where
- T: ArrowNumericType,
- F: Fn(&T::Native, &T::Native) -> bool,
-{
- let null_count = array.null_count();
-
- // Includes case array.len() == 0
- if null_count == array.len() {
- return None;
- }
-
- let data = array.data();
- let m = array.values();
- let mut n;
-
- if null_count == 0 {
- // optimized path for arrays without null values
- n = m[1..]
- .iter()
- .fold(m[0], |max, item| if cmp(&max, item) { *item } else { max });
- } else {
- n = T::default_value();
- let mut has_value = false;
- for (i, item) in m.iter().enumerate() {
- if data.is_valid(i) && (!has_value || cmp(&n, item)) {
- has_value = true;
- n = *item
- }
- }
- }
- Some(n)
-}
-
-/// Returns the minimum value in the boolean array.
-///
-/// ```
-/// use arrow::{
-/// array::BooleanArray,
-/// compute::min_boolean,
-/// };
-///
-/// let a = BooleanArray::from(vec![Some(true), None, Some(false)]);
-/// assert_eq!(min_boolean(&a), Some(false))
-/// ```
-pub fn min_boolean(array: &BooleanArray) -> Option<bool> {
- // short circuit if all nulls / zero length array
- if array.null_count() == array.len() {
- return None;
- }
-
- // Note the min bool is false (0), so short circuit as soon as we see it
- array
- .iter()
- .find(|&b| b == Some(false))
- .flatten()
- .or(Some(true))
-}
-
-/// Returns the maximum value in the boolean array
-///
-/// ```
-/// use arrow::{
-/// array::BooleanArray,
-/// compute::max_boolean,
-/// };
-///
-/// let a = BooleanArray::from(vec![Some(true), None, Some(false)]);
-/// assert_eq!(max_boolean(&a), Some(true))
-/// ```
-pub fn max_boolean(array: &BooleanArray) -> Option<bool> {
- // short circuit if all nulls / zero length array
- if array.null_count() == array.len() {
- return None;
- }
-
- // Note the max bool is true (1), so short circuit as soon as we see it
- array
- .iter()
- .find(|&b| b == Some(true))
- .flatten()
- .or(Some(false))
-}
-
-/// Returns the sum of values in the array.
-///
-/// Returns `None` if the array is empty or only contains null values.
-#[cfg(not(feature = "simd"))]
-pub fn sum<T>(array: &PrimitiveArray<T>) -> Option<T::Native>
-where
- T: ArrowNumericType,
- T::Native: Add<Output = T::Native>,
-{
- let null_count = array.null_count();
-
- if null_count == array.len() {
- return None;
- }
-
- let data: &[T::Native] = array.values();
-
- match array.data().null_buffer() {
- None => {
- let sum = data.iter().fold(T::default_value(), |accumulator, value| {
- accumulator + *value
- });
-
- Some(sum)
- }
- Some(buffer) => {
- let mut sum = T::default_value();
- let data_chunks = data.chunks_exact(64);
- let remainder = data_chunks.remainder();
-
- let bit_chunks = buffer.bit_chunks(array.offset(), array.len());
- data_chunks
- .zip(bit_chunks.iter())
- .for_each(|(chunk, mask)| {
- // index_mask has value 1 << i in the loop
- let mut index_mask = 1;
- chunk.iter().for_each(|value| {
- if (mask & index_mask) != 0 {
- sum = sum + *value;
- }
- index_mask <<= 1;
- });
- });
-
- let remainder_bits = bit_chunks.remainder_bits();
-
- remainder.iter().enumerate().for_each(|(i, value)| {
- if remainder_bits & (1 << i) != 0 {
- sum = sum + *value;
- }
- });
-
- Some(sum)
- }
- }
-}
-
-#[cfg(feature = "simd")]
-mod simd {
- use super::is_nan;
- use crate::array::{Array, PrimitiveArray};
- use crate::datatypes::ArrowNumericType;
- use std::marker::PhantomData;
- use std::ops::Add;
-
- pub(super) trait SimdAggregate<T: ArrowNumericType> {
- type ScalarAccumulator;
- type SimdAccumulator;
-
- /// Returns the accumulator for aggregating scalar values
- fn init_accumulator_scalar() -> Self::ScalarAccumulator;
-
- /// Returns the accumulator for aggregating simd chunks of values
- fn init_accumulator_chunk() -> Self::SimdAccumulator;
-
- /// Updates the accumulator with the values of one chunk
- fn accumulate_chunk_non_null(
- accumulator: &mut Self::SimdAccumulator,
- chunk: T::Simd,
- );
-
- /// Updates the accumulator with the values of one chunk according to the given vector mask
- fn accumulate_chunk_nullable(
- accumulator: &mut Self::SimdAccumulator,
- chunk: T::Simd,
- mask: T::SimdMask,
- );
-
- /// Updates the accumulator with one value
- fn accumulate_scalar(accumulator: &mut Self::ScalarAccumulator, value: T::Native);
-
- /// Reduces the vector lanes of the simd accumulator and the scalar accumulator to a single value
- fn reduce(
- simd_accumulator: Self::SimdAccumulator,
- scalar_accumulator: Self::ScalarAccumulator,
- ) -> Option<T::Native>;
- }
-
- pub(super) struct SumAggregate<T: ArrowNumericType> {
- phantom: PhantomData<T>,
- }
-
- impl<T: ArrowNumericType> SimdAggregate<T> for SumAggregate<T>
- where
- T::Native: Add<Output = T::Native>,
- {
- type ScalarAccumulator = T::Native;
- type SimdAccumulator = T::Simd;
-
- fn init_accumulator_scalar() -> Self::ScalarAccumulator {
- T::default_value()
- }
-
- fn init_accumulator_chunk() -> Self::SimdAccumulator {
- T::init(Self::init_accumulator_scalar())
- }
-
- fn accumulate_chunk_non_null(accumulator: &mut T::Simd, chunk: T::Simd) {
- *accumulator = *accumulator + chunk;
- }
-
- fn accumulate_chunk_nullable(
- accumulator: &mut T::Simd,
- chunk: T::Simd,
- vecmask: T::SimdMask,
- ) {
- let zero = T::init(T::default_value());
- let blended = T::mask_select(vecmask, chunk, zero);
-
- *accumulator = *accumulator + blended;
- }
-
- fn accumulate_scalar(accumulator: &mut T::Native, value: T::Native) {
- *accumulator = *accumulator + value
- }
-
- fn reduce(
- simd_accumulator: Self::SimdAccumulator,
- scalar_accumulator: Self::ScalarAccumulator,
- ) -> Option<T::Native> {
- // we can't use T::lanes() as the slice len because it is not const,
- // instead always reserve the maximum number of lanes
- let mut tmp = [T::default_value(); 64];
- let slice = &mut tmp[0..T::lanes()];
- T::write(simd_accumulator, slice);
-
- let mut reduced = Self::init_accumulator_scalar();
- slice
- .iter()
- .for_each(|value| Self::accumulate_scalar(&mut reduced, *value));
-
- Self::accumulate_scalar(&mut reduced, scalar_accumulator);
-
- // result can not be None because we checked earlier for the null count
- Some(reduced)
- }
- }
-
- pub(super) struct MinAggregate<T: ArrowNumericType> {
- phantom: PhantomData<T>,
- }
-
- impl<T: ArrowNumericType> SimdAggregate<T> for MinAggregate<T>
- where
- T::Native: PartialOrd,
- {
- type ScalarAccumulator = (T::Native, bool);
- type SimdAccumulator = (T::Simd, T::SimdMask);
-
- fn init_accumulator_scalar() -> Self::ScalarAccumulator {
- (T::default_value(), false)
- }
-
- fn init_accumulator_chunk() -> Self::SimdAccumulator {
- (T::init(T::default_value()), T::mask_init(false))
- }
-
- fn accumulate_chunk_non_null(
- accumulator: &mut Self::SimdAccumulator,
- chunk: T::Simd,
- ) {
- let acc_is_nan = !T::eq(accumulator.0, accumulator.0);
- let is_lt = acc_is_nan | T::lt(chunk, accumulator.0);
- let first_or_lt = !accumulator.1 | is_lt;
-
- accumulator.0 = T::mask_select(first_or_lt, chunk, accumulator.0);
- accumulator.1 = T::mask_init(true);
- }
-
- fn accumulate_chunk_nullable(
- accumulator: &mut Self::SimdAccumulator,
- chunk: T::Simd,
- vecmask: T::SimdMask,
- ) {
- let acc_is_nan = !T::eq(accumulator.0, accumulator.0);
- let is_lt = vecmask & (acc_is_nan | T::lt(chunk, accumulator.0));
- let first_or_lt = !accumulator.1 | is_lt;
-
- accumulator.0 = T::mask_select(first_or_lt, chunk, accumulator.0);
- accumulator.1 |= vecmask;
- }
-
- fn accumulate_scalar(
- accumulator: &mut Self::ScalarAccumulator,
- value: T::Native,
- ) {
- if !accumulator.1 {
- accumulator.0 = value;
- } else {
- let acc_is_nan = is_nan(accumulator.0);
- if acc_is_nan || value < accumulator.0 {
- accumulator.0 = value
- }
- }
- accumulator.1 = true
- }
-
- fn reduce(
- simd_accumulator: Self::SimdAccumulator,
- scalar_accumulator: Self::ScalarAccumulator,
- ) -> Option<T::Native> {
- // we can't use T::lanes() as the slice len because it is not const,
- // instead always reserve the maximum number of lanes
- let mut tmp = [T::default_value(); 64];
- let slice = &mut tmp[0..T::lanes()];
- T::write(simd_accumulator.0, slice);
-
- let mut reduced = Self::init_accumulator_scalar();
- slice
- .iter()
- .enumerate()
- .filter(|(i, _value)| T::mask_get(&simd_accumulator.1, *i))
- .for_each(|(_i, value)| Self::accumulate_scalar(&mut reduced, *value));
-
- if scalar_accumulator.1 {
- Self::accumulate_scalar(&mut reduced, scalar_accumulator.0);
- }
-
- if reduced.1 {
- Some(reduced.0)
- } else {
- None
- }
- }
- }
-
- pub(super) struct MaxAggregate<T: ArrowNumericType> {
- phantom: PhantomData<T>,
- }
-
- impl<T: ArrowNumericType> SimdAggregate<T> for MaxAggregate<T>
- where
- T::Native: PartialOrd,
- {
- type ScalarAccumulator = (T::Native, bool);
- type SimdAccumulator = (T::Simd, T::SimdMask);
-
- fn init_accumulator_scalar() -> Self::ScalarAccumulator {
- (T::default_value(), false)
- }
-
- fn init_accumulator_chunk() -> Self::SimdAccumulator {
- (T::init(T::default_value()), T::mask_init(false))
- }
-
- fn accumulate_chunk_non_null(
- accumulator: &mut Self::SimdAccumulator,
- chunk: T::Simd,
- ) {
- let chunk_is_nan = !T::eq(chunk, chunk);
- let is_gt = chunk_is_nan | T::gt(chunk, accumulator.0);
- let first_or_gt = !accumulator.1 | is_gt;
-
- accumulator.0 = T::mask_select(first_or_gt, chunk, accumulator.0);
- accumulator.1 = T::mask_init(true);
- }
-
- fn accumulate_chunk_nullable(
- accumulator: &mut Self::SimdAccumulator,
- chunk: T::Simd,
- vecmask: T::SimdMask,
- ) {
- let chunk_is_nan = !T::eq(chunk, chunk);
- let is_gt = vecmask & (chunk_is_nan | T::gt(chunk, accumulator.0));
- let first_or_gt = !accumulator.1 | is_gt;
-
- accumulator.0 = T::mask_select(first_or_gt, chunk, accumulator.0);
- accumulator.1 |= vecmask;
- }
-
- fn accumulate_scalar(
- accumulator: &mut Self::ScalarAccumulator,
- value: T::Native,
- ) {
- if !accumulator.1 {
- accumulator.0 = value;
- } else {
- let value_is_nan = is_nan(value);
- if value_is_nan || value > accumulator.0 {
- accumulator.0 = value
- }
- }
- accumulator.1 = true;
- }
-
- fn reduce(
- simd_accumulator: Self::SimdAccumulator,
- scalar_accumulator: Self::ScalarAccumulator,
- ) -> Option<T::Native> {
- // we can't use T::lanes() as the slice len because it is not const,
- // instead always reserve the maximum number of lanes
- let mut tmp = [T::default_value(); 64];
- let slice = &mut tmp[0..T::lanes()];
- T::write(simd_accumulator.0, slice);
-
- let mut reduced = Self::init_accumulator_scalar();
- slice
- .iter()
- .enumerate()
- .filter(|(i, _value)| T::mask_get(&simd_accumulator.1, *i))
- .for_each(|(_i, value)| Self::accumulate_scalar(&mut reduced, *value));
-
- if scalar_accumulator.1 {
- Self::accumulate_scalar(&mut reduced, scalar_accumulator.0);
- }
-
- if reduced.1 {
- Some(reduced.0)
- } else {
- None
- }
- }
- }
-
- pub(super) fn simd_aggregation<T: ArrowNumericType, A: SimdAggregate<T>>(
- array: &PrimitiveArray<T>,
- ) -> Option<T::Native> {
- let null_count = array.null_count();
-
- if null_count == array.len() {
- return None;
- }
-
- let data: &[T::Native] = array.values();
-
- let mut chunk_acc = A::init_accumulator_chunk();
- let mut rem_acc = A::init_accumulator_scalar();
-
- match array.data().null_buffer() {
- None => {
- let data_chunks = data.chunks_exact(64);
- let remainder = data_chunks.remainder();
-
- data_chunks.for_each(|chunk| {
- chunk.chunks_exact(T::lanes()).for_each(|chunk| {
- let chunk = T::load(&chunk);
- A::accumulate_chunk_non_null(&mut chunk_acc, chunk);
- });
- });
-
- remainder.iter().for_each(|value| {
- A::accumulate_scalar(&mut rem_acc, *value);
- });
- }
- Some(buffer) => {
- // process data in chunks of 64 elements since we also get 64 bits of validity information at a time
- let data_chunks = data.chunks_exact(64);
- let remainder = data_chunks.remainder();
-
- let bit_chunks = buffer.bit_chunks(array.offset(), array.len());
- let remainder_bits = bit_chunks.remainder_bits();
-
- data_chunks.zip(bit_chunks).for_each(|(chunk, mut mask)| {
- // split chunks further into slices corresponding to the vector length
- // the compiler is able to unroll this inner loop and remove bounds checks
- // since the outer chunk size (64) is always a multiple of the number of lanes
- chunk.chunks_exact(T::lanes()).for_each(|chunk| {
- let vecmask = T::mask_from_u64(mask);
- let chunk = T::load(&chunk);
-
- A::accumulate_chunk_nullable(&mut chunk_acc, chunk, vecmask);
-
- // skip the shift and avoid overflow for u8 type, which uses 64 lanes.
- mask >>= T::lanes() % 64;
- });
- });
-
- remainder.iter().enumerate().for_each(|(i, value)| {
- if remainder_bits & (1 << i) != 0 {
- A::accumulate_scalar(&mut rem_acc, *value)
- }
- });
- }
- }
-
- A::reduce(chunk_acc, rem_acc)
- }
-}
-
-/// Returns the sum of values in the array.
-///
-/// Returns `None` if the array is empty or only contains null values.
-#[cfg(feature = "simd")]
-pub fn sum<T: ArrowNumericType>(array: &PrimitiveArray<T>) -> Option<T::Native>
-where
- T::Native: Add<Output = T::Native>,
-{
- use simd::*;
-
- simd::simd_aggregation::<T, SumAggregate<T>>(&array)
-}
-
-#[cfg(feature = "simd")]
-/// Returns the minimum value in the array, according to the natural order.
-/// For floating point arrays any NaN values are considered to be greater than any other non-null value
-pub fn min<T: ArrowNumericType>(array: &PrimitiveArray<T>) -> Option<T::Native>
-where
- T::Native: PartialOrd,
-{
- use simd::*;
-
- simd::simd_aggregation::<T, MinAggregate<T>>(&array)
-}
-
-#[cfg(feature = "simd")]
-/// Returns the maximum value in the array, according to the natural order.
-/// For floating point arrays any NaN values are considered to be greater than any other non-null value
-pub fn max<T: ArrowNumericType>(array: &PrimitiveArray<T>) -> Option<T::Native>
-where
- T::Native: PartialOrd,
-{
- use simd::*;
-
- simd::simd_aggregation::<T, MaxAggregate<T>>(&array)
-}
-
-#[cfg(test)]
-mod tests {
- use super::*;
- use crate::array::*;
- use crate::compute::add;
-
- #[test]
- fn test_primitive_array_sum() {
- let a = Int32Array::from(vec![1, 2, 3, 4, 5]);
- assert_eq!(15, sum(&a).unwrap());
- }
-
- #[test]
- fn test_primitive_array_float_sum() {
- let a = Float64Array::from(vec![1.1, 2.2, 3.3, 4.4, 5.5]);
- assert!(16.5 - sum(&a).unwrap() < f64::EPSILON);
- }
-
- #[test]
- fn test_primitive_array_sum_with_nulls() {
- let a = Int32Array::from(vec![None, Some(2), Some(3), None, Some(5)]);
- assert_eq!(10, sum(&a).unwrap());
- }
-
- #[test]
- fn test_primitive_array_sum_all_nulls() {
- let a = Int32Array::from(vec![None, None, None]);
- assert_eq!(None, sum(&a));
- }
-
- #[test]
- fn test_primitive_array_sum_large_64() {
- let a: Int64Array = (1..=100)
- .map(|i| if i % 3 == 0 { Some(i) } else { None })
- .collect();
- let b: Int64Array = (1..=100)
- .map(|i| if i % 3 == 0 { Some(0) } else { Some(i) })
- .collect();
- // create an array that actually has non-zero values at the invalid indices
- let c = add(&a, &b).unwrap();
- assert_eq!(Some((1..=100).filter(|i| i % 3 == 0).sum()), sum(&c));
- }
-
- #[test]
- fn test_primitive_array_sum_large_32() {
- let a: Int32Array = (1..=100)
- .map(|i| if i % 3 == 0 { Some(i) } else { None })
- .collect();
- let b: Int32Array = (1..=100)
- .map(|i| if i % 3 == 0 { Some(0) } else { Some(i) })
- .collect();
- // create an array that actually has non-zero values at the invalid indices
- let c = add(&a, &b).unwrap();
- assert_eq!(Some((1..=100).filter(|i| i % 3 == 0).sum()), sum(&c));
- }
-
- #[test]
- fn test_primitive_array_sum_large_16() {
- let a: Int16Array = (1..=100)
- .map(|i| if i % 3 == 0 { Some(i) } else { None })
- .collect();
- let b: Int16Array = (1..=100)
- .map(|i| if i % 3 == 0 { Some(0) } else { Some(i) })
- .collect();
- // create an array that actually has non-zero values at the invalid indices
- let c = add(&a, &b).unwrap();
- assert_eq!(Some((1..=100).filter(|i| i % 3 == 0).sum()), sum(&c));
- }
-
- #[test]
- fn test_primitive_array_sum_large_8() {
- // include fewer values than other large tests so the result does not overflow the u8
- let a: UInt8Array = (1..=100)
- .map(|i| if i % 33 == 0 { Some(i) } else { None })
- .collect();
- let b: UInt8Array = (1..=100)
- .map(|i| if i % 33 == 0 { Some(0) } else { Some(i) })
- .collect();
- // create an array that actually has non-zero values at the invalid indices
- let c = add(&a, &b).unwrap();
- assert_eq!(Some((1..=100).filter(|i| i % 33 == 0).sum()), sum(&c));
- }
-
- #[test]
- fn test_primitive_array_min_max() {
- let a = Int32Array::from(vec![5, 6, 7, 8, 9]);
- assert_eq!(5, min(&a).unwrap());
- assert_eq!(9, max(&a).unwrap());
- }
-
- #[test]
- fn test_primitive_array_min_max_with_nulls() {
- let a = Int32Array::from(vec![Some(5), None, None, Some(8), Some(9)]);
- assert_eq!(5, min(&a).unwrap());
- assert_eq!(9, max(&a).unwrap());
- }
-
- #[test]
- fn test_primitive_min_max_1() {
- let a = Int32Array::from(vec![None, None, Some(5), Some(2)]);
- assert_eq!(Some(2), min(&a));
- assert_eq!(Some(5), max(&a));
- }
-
- #[test]
- fn test_primitive_min_max_float_large_nonnull_array() {
- let a: Float64Array = (0..256).map(|i| Some((i + 1) as f64)).collect();
- // min/max are on boundaries of chunked data
- assert_eq!(Some(1.0), min(&a));
- assert_eq!(Some(256.0), max(&a));
-
- // max is last value in remainder after chunking
- let a: Float64Array = (0..255).map(|i| Some((i + 1) as f64)).collect();
- assert_eq!(Some(255.0), max(&a));
-
- // max is first value in remainder after chunking
- let a: Float64Array = (0..257).map(|i| Some((i + 1) as f64)).collect();
- assert_eq!(Some(257.0), max(&a));
- }
-
- #[test]
- fn test_primitive_min_max_float_large_nullable_array() {
- let a: Float64Array = (0..256)
- .map(|i| {
- if (i + 1) % 3 == 0 {
- None
- } else {
- Some((i + 1) as f64)
- }
- })
- .collect();
- // min/max are on boundaries of chunked data
- assert_eq!(Some(1.0), min(&a));
- assert_eq!(Some(256.0), max(&a));
-
- let a: Float64Array = (0..256)
- .map(|i| {
- if i == 0 || i == 255 {
- None
- } else {
- Some((i + 1) as f64)
- }
- })
- .collect();
- // boundaries of chunked data are null
- assert_eq!(Some(2.0), min(&a));
- assert_eq!(Some(255.0), max(&a));
-
- let a: Float64Array = (0..256)
- .map(|i| if i != 100 { None } else { Some((i) as f64) })
- .collect();
- // a single non-null value somewhere in the middle
- assert_eq!(Some(100.0), min(&a));
- assert_eq!(Some(100.0), max(&a));
-
- // max is last value in remainder after chunking
- let a: Float64Array = (0..255).map(|i| Some((i + 1) as f64)).collect();
- assert_eq!(Some(255.0), max(&a));
-
- // max is first value in remainder after chunking
- let a: Float64Array = (0..257).map(|i| Some((i + 1) as f64)).collect();
- assert_eq!(Some(257.0), max(&a));
- }
-
- #[test]
- fn test_primitive_min_max_float_edge_cases() {
- let a: Float64Array = (0..100).map(|_| Some(f64::NEG_INFINITY)).collect();
- assert_eq!(Some(f64::NEG_INFINITY), min(&a));
- assert_eq!(Some(f64::NEG_INFINITY), max(&a));
-
- let a: Float64Array = (0..100).map(|_| Some(f64::MIN)).collect();
- assert_eq!(Some(f64::MIN), min(&a));
- assert_eq!(Some(f64::MIN), max(&a));
-
- let a: Float64Array = (0..100).map(|_| Some(f64::MAX)).collect();
- assert_eq!(Some(f64::MAX), min(&a));
- assert_eq!(Some(f64::MAX), max(&a));
-
- let a: Float64Array = (0..100).map(|_| Some(f64::INFINITY)).collect();
- assert_eq!(Some(f64::INFINITY), min(&a));
- assert_eq!(Some(f64::INFINITY), max(&a));
- }
-
- #[test]
- fn test_primitive_min_max_float_all_nans_non_null() {
- let a: Float64Array = (0..100).map(|_| Some(f64::NAN)).collect();
- assert!(max(&a).unwrap().is_nan());
- assert!(min(&a).unwrap().is_nan());
- }
-
- #[test]
- fn test_primitive_min_max_float_first_nan_nonnull() {
- let a: Float64Array = (0..100)
- .map(|i| {
- if i == 0 {
- Some(f64::NAN)
- } else {
- Some(i as f64)
- }
- })
- .collect();
- assert_eq!(Some(1.0), min(&a));
- assert!(max(&a).unwrap().is_nan());
- }
-
- #[test]
- fn test_primitive_min_max_float_last_nan_nonnull() {
- let a: Float64Array = (0..100)
- .map(|i| {
- if i == 99 {
- Some(f64::NAN)
- } else {
- Some((i + 1) as f64)
- }
- })
- .collect();
- assert_eq!(Some(1.0), min(&a));
- assert!(max(&a).unwrap().is_nan());
- }
-
- #[test]
- fn test_primitive_min_max_float_first_nan_nullable() {
- let a: Float64Array = (0..100)
- .map(|i| {
- if i == 0 {
- Some(f64::NAN)
- } else if i % 2 == 0 {
- None
- } else {
- Some(i as f64)
- }
- })
- .collect();
- assert_eq!(Some(1.0), min(&a));
- assert!(max(&a).unwrap().is_nan());
- }
-
- #[test]
- fn test_primitive_min_max_float_last_nan_nullable() {
- let a: Float64Array = (0..100)
- .map(|i| {
- if i == 99 {
- Some(f64::NAN)
- } else if i % 2 == 0 {
- None
- } else {
- Some(i as f64)
- }
- })
- .collect();
- assert_eq!(Some(1.0), min(&a));
- assert!(max(&a).unwrap().is_nan());
- }
-
- #[test]
- fn test_primitive_min_max_float_inf_and_nans() {
- let a: Float64Array = (0..100)
- .map(|i| {
- let x = match i % 10 {
- 0 => f64::NEG_INFINITY,
- 1 => f64::MIN,
- 2 => f64::MAX,
- 4 => f64::INFINITY,
- 5 => f64::NAN,
- _ => i as f64,
- };
- Some(x)
- })
- .collect();
- assert_eq!(Some(f64::NEG_INFINITY), min(&a));
- assert!(max(&a).unwrap().is_nan());
- }
-
- #[test]
- fn test_string_min_max_with_nulls() {
- let a = StringArray::from(vec![Some("b"), None, None, Some("a"), Some("c")]);
- assert_eq!("a", min_string(&a).unwrap());
- assert_eq!("c", max_string(&a).unwrap());
- }
-
- #[test]
- fn test_string_min_max_all_nulls() {
- let a = StringArray::from(vec![None, None]);
- assert_eq!(None, min_string(&a));
- assert_eq!(None, max_string(&a));
- }
-
- #[test]
- fn test_string_min_max_1() {
- let a = StringArray::from(vec![None, None, Some("b"), Some("a")]);
- assert_eq!(Some("a"), min_string(&a));
- assert_eq!(Some("b"), max_string(&a));
- }
-
- #[test]
- fn test_boolean_min_max_empty() {
- let a = BooleanArray::from(vec![] as Vec<Option<bool>>);
- assert_eq!(None, min_boolean(&a));
- assert_eq!(None, max_boolean(&a));
- }
-
- #[test]
- fn test_boolean_min_max_all_null() {
- let a = BooleanArray::from(vec![None, None]);
- assert_eq!(None, min_boolean(&a));
- assert_eq!(None, max_boolean(&a));
- }
-
- #[test]
- fn test_boolean_min_max_no_null() {
- let a = BooleanArray::from(vec![Some(true), Some(false), Some(true)]);
- assert_eq!(Some(false), min_boolean(&a));
- assert_eq!(Some(true), max_boolean(&a));
- }
-
- #[test]
- fn test_boolean_min_max() {
- let a = BooleanArray::from(vec![Some(true), Some(true), None, Some(false), None]);
- assert_eq!(Some(false), min_boolean(&a));
- assert_eq!(Some(true), max_boolean(&a));
-
- let a = BooleanArray::from(vec![None, Some(true), None, Some(false), None]);
- assert_eq!(Some(false), min_boolean(&a));
- assert_eq!(Some(true), max_boolean(&a));
-
- let a =
- BooleanArray::from(vec![Some(false), Some(true), None, Some(false), None]);
- assert_eq!(Some(false), min_boolean(&a));
- assert_eq!(Some(true), max_boolean(&a));
- }
-
- #[test]
- fn test_boolean_min_max_smaller() {
- let a = BooleanArray::from(vec![Some(false)]);
- assert_eq!(Some(false), min_boolean(&a));
- assert_eq!(Some(false), max_boolean(&a));
-
- let a = BooleanArray::from(vec![None, Some(false)]);
- assert_eq!(Some(false), min_boolean(&a));
- assert_eq!(Some(false), max_boolean(&a));
-
- let a = BooleanArray::from(vec![None, Some(true)]);
- assert_eq!(Some(true), min_boolean(&a));
- assert_eq!(Some(true), max_boolean(&a));
-
- let a = BooleanArray::from(vec![Some(true)]);
- assert_eq!(Some(true), min_boolean(&a));
- assert_eq!(Some(true), max_boolean(&a));
- }
-}
diff --git a/arrow/src/compute/kernels/arithmetic.rs b/arrow/src/compute/kernels/arithmetic.rs
deleted file mode 100644
index 50c06b0..0000000
--- a/arrow/src/compute/kernels/arithmetic.rs
+++ /dev/null
@@ -1,1539 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-//! Defines basic arithmetic kernels for `PrimitiveArrays`.
-//!
-//! These kernels can leverage SIMD if available on your system. Currently no runtime
-//! detection is provided, you should enable the specific SIMD intrinsics using
-//! `RUSTFLAGS="-C target-feature=+avx2"` for example. See the documentation
-//! [here](https://doc.rust-lang.org/stable/core/arch/) for more information.
-
-use std::ops::{Add, Div, Mul, Neg, Rem, Sub};
-
-use num::{One, Zero};
-
-use crate::buffer::Buffer;
-#[cfg(feature = "simd")]
-use crate::buffer::MutableBuffer;
-#[cfg(not(feature = "simd"))]
-use crate::compute::kernels::arity::unary;
-use crate::compute::util::combine_option_bitmap;
-use crate::datatypes;
-use crate::datatypes::ArrowNumericType;
-use crate::error::{ArrowError, Result};
-use crate::{array::*, util::bit_util};
-use num::traits::Pow;
-#[cfg(feature = "simd")]
-use std::borrow::BorrowMut;
-#[cfg(feature = "simd")]
-use std::slice::{ChunksExact, ChunksExactMut};
-
-/// SIMD vectorized version of `unary_math_op` above specialized for signed numerical values.
-#[cfg(feature = "simd")]
-fn simd_signed_unary_math_op<T, SIMD_OP, SCALAR_OP>(
- array: &PrimitiveArray<T>,
- simd_op: SIMD_OP,
- scalar_op: SCALAR_OP,
-) -> Result<PrimitiveArray<T>>
-where
- T: datatypes::ArrowSignedNumericType,
- SIMD_OP: Fn(T::SignedSimd) -> T::SignedSimd,
- SCALAR_OP: Fn(T::Native) -> T::Native,
-{
- let lanes = T::lanes();
- let buffer_size = array.len() * std::mem::size_of::<T::Native>();
- let mut result = MutableBuffer::new(buffer_size).with_bitset(buffer_size, false);
-
- let mut result_chunks = result.typed_data_mut().chunks_exact_mut(lanes);
- let mut array_chunks = array.values().chunks_exact(lanes);
-
- result_chunks
- .borrow_mut()
- .zip(array_chunks.borrow_mut())
- .for_each(|(result_slice, input_slice)| {
- let simd_input = T::load_signed(input_slice);
- let simd_result = T::signed_unary_op(simd_input, &simd_op);
- T::write_signed(simd_result, result_slice);
- });
-
- let result_remainder = result_chunks.into_remainder();
- let array_remainder = array_chunks.remainder();
-
- result_remainder.into_iter().zip(array_remainder).for_each(
- |(scalar_result, scalar_input)| {
- *scalar_result = scalar_op(*scalar_input);
- },
- );
-
- let data = ArrayData::new(
- T::DATA_TYPE,
- array.len(),
- None,
- array.data_ref().null_buffer().cloned(),
- 0,
- vec![result.into()],
- vec![],
- );
- Ok(PrimitiveArray::<T>::from(data))
-}
-
-#[cfg(feature = "simd")]
-fn simd_float_unary_math_op<T, SIMD_OP, SCALAR_OP>(
- array: &PrimitiveArray<T>,
- simd_op: SIMD_OP,
- scalar_op: SCALAR_OP,
-) -> Result<PrimitiveArray<T>>
-where
- T: datatypes::ArrowFloatNumericType,
- SIMD_OP: Fn(T::Simd) -> T::Simd,
- SCALAR_OP: Fn(T::Native) -> T::Native,
-{
- let lanes = T::lanes();
- let buffer_size = array.len() * std::mem::size_of::<T::Native>();
-
- let mut result = MutableBuffer::new(buffer_size).with_bitset(buffer_size, false);
-
- let mut result_chunks = result.typed_data_mut().chunks_exact_mut(lanes);
- let mut array_chunks = array.values().chunks_exact(lanes);
-
- result_chunks
- .borrow_mut()
- .zip(array_chunks.borrow_mut())
- .for_each(|(result_slice, input_slice)| {
- let simd_input = T::load(input_slice);
- let simd_result = T::unary_op(simd_input, &simd_op);
- T::write(simd_result, result_slice);
- });
-
- let result_remainder = result_chunks.into_remainder();
- let array_remainder = array_chunks.remainder();
-
- result_remainder.into_iter().zip(array_remainder).for_each(
- |(scalar_result, scalar_input)| {
- *scalar_result = scalar_op(*scalar_input);
- },
- );
-
- let data = ArrayData::new(
- T::DATA_TYPE,
- array.len(),
- None,
- array.data_ref().null_buffer().cloned(),
- 0,
- vec![result.into()],
- vec![],
- );
- Ok(PrimitiveArray::<T>::from(data))
-}
-
-/// Helper function to perform math lambda function on values from two arrays. If either
-/// left or right value is null then the output value is also null, so `1 + null` is
-/// `null`.
-///
-/// # Errors
-///
-/// This function errors if the arrays have different lengths
-pub fn math_op<T, F>(
- left: &PrimitiveArray<T>,
- right: &PrimitiveArray<T>,
- op: F,
-) -> Result<PrimitiveArray<T>>
-where
- T: ArrowNumericType,
- F: Fn(T::Native, T::Native) -> T::Native,
-{
- if left.len() != right.len() {
- return Err(ArrowError::ComputeError(
- "Cannot perform math operation on arrays of different length".to_string(),
- ));
- }
-
- let null_bit_buffer =
- combine_option_bitmap(left.data_ref(), right.data_ref(), left.len())?;
-
- let values = left
- .values()
- .iter()
- .zip(right.values().iter())
- .map(|(l, r)| op(*l, *r));
- // JUSTIFICATION
- // Benefit
- // ~60% speedup
- // Soundness
- // `values` is an iterator with a known size.
- let buffer = unsafe { Buffer::from_trusted_len_iter(values) };
-
- let data = ArrayData::new(
- T::DATA_TYPE,
- left.len(),
- None,
- null_bit_buffer,
- 0,
- vec![buffer],
- vec![],
- );
- Ok(PrimitiveArray::<T>::from(data))
-}
-
-/// Helper function to modulus two arrays.
-///
-/// # Errors
-///
-/// This function errors if:
-/// * the arrays have different lengths
-/// * a division by zero is found
-fn math_modulus<T>(
- left: &PrimitiveArray<T>,
- right: &PrimitiveArray<T>,
-) -> Result<PrimitiveArray<T>>
-where
- T: ArrowNumericType,
- T::Native: Rem<Output = T::Native> + Zero,
-{
- if left.len() != right.len() {
- return Err(ArrowError::ComputeError(
- "Cannot perform math operation on arrays of different length".to_string(),
- ));
- }
-
- let null_bit_buffer =
- combine_option_bitmap(left.data_ref(), right.data_ref(), left.len())?;
-
- let buffer = if let Some(b) = &null_bit_buffer {
- let values = left.values().iter().zip(right.values()).enumerate().map(
- |(i, (left, right))| {
- let is_valid = unsafe { bit_util::get_bit_raw(b.as_ptr(), i) };
- if is_valid {
- if right.is_zero() {
- Err(ArrowError::DivideByZero)
- } else {
- Ok(*left % *right)
- }
- } else {
- Ok(T::default_value())
- }
- },
- );
- unsafe { Buffer::try_from_trusted_len_iter(values) }
- } else {
- // no value is null
- let values = left
- .values()
- .iter()
- .zip(right.values())
- .map(|(left, right)| {
- if right.is_zero() {
- Err(ArrowError::DivideByZero)
- } else {
- Ok(*left % *right)
- }
- });
- unsafe { Buffer::try_from_trusted_len_iter(values) }
- }?;
-
- let data = ArrayData::new(
- T::DATA_TYPE,
- left.len(),
- None,
- null_bit_buffer,
- 0,
- vec![buffer],
- vec![],
- );
- Ok(PrimitiveArray::<T>::from(data))
-}
-
-/// Helper function to divide two arrays.
-///
-/// # Errors
-///
-/// This function errors if:
-/// * the arrays have different lengths
-/// * a division by zero is found
-fn math_divide<T>(
- left: &PrimitiveArray<T>,
- right: &PrimitiveArray<T>,
-) -> Result<PrimitiveArray<T>>
-where
- T: ArrowNumericType,
- T::Native: Div<Output = T::Native> + Zero,
-{
- if left.len() != right.len() {
- return Err(ArrowError::ComputeError(
- "Cannot perform math operation on arrays of different length".to_string(),
- ));
- }
-
- let null_bit_buffer =
- combine_option_bitmap(left.data_ref(), right.data_ref(), left.len())?;
-
- let buffer = if let Some(b) = &null_bit_buffer {
- let values = left.values().iter().zip(right.values()).enumerate().map(
- |(i, (left, right))| {
- let is_valid = unsafe { bit_util::get_bit_raw(b.as_ptr(), i) };
- if is_valid {
- if right.is_zero() {
- Err(ArrowError::DivideByZero)
- } else {
- Ok(*left / *right)
- }
- } else {
- Ok(T::default_value())
- }
- },
- );
- unsafe { Buffer::try_from_trusted_len_iter(values) }
- } else {
- // no value is null
- let values = left
- .values()
- .iter()
- .zip(right.values())
- .map(|(left, right)| {
- if right.is_zero() {
- Err(ArrowError::DivideByZero)
- } else {
- Ok(*left / *right)
- }
- });
- unsafe { Buffer::try_from_trusted_len_iter(values) }
- }?;
-
- let data = ArrayData::new(
- T::DATA_TYPE,
- left.len(),
- None,
- null_bit_buffer,
- 0,
- vec![buffer],
- vec![],
- );
- Ok(PrimitiveArray::<T>::from(data))
-}
-
-/// Scalar-modulo version of `math_modulus`.
-fn math_modulus_scalar<T>(
- array: &PrimitiveArray<T>,
- modulo: T::Native,
-) -> Result<PrimitiveArray<T>>
-where
- T: ArrowNumericType,
- T::Native: Rem<Output = T::Native> + Zero,
-{
- if modulo.is_zero() {
- return Err(ArrowError::DivideByZero);
- }
-
- let values = array.values().iter().map(|value| *value % modulo);
- let buffer = unsafe { Buffer::from_trusted_len_iter(values) };
-
- let data = ArrayData::new(
- T::DATA_TYPE,
- array.len(),
- None,
- array.data_ref().null_buffer().cloned(),
- 0,
- vec![buffer],
- vec![],
- );
- Ok(PrimitiveArray::<T>::from(data))
-}
-
-/// Scalar-divisor version of `math_divide`.
-fn math_divide_scalar<T>(
- array: &PrimitiveArray<T>,
- divisor: T::Native,
-) -> Result<PrimitiveArray<T>>
-where
- T: ArrowNumericType,
- T::Native: Div<Output = T::Native> + Zero,
-{
- if divisor.is_zero() {
- return Err(ArrowError::DivideByZero);
- }
-
- let values = array.values().iter().map(|value| *value / divisor);
- let buffer = unsafe { Buffer::from_trusted_len_iter(values) };
-
- let data = ArrayData::new(
- T::DATA_TYPE,
- array.len(),
- None,
- array.data_ref().null_buffer().cloned(),
- 0,
- vec![buffer],
- vec![],
- );
- Ok(PrimitiveArray::<T>::from(data))
-}
-
-/// SIMD vectorized version of `math_op` above.
-#[cfg(feature = "simd")]
-fn simd_math_op<T, SIMD_OP, SCALAR_OP>(
- left: &PrimitiveArray<T>,
- right: &PrimitiveArray<T>,
- simd_op: SIMD_OP,
- scalar_op: SCALAR_OP,
-) -> Result<PrimitiveArray<T>>
-where
- T: ArrowNumericType,
- SIMD_OP: Fn(T::Simd, T::Simd) -> T::Simd,
- SCALAR_OP: Fn(T::Native, T::Native) -> T::Native,
-{
- if left.len() != right.len() {
- return Err(ArrowError::ComputeError(
- "Cannot perform math operation on arrays of different length".to_string(),
- ));
- }
-
- let null_bit_buffer =
- combine_option_bitmap(left.data_ref(), right.data_ref(), left.len())?;
-
- let lanes = T::lanes();
- let buffer_size = left.len() * std::mem::size_of::<T::Native>();
- let mut result = MutableBuffer::new(buffer_size).with_bitset(buffer_size, false);
-
- let mut result_chunks = result.typed_data_mut().chunks_exact_mut(lanes);
- let mut left_chunks = left.values().chunks_exact(lanes);
- let mut right_chunks = right.values().chunks_exact(lanes);
-
- result_chunks
- .borrow_mut()
- .zip(left_chunks.borrow_mut().zip(right_chunks.borrow_mut()))
- .for_each(|(result_slice, (left_slice, right_slice))| {
- let simd_left = T::load(left_slice);
- let simd_right = T::load(right_slice);
- let simd_result = T::bin_op(simd_left, simd_right, &simd_op);
- T::write(simd_result, result_slice);
- });
-
- let result_remainder = result_chunks.into_remainder();
- let left_remainder = left_chunks.remainder();
- let right_remainder = right_chunks.remainder();
-
- result_remainder
- .iter_mut()
- .zip(left_remainder.iter().zip(right_remainder.iter()))
- .for_each(|(scalar_result, (scalar_left, scalar_right))| {
- *scalar_result = scalar_op(*scalar_left, *scalar_right);
- });
-
- let data = ArrayData::new(
- T::DATA_TYPE,
- left.len(),
- None,
- null_bit_buffer,
- 0,
- vec![result.into()],
- vec![],
- );
- Ok(PrimitiveArray::<T>::from(data))
-}
-
-/// SIMD vectorized implementation of `left % right`.
-/// If any of the lanes marked as valid in `valid_mask` are `0` then an `ArrowError::DivideByZero`
-/// is returned. The contents of no-valid lanes are undefined.
-#[cfg(feature = "simd")]
-#[inline]
-fn simd_checked_modulus<T: ArrowNumericType>(
- valid_mask: Option<u64>,
- left: T::Simd,
- right: T::Simd,
-) -> Result<T::Simd>
-where
- T::Native: One + Zero,
-{
- let zero = T::init(T::Native::zero());
- let one = T::init(T::Native::one());
-
- let right_no_invalid_zeros = match valid_mask {
- Some(mask) => {
- let simd_mask = T::mask_from_u64(mask);
- // select `1` for invalid lanes, which will be a no-op during division later
- T::mask_select(simd_mask, right, one)
- }
- None => right,
- };
-
- let zero_mask = T::eq(right_no_invalid_zeros, zero);
-
- if T::mask_any(zero_mask) {
- Err(ArrowError::DivideByZero)
- } else {
- Ok(T::bin_op(left, right_no_invalid_zeros, |a, b| a % b))
- }
-}
-
-/// SIMD vectorized implementation of `left / right`.
-/// If any of the lanes marked as valid in `valid_mask` are `0` then an `ArrowError::DivideByZero`
-/// is returned. The contents of no-valid lanes are undefined.
-#[cfg(feature = "simd")]
-#[inline]
-fn simd_checked_divide<T: ArrowNumericType>(
- valid_mask: Option<u64>,
- left: T::Simd,
- right: T::Simd,
-) -> Result<T::Simd>
-where
- T::Native: One + Zero,
-{
- let zero = T::init(T::Native::zero());
- let one = T::init(T::Native::one());
-
- let right_no_invalid_zeros = match valid_mask {
- Some(mask) => {
- let simd_mask = T::mask_from_u64(mask);
- // select `1` for invalid lanes, which will be a no-op during division later
- T::mask_select(simd_mask, right, one)
- }
- None => right,
- };
-
- let zero_mask = T::eq(right_no_invalid_zeros, zero);
-
- if T::mask_any(zero_mask) {
- Err(ArrowError::DivideByZero)
- } else {
- Ok(T::bin_op(left, right_no_invalid_zeros, |a, b| a / b))
- }
-}
-
-/// Scalar implementation of `left % right` for the remainder elements after complete chunks have been processed using SIMD.
-/// If any of the values marked as valid in `valid_mask` are `0` then an `ArrowError::DivideByZero` is returned.
-#[cfg(feature = "simd")]
-#[inline]
-fn simd_checked_modulus_remainder<T: ArrowNumericType>(
- valid_mask: Option<u64>,
- left_chunks: ChunksExact<T::Native>,
- right_chunks: ChunksExact<T::Native>,
- result_chunks: ChunksExactMut<T::Native>,
-) -> Result<()>
-where
- T::Native: Zero + Rem<Output = T::Native>,
-{
- let result_remainder = result_chunks.into_remainder();
- let left_remainder = left_chunks.remainder();
- let right_remainder = right_chunks.remainder();
-
- result_remainder
- .iter_mut()
- .zip(left_remainder.iter().zip(right_remainder.iter()))
- .enumerate()
- .try_for_each(|(i, (result_scalar, (left_scalar, right_scalar)))| {
- if valid_mask.map(|mask| mask & (1 << i) != 0).unwrap_or(true) {
- if *right_scalar == T::Native::zero() {
- return Err(ArrowError::DivideByZero);
- }
- *result_scalar = *left_scalar % *right_scalar;
- }
- Ok(())
- })?;
-
- Ok(())
-}
-
-/// Scalar implementation of `left / right` for the remainder elements after complete chunks have been processed using SIMD.
-/// If any of the values marked as valid in `valid_mask` are `0` then an `ArrowError::DivideByZero` is returned.
-#[cfg(feature = "simd")]
-#[inline]
-fn simd_checked_divide_remainder<T: ArrowNumericType>(
- valid_mask: Option<u64>,
- left_chunks: ChunksExact<T::Native>,
- right_chunks: ChunksExact<T::Native>,
- result_chunks: ChunksExactMut<T::Native>,
-) -> Result<()>
-where
- T::Native: Zero + Div<Output = T::Native>,
-{
- let result_remainder = result_chunks.into_remainder();
- let left_remainder = left_chunks.remainder();
- let right_remainder = right_chunks.remainder();
-
- result_remainder
- .iter_mut()
- .zip(left_remainder.iter().zip(right_remainder.iter()))
- .enumerate()
- .try_for_each(|(i, (result_scalar, (left_scalar, right_scalar)))| {
- if valid_mask.map(|mask| mask & (1 << i) != 0).unwrap_or(true) {
- if *right_scalar == T::Native::zero() {
- return Err(ArrowError::DivideByZero);
- }
- *result_scalar = *left_scalar / *right_scalar;
- }
- Ok(())
- })?;
-
- Ok(())
-}
-
-/// Scalar-modulo version of `simd_checked_modulus_remainder`.
-#[cfg(feature = "simd")]
-#[inline]
-fn simd_checked_modulus_scalar_remainder<T: ArrowNumericType>(
- array_chunks: ChunksExact<T::Native>,
- modulo: T::Native,
- result_chunks: ChunksExactMut<T::Native>,
-) -> Result<()>
-where
- T::Native: Zero + Rem<Output = T::Native>,
-{
- if modulo.is_zero() {
- return Err(ArrowError::DivideByZero);
- }
-
- let result_remainder = result_chunks.into_remainder();
- let array_remainder = array_chunks.remainder();
-
- result_remainder
- .iter_mut()
- .zip(array_remainder.iter())
- .for_each(|(result_scalar, array_scalar)| {
- *result_scalar = *array_scalar % modulo;
- });
-
- Ok(())
-}
-
-/// Scalar-divisor version of `simd_checked_divide_remainder`.
-#[cfg(feature = "simd")]
-#[inline]
-fn simd_checked_divide_scalar_remainder<T: ArrowNumericType>(
- array_chunks: ChunksExact<T::Native>,
- divisor: T::Native,
- result_chunks: ChunksExactMut<T::Native>,
-) -> Result<()>
-where
- T::Native: Zero + Div<Output = T::Native>,
-{
- if divisor.is_zero() {
- return Err(ArrowError::DivideByZero);
- }
-
- let result_remainder = result_chunks.into_remainder();
- let array_remainder = array_chunks.remainder();
-
- result_remainder
- .iter_mut()
- .zip(array_remainder.iter())
- .for_each(|(result_scalar, array_scalar)| {
- *result_scalar = *array_scalar / divisor;
- });
-
- Ok(())
-}
-
-/// SIMD vectorized version of `modulus`.
-///
-/// The modulus kernels need their own implementation as there is a need to handle situations
-/// where a modulus by `0` occurs. This is complicated by `NULL` slots and padding.
-#[cfg(feature = "simd")]
-fn simd_modulus<T>(
- left: &PrimitiveArray<T>,
- right: &PrimitiveArray<T>,
-) -> Result<PrimitiveArray<T>>
-where
- T: ArrowNumericType,
- T::Native: One + Zero + Rem<Output = T::Native>,
-{
- if left.len() != right.len() {
- return Err(ArrowError::ComputeError(
- "Cannot perform math operation on arrays of different length".to_string(),
- ));
- }
-
- // Create the combined `Bitmap`
- let null_bit_buffer =
- combine_option_bitmap(left.data_ref(), right.data_ref(), left.len())?;
-
- let lanes = T::lanes();
- let buffer_size = left.len() * std::mem::size_of::<T::Native>();
- let mut result = MutableBuffer::new(buffer_size).with_bitset(buffer_size, false);
-
- match &null_bit_buffer {
- Some(b) => {
- // combine_option_bitmap returns a slice or new buffer starting at 0
- let valid_chunks = b.bit_chunks(0, left.len());
-
- // process data in chunks of 64 elements since we also get 64 bits of validity information at a time
- let mut result_chunks = result.typed_data_mut().chunks_exact_mut(64);
- let mut left_chunks = left.values().chunks_exact(64);
- let mut right_chunks = right.values().chunks_exact(64);
-
- valid_chunks
- .iter()
- .zip(
- result_chunks
- .borrow_mut()
- .zip(left_chunks.borrow_mut().zip(right_chunks.borrow_mut())),
- )
- .try_for_each(
- |(mut mask, (result_slice, (left_slice, right_slice)))| {
- // split chunks further into slices corresponding to the vector length
- // the compiler is able to unroll this inner loop and remove bounds checks
- // since the outer chunk size (64) is always a multiple of the number of lanes
- result_slice
- .chunks_exact_mut(lanes)
- .zip(left_slice.chunks_exact(lanes).zip(right_slice.chunks_exact(lanes)))
- .try_for_each(|(result_slice, (left_slice, right_slice))| -> Result<()> {
- let simd_left = T::load(left_slice);
- let simd_right = T::load(right_slice);
-
- let simd_result = simd_checked_modulus::<T>(Some(mask), simd_left, simd_right)?;
-
- T::write(simd_result, result_slice);
-
- // skip the shift and avoid overflow for u8 type, which uses 64 lanes.
- mask >>= T::lanes() % 64;
-
- Ok(())
- })
- },
- )?;
-
- let valid_remainder = valid_chunks.remainder_bits();
-
- simd_checked_modulus_remainder::<T>(
- Some(valid_remainder),
- left_chunks,
- right_chunks,
- result_chunks,
- )?;
- }
- None => {
- let mut result_chunks = result.typed_data_mut().chunks_exact_mut(lanes);
- let mut left_chunks = left.values().chunks_exact(lanes);
- let mut right_chunks = right.values().chunks_exact(lanes);
-
- result_chunks
- .borrow_mut()
- .zip(left_chunks.borrow_mut().zip(right_chunks.borrow_mut()))
- .try_for_each(
- |(result_slice, (left_slice, right_slice))| -> Result<()> {
- let simd_left = T::load(left_slice);
- let simd_right = T::load(right_slice);
-
- let simd_result =
- simd_checked_modulus::<T>(None, simd_left, simd_right)?;
-
- T::write(simd_result, result_slice);
-
- Ok(())
- },
- )?;
-
- simd_checked_modulus_remainder::<T>(
- None,
- left_chunks,
- right_chunks,
- result_chunks,
- )?;
- }
- }
-
- let data = ArrayData::new(
- T::DATA_TYPE,
- left.len(),
- None,
- null_bit_buffer,
- 0,
- vec![result.into()],
- vec![],
- );
- Ok(PrimitiveArray::<T>::from(data))
-}
-
-/// SIMD vectorized version of `divide`.
-///
-/// The divide kernels need their own implementation as there is a need to handle situations
-/// where a divide by `0` occurs. This is complicated by `NULL` slots and padding.
-#[cfg(feature = "simd")]
-fn simd_divide<T>(
- left: &PrimitiveArray<T>,
- right: &PrimitiveArray<T>,
-) -> Result<PrimitiveArray<T>>
-where
- T: ArrowNumericType,
- T::Native: One + Zero + Div<Output = T::Native>,
-{
- if left.len() != right.len() {
- return Err(ArrowError::ComputeError(
- "Cannot perform math operation on arrays of different length".to_string(),
- ));
- }
-
- // Create the combined `Bitmap`
- let null_bit_buffer =
- combine_option_bitmap(left.data_ref(), right.data_ref(), left.len())?;
-
- let lanes = T::lanes();
- let buffer_size = left.len() * std::mem::size_of::<T::Native>();
- let mut result = MutableBuffer::new(buffer_size).with_bitset(buffer_size, false);
-
- match &null_bit_buffer {
- Some(b) => {
- // combine_option_bitmap returns a slice or new buffer starting at 0
- let valid_chunks = b.bit_chunks(0, left.len());
-
- // process data in chunks of 64 elements since we also get 64 bits of validity information at a time
- let mut result_chunks = result.typed_data_mut().chunks_exact_mut(64);
- let mut left_chunks = left.values().chunks_exact(64);
- let mut right_chunks = right.values().chunks_exact(64);
-
- valid_chunks
- .iter()
- .zip(
- result_chunks
- .borrow_mut()
- .zip(left_chunks.borrow_mut().zip(right_chunks.borrow_mut())),
- )
- .try_for_each(
- |(mut mask, (result_slice, (left_slice, right_slice)))| {
- // split chunks further into slices corresponding to the vector length
- // the compiler is able to unroll this inner loop and remove bounds checks
- // since the outer chunk size (64) is always a multiple of the number of lanes
- result_slice
- .chunks_exact_mut(lanes)
- .zip(left_slice.chunks_exact(lanes).zip(right_slice.chunks_exact(lanes)))
- .try_for_each(|(result_slice, (left_slice, right_slice))| -> Result<()> {
- let simd_left = T::load(left_slice);
- let simd_right = T::load(right_slice);
-
- let simd_result = simd_checked_divide::<T>(Some(mask), simd_left, simd_right)?;
-
- T::write(simd_result, result_slice);
-
- // skip the shift and avoid overflow for u8 type, which uses 64 lanes.
- mask >>= T::lanes() % 64;
-
- Ok(())
- })
- },
- )?;
-
- let valid_remainder = valid_chunks.remainder_bits();
-
- simd_checked_divide_remainder::<T>(
- Some(valid_remainder),
- left_chunks,
- right_chunks,
- result_chunks,
- )?;
- }
- None => {
- let mut result_chunks = result.typed_data_mut().chunks_exact_mut(lanes);
- let mut left_chunks = left.values().chunks_exact(lanes);
- let mut right_chunks = right.values().chunks_exact(lanes);
-
- result_chunks
- .borrow_mut()
- .zip(left_chunks.borrow_mut().zip(right_chunks.borrow_mut()))
- .try_for_each(
- |(result_slice, (left_slice, right_slice))| -> Result<()> {
- let simd_left = T::load(left_slice);
- let simd_right = T::load(right_slice);
-
- let simd_result =
- simd_checked_divide::<T>(None, simd_left, simd_right)?;
-
- T::write(simd_result, result_slice);
-
- Ok(())
- },
- )?;
-
- simd_checked_divide_remainder::<T>(
- None,
- left_chunks,
- right_chunks,
- result_chunks,
- )?;
- }
- }
-
- let data = ArrayData::new(
- T::DATA_TYPE,
- left.len(),
- None,
- null_bit_buffer,
- 0,
- vec![result.into()],
- vec![],
- );
- Ok(PrimitiveArray::<T>::from(data))
-}
-
-/// SIMD vectorized version of `modulus_scalar`.
-#[cfg(feature = "simd")]
-fn simd_modulus_scalar<T>(
- array: &PrimitiveArray<T>,
- modulo: T::Native,
-) -> Result<PrimitiveArray<T>>
-where
- T: ArrowNumericType,
- T::Native: One + Zero + Rem<Output = T::Native>,
-{
- if modulo.is_zero() {
- return Err(ArrowError::DivideByZero);
- }
-
- let lanes = T::lanes();
- let buffer_size = array.len() * std::mem::size_of::<T::Native>();
- let mut result = MutableBuffer::new(buffer_size).with_bitset(buffer_size, false);
-
- let mut result_chunks = result.typed_data_mut().chunks_exact_mut(lanes);
- let mut array_chunks = array.values().chunks_exact(lanes);
-
- result_chunks
- .borrow_mut()
- .zip(array_chunks.borrow_mut())
- .for_each(|(result_slice, array_slice)| {
- let simd_left = T::load(array_slice);
- let simd_right = T::init(modulo);
-
- let simd_result = T::bin_op(simd_left, simd_right, |a, b| a % b);
- T::write(simd_result, result_slice);
- });
-
- simd_checked_modulus_scalar_remainder::<T>(array_chunks, modulo, result_chunks)?;
-
- let data = ArrayData::new(
- T::DATA_TYPE,
- array.len(),
- None,
- array.data_ref().null_buffer().cloned(),
- 0,
- vec![result.into()],
- vec![],
- );
- Ok(PrimitiveArray::<T>::from(data))
-}
-
-/// SIMD vectorized version of `divide_scalar`.
-#[cfg(feature = "simd")]
-fn simd_divide_scalar<T>(
- array: &PrimitiveArray<T>,
- divisor: T::Native,
-) -> Result<PrimitiveArray<T>>
-where
- T: ArrowNumericType,
- T::Native: One + Zero + Div<Output = T::Native>,
-{
- if divisor.is_zero() {
- return Err(ArrowError::DivideByZero);
- }
-
- let lanes = T::lanes();
- let buffer_size = array.len() * std::mem::size_of::<T::Native>();
- let mut result = MutableBuffer::new(buffer_size).with_bitset(buffer_size, false);
-
- let mut result_chunks = result.typed_data_mut().chunks_exact_mut(lanes);
- let mut array_chunks = array.values().chunks_exact(lanes);
-
- result_chunks
- .borrow_mut()
- .zip(array_chunks.borrow_mut())
- .for_each(|(result_slice, array_slice)| {
- let simd_left = T::load(array_slice);
- let simd_right = T::init(divisor);
-
- let simd_result = T::bin_op(simd_left, simd_right, |a, b| a / b);
- T::write(simd_result, result_slice);
- });
-
- simd_checked_divide_scalar_remainder::<T>(array_chunks, divisor, result_chunks)?;
-
- let data = ArrayData::new(
- T::DATA_TYPE,
- array.len(),
- None,
- array.data_ref().null_buffer().cloned(),
- 0,
- vec![result.into()],
- vec![],
- );
- Ok(PrimitiveArray::<T>::from(data))
-}
-
-/// Perform `left + right` operation on two arrays. If either left or right value is null
-/// then the result is also null.
-pub fn add<T>(
- left: &PrimitiveArray<T>,
- right: &PrimitiveArray<T>,
-) -> Result<PrimitiveArray<T>>
-where
- T: ArrowNumericType,
- T::Native: Add<Output = T::Native>
- + Sub<Output = T::Native>
- + Mul<Output = T::Native>
- + Div<Output = T::Native>
- + Zero,
-{
- #[cfg(feature = "simd")]
- return simd_math_op(&left, &right, |a, b| a + b, |a, b| a + b);
- #[cfg(not(feature = "simd"))]
- return math_op(left, right, |a, b| a + b);
-}
-
-/// Perform `left - right` operation on two arrays. If either left or right value is null
-/// then the result is also null.
-pub fn subtract<T>(
- left: &PrimitiveArray<T>,
- right: &PrimitiveArray<T>,
-) -> Result<PrimitiveArray<T>>
-where
- T: datatypes::ArrowNumericType,
- T::Native: Add<Output = T::Native>
- + Sub<Output = T::Native>
- + Mul<Output = T::Native>
- + Div<Output = T::Native>
- + Zero,
-{
- #[cfg(feature = "simd")]
- return simd_math_op(&left, &right, |a, b| a - b, |a, b| a - b);
- #[cfg(not(feature = "simd"))]
- return math_op(left, right, |a, b| a - b);
-}
-
-/// Perform `-` operation on an array. If value is null then the result is also null.
-pub fn negate<T>(array: &PrimitiveArray<T>) -> Result<PrimitiveArray<T>>
-where
- T: datatypes::ArrowSignedNumericType,
- T::Native: Neg<Output = T::Native>,
-{
- #[cfg(feature = "simd")]
- return simd_signed_unary_math_op(array, |x| -x, |x| -x);
- #[cfg(not(feature = "simd"))]
- return Ok(unary(array, |x| -x));
-}
-
-/// Raise array with floating point values to the power of a scalar.
-pub fn powf_scalar<T>(
- array: &PrimitiveArray<T>,
- raise: T::Native,
-) -> Result<PrimitiveArray<T>>
-where
- T: datatypes::ArrowFloatNumericType,
- T::Native: Pow<T::Native, Output = T::Native>,
-{
- #[cfg(feature = "simd")]
- {
- let raise_vector = T::init(raise);
- return simd_float_unary_math_op(
- array,
- |x| T::pow(x, raise_vector),
- |x| x.pow(raise),
- );
- }
- #[cfg(not(feature = "simd"))]
- return Ok(unary(array, |x| x.pow(raise)));
-}
-
-/// Perform `left * right` operation on two arrays. If either left or right value is null
-/// then the result is also null.
-pub fn multiply<T>(
- left: &PrimitiveArray<T>,
- right: &PrimitiveArray<T>,
-) -> Result<PrimitiveArray<T>>
-where
- T: datatypes::ArrowNumericType,
- T::Native: Add<Output = T::Native>
- + Sub<Output = T::Native>
- + Mul<Output = T::Native>
- + Div<Output = T::Native>
- + Rem<Output = T::Native>
- + Zero,
-{
- #[cfg(feature = "simd")]
- return simd_math_op(&left, &right, |a, b| a * b, |a, b| a * b);
- #[cfg(not(feature = "simd"))]
- return math_op(left, right, |a, b| a * b);
-}
-
-/// Perform `left % right` operation on two arrays. If either left or right value is null
-/// then the result is also null. If any right hand value is zero then the result of this
-/// operation will be `Err(ArrowError::DivideByZero)`.
-pub fn modulus<T>(
- left: &PrimitiveArray<T>,
- right: &PrimitiveArray<T>,
-) -> Result<PrimitiveArray<T>>
-where
- T: datatypes::ArrowNumericType,
- T::Native: Add<Output = T::Native>
- + Sub<Output = T::Native>
- + Mul<Output = T::Native>
- + Div<Output = T::Native>
- + Rem<Output = T::Native>
- + Zero
- + One,
-{
- #[cfg(feature = "simd")]
- return simd_modulus(&left, &right);
- #[cfg(not(feature = "simd"))]
- return math_modulus(&left, &right);
-}
-
-/// Perform `left / right` operation on two arrays. If either left or right value is null
-/// then the result is also null. If any right hand value is zero then the result of this
-/// operation will be `Err(ArrowError::DivideByZero)`.
-pub fn divide<T>(
- left: &PrimitiveArray<T>,
- right: &PrimitiveArray<T>,
-) -> Result<PrimitiveArray<T>>
-where
- T: datatypes::ArrowNumericType,
- T::Native: Add<Output = T::Native>
- + Sub<Output = T::Native>
- + Mul<Output = T::Native>
- + Div<Output = T::Native>
- + Rem<Output = T::Native>
- + Zero
- + One,
-{
- #[cfg(feature = "simd")]
- return simd_divide(&left, &right);
- #[cfg(not(feature = "simd"))]
- return math_divide(&left, &right);
-}
-
-/// Modulus every value in an array by a scalar. If any value in the array is null then the
-/// result is also null. If the scalar is zero then the result of this operation will be
-/// `Err(ArrowError::DivideByZero)`.
-pub fn modulus_scalar<T>(
- array: &PrimitiveArray<T>,
- modulo: T::Native,
-) -> Result<PrimitiveArray<T>>
-where
- T: datatypes::ArrowNumericType,
- T::Native: Add<Output = T::Native>
- + Sub<Output = T::Native>
- + Mul<Output = T::Native>
- + Div<Output = T::Native>
- + Rem<Output = T::Native>
- + Zero
- + One,
-{
- #[cfg(feature = "simd")]
- return simd_modulus_scalar(&array, modulo);
- #[cfg(not(feature = "simd"))]
- return math_modulus_scalar(&array, modulo);
-}
-
-/// Divide every value in an array by a scalar. If any value in the array is null then the
-/// result is also null. If the scalar is zero then the result of this operation will be
-/// `Err(ArrowError::DivideByZero)`.
-pub fn divide_scalar<T>(
- array: &PrimitiveArray<T>,
- divisor: T::Native,
-) -> Result<PrimitiveArray<T>>
-where
- T: datatypes::ArrowNumericType,
- T::Native: Add<Output = T::Native>
- + Sub<Output = T::Native>
- + Mul<Output = T::Native>
- + Div<Output = T::Native>
- + Rem<Output = T::Native>
- + Zero
- + One,
-{
- #[cfg(feature = "simd")]
- return simd_divide_scalar(&array, divisor);
- #[cfg(not(feature = "simd"))]
- return math_divide_scalar(&array, divisor);
-}
-
-#[cfg(test)]
-mod tests {
- use super::*;
- use crate::array::Int32Array;
-
- #[test]
- fn test_primitive_array_add() {
- let a = Int32Array::from(vec![5, 6, 7, 8, 9]);
- let b = Int32Array::from(vec![6, 7, 8, 9, 8]);
- let c = add(&a, &b).unwrap();
- assert_eq!(11, c.value(0));
- assert_eq!(13, c.value(1));
- assert_eq!(15, c.value(2));
- assert_eq!(17, c.value(3));
- assert_eq!(17, c.value(4));
- }
-
- #[test]
- fn test_primitive_array_add_sliced() {
- let a = Int32Array::from(vec![0, 0, 0, 5, 6, 7, 8, 9, 0]);
- let b = Int32Array::from(vec![0, 0, 0, 6, 7, 8, 9, 8, 0]);
- let a = a.slice(3, 5);
- let b = b.slice(3, 5);
- let a = a.as_any().downcast_ref::<Int32Array>().unwrap();
- let b = b.as_any().downcast_ref::<Int32Array>().unwrap();
-
- assert_eq!(5, a.value(0));
- assert_eq!(6, b.value(0));
-
- let c = add(&a, &b).unwrap();
- assert_eq!(5, c.len());
- assert_eq!(11, c.value(0));
- assert_eq!(13, c.value(1));
- assert_eq!(15, c.value(2));
- assert_eq!(17, c.value(3));
- assert_eq!(17, c.value(4));
- }
-
- #[test]
- fn test_primitive_array_add_mismatched_length() {
- let a = Int32Array::from(vec![5, 6, 7, 8, 9]);
- let b = Int32Array::from(vec![6, 7, 8]);
- let e = add(&a, &b)
- .err()
- .expect("should have failed due to different lengths");
- assert_eq!(
- "ComputeError(\"Cannot perform math operation on arrays of different length\")",
- format!("{:?}", e)
- );
- }
-
- #[test]
- fn test_primitive_array_subtract() {
- let a = Int32Array::from(vec![1, 2, 3, 4, 5]);
- let b = Int32Array::from(vec![5, 4, 3, 2, 1]);
- let c = subtract(&a, &b).unwrap();
- assert_eq!(-4, c.value(0));
- assert_eq!(-2, c.value(1));
- assert_eq!(0, c.value(2));
- assert_eq!(2, c.value(3));
- assert_eq!(4, c.value(4));
- }
-
- #[test]
- fn test_primitive_array_multiply() {
- let a = Int32Array::from(vec![5, 6, 7, 8, 9]);
- let b = Int32Array::from(vec![6, 7, 8, 9, 8]);
- let c = multiply(&a, &b).unwrap();
- assert_eq!(30, c.value(0));
- assert_eq!(42, c.value(1));
- assert_eq!(56, c.value(2));
- assert_eq!(72, c.value(3));
- assert_eq!(72, c.value(4));
- }
-
- #[test]
- fn test_primitive_array_divide() {
- let a = Int32Array::from(vec![15, 15, 8, 1, 9]);
- let b = Int32Array::from(vec![5, 6, 8, 9, 1]);
- let c = divide(&a, &b).unwrap();
- assert_eq!(3, c.value(0));
- assert_eq!(2, c.value(1));
- assert_eq!(1, c.value(2));
- assert_eq!(0, c.value(3));
- assert_eq!(9, c.value(4));
- }
-
- #[test]
- fn test_primitive_array_modulus() {
- let a = Int32Array::from(vec![15, 15, 8, 1, 9]);
- let b = Int32Array::from(vec![5, 6, 8, 9, 1]);
- let c = modulus(&a, &b).unwrap();
- assert_eq!(0, c.value(0));
- assert_eq!(3, c.value(1));
- assert_eq!(0, c.value(2));
- assert_eq!(1, c.value(3));
- assert_eq!(0, c.value(4));
- }
-
- #[test]
- fn test_primitive_array_divide_scalar() {
- let a = Int32Array::from(vec![15, 14, 9, 8, 1]);
- let b = 3;
- let c = divide_scalar(&a, b).unwrap();
- let expected = Int32Array::from(vec![5, 4, 3, 2, 0]);
- assert_eq!(c, expected);
- }
-
- #[test]
- fn test_primitive_array_modulus_scalar() {
- let a = Int32Array::from(vec![15, 14, 9, 8, 1]);
- let b = 3;
- let c = modulus_scalar(&a, b).unwrap();
- let expected = Int32Array::from(vec![0, 2, 0, 2, 1]);
- assert_eq!(c, expected);
- }
-
- #[test]
- fn test_primitive_array_divide_sliced() {
- let a = Int32Array::from(vec![0, 0, 0, 15, 15, 8, 1, 9, 0]);
- let b = Int32Array::from(vec![0, 0, 0, 5, 6, 8, 9, 1, 0]);
- let a = a.slice(3, 5);
- let b = b.slice(3, 5);
- let a = a.as_any().downcast_ref::<Int32Array>().unwrap();
- let b = b.as_any().downcast_ref::<Int32Array>().unwrap();
-
- let c = divide(&a, &b).unwrap();
- assert_eq!(5, c.len());
- assert_eq!(3, c.value(0));
- assert_eq!(2, c.value(1));
- assert_eq!(1, c.value(2));
- assert_eq!(0, c.value(3));
- assert_eq!(9, c.value(4));
- }
-
- #[test]
- fn test_primitive_array_modulus_sliced() {
- let a = Int32Array::from(vec![0, 0, 0, 15, 15, 8, 1, 9, 0]);
- let b = Int32Array::from(vec![0, 0, 0, 5, 6, 8, 9, 1, 0]);
- let a = a.slice(3, 5);
- let b = b.slice(3, 5);
- let a = a.as_any().downcast_ref::<Int32Array>().unwrap();
- let b = b.as_any().downcast_ref::<Int32Array>().unwrap();
-
- let c = modulus(&a, &b).unwrap();
- assert_eq!(5, c.len());
- assert_eq!(0, c.value(0));
- assert_eq!(3, c.value(1));
- assert_eq!(0, c.value(2));
- assert_eq!(1, c.value(3));
- assert_eq!(0, c.value(4));
- }
-
- #[test]
- fn test_primitive_array_divide_with_nulls() {
- let a = Int32Array::from(vec![Some(15), None, Some(8), Some(1), Some(9), None]);
- let b = Int32Array::from(vec![Some(5), Some(6), Some(8), Some(9), None, None]);
- let c = divide(&a, &b).unwrap();
- assert_eq!(3, c.value(0));
- assert_eq!(true, c.is_null(1));
- assert_eq!(1, c.value(2));
- assert_eq!(0, c.value(3));
- assert_eq!(true, c.is_null(4));
- assert_eq!(true, c.is_null(5));
- }
-
- #[test]
- fn test_primitive_array_modulus_with_nulls() {
- let a = Int32Array::from(vec![Some(15), None, Some(8), Some(1), Some(9), None]);
- let b = Int32Array::from(vec![Some(5), Some(6), Some(8), Some(9), None, None]);
- let c = modulus(&a, &b).unwrap();
- assert_eq!(0, c.value(0));
- assert_eq!(true, c.is_null(1));
- assert_eq!(0, c.value(2));
- assert_eq!(1, c.value(3));
- assert_eq!(true, c.is_null(4));
- assert_eq!(true, c.is_null(5));
- }
-
- #[test]
- fn test_primitive_array_divide_scalar_with_nulls() {
- let a = Int32Array::from(vec![Some(15), None, Some(8), Some(1), Some(9), None]);
- let b = 3;
- let c = divide_scalar(&a, b).unwrap();
- let expected =
- Int32Array::from(vec![Some(5), None, Some(2), Some(0), Some(3), None]);
- assert_eq!(c, expected);
- }
-
- #[test]
- fn test_primitive_array_modulus_scalar_with_nulls() {
- let a = Int32Array::from(vec![Some(15), None, Some(8), Some(1), Some(9), None]);
- let b = 3;
- let c = modulus_scalar(&a, b).unwrap();
- let expected =
- Int32Array::from(vec![Some(0), None, Some(2), Some(1), Some(0), None]);
- assert_eq!(c, expected);
- }
-
- #[test]
- fn test_primitive_array_divide_with_nulls_sliced() {
- let a = Int32Array::from(vec![
- None,
- None,
- None,
- None,
- None,
- None,
- None,
- None,
- Some(15),
- None,
- Some(8),
- Some(1),
- Some(9),
- None,
- None,
- ]);
- let b = Int32Array::from(vec![
- None,
- None,
- None,
- None,
- None,
- None,
- None,
- None,
- Some(5),
- Some(6),
- Some(8),
- Some(9),
- None,
- None,
- None,
- ]);
-
- let a = a.slice(8, 6);
- let a = a.as_any().downcast_ref::<Int32Array>().unwrap();
-
- let b = b.slice(8, 6);
- let b = b.as_any().downcast_ref::<Int32Array>().unwrap();
-
- let c = divide(&a, &b).unwrap();
- assert_eq!(6, c.len());
- assert_eq!(3, c.value(0));
- assert_eq!(true, c.is_null(1));
- assert_eq!(1, c.value(2));
- assert_eq!(0, c.value(3));
- assert_eq!(true, c.is_null(4));
- assert_eq!(true, c.is_null(5));
- }
-
- #[test]
- fn test_primitive_array_modulus_with_nulls_sliced() {
- let a = Int32Array::from(vec![
- None,
- None,
- None,
- None,
- None,
- None,
- None,
- None,
- Some(15),
- None,
- Some(8),
- Some(1),
- Some(9),
- None,
- None,
- ]);
- let b = Int32Array::from(vec![
- None,
- None,
- None,
- None,
- None,
- None,
- None,
- None,
- Some(5),
- Some(6),
- Some(8),
- Some(9),
- None,
- None,
- None,
- ]);
-
- let a = a.slice(8, 6);
- let a = a.as_any().downcast_ref::<Int32Array>().unwrap();
-
- let b = b.slice(8, 6);
- let b = b.as_any().downcast_ref::<Int32Array>().unwrap();
-
- let c = modulus(&a, &b).unwrap();
- assert_eq!(6, c.len());
- assert_eq!(0, c.value(0));
- assert_eq!(true, c.is_null(1));
- assert_eq!(0, c.value(2));
- assert_eq!(1, c.value(3));
- assert_eq!(true, c.is_null(4));
- assert_eq!(true, c.is_null(5));
- }
-
- #[test]
- #[should_panic(expected = "DivideByZero")]
- fn test_primitive_array_divide_by_zero() {
- let a = Int32Array::from(vec![15]);
- let b = Int32Array::from(vec![0]);
- divide(&a, &b).unwrap();
- }
-
- #[test]
- #[should_panic(expected = "DivideByZero")]
- fn test_primitive_array_modulus_by_zero() {
- let a = Int32Array::from(vec![15]);
- let b = Int32Array::from(vec![0]);
- modulus(&a, &b).unwrap();
- }
-
- #[test]
- fn test_primitive_array_divide_f64() {
- let a = Float64Array::from(vec![15.0, 15.0, 8.0]);
- let b = Float64Array::from(vec![5.0, 6.0, 8.0]);
- let c = divide(&a, &b).unwrap();
- assert!(3.0 - c.value(0) < f64::EPSILON);
- assert!(2.5 - c.value(1) < f64::EPSILON);
- assert!(1.0 - c.value(2) < f64::EPSILON);
- }
-
- #[test]
- fn test_primitive_array_add_with_nulls() {
- let a = Int32Array::from(vec![Some(5), None, Some(7), None]);
- let b = Int32Array::from(vec![None, None, Some(6), Some(7)]);
- let c = add(&a, &b).unwrap();
- assert_eq!(true, c.is_null(0));
- assert_eq!(true, c.is_null(1));
- assert_eq!(false, c.is_null(2));
- assert_eq!(true, c.is_null(3));
- assert_eq!(13, c.value(2));
- }
-
- #[test]
- fn test_primitive_array_negate() {
- let a: Int64Array = (0..100).into_iter().map(Some).collect();
- let actual = negate(&a).unwrap();
- let expected: Int64Array = (0..100).into_iter().map(|i| Some(-i)).collect();
- assert_eq!(expected, actual);
- }
-
- #[test]
- fn test_arithmetic_kernel_should_not_rely_on_padding() {
- let a: UInt8Array = (0..128_u8).into_iter().map(Some).collect();
- let a = a.slice(63, 65);
- let a = a.as_any().downcast_ref::<UInt8Array>().unwrap();
-
- let b: UInt8Array = (0..128_u8).into_iter().map(Some).collect();
- let b = b.slice(63, 65);
- let b = b.as_any().downcast_ref::<UInt8Array>().unwrap();
-
- let actual = add(&a, &b).unwrap();
- let actual: Vec<Option<u8>> = actual.iter().collect();
- let expected: Vec<Option<u8>> = (63..63_u8 + 65_u8)
- .into_iter()
- .map(|i| Some(i + i))
- .collect();
- assert_eq!(expected, actual);
- }
-
- #[test]
- fn test_primitive_array_raise_power_scalar() {
- let a = Float64Array::from(vec![1.0, 2.0, 3.0]);
- let actual = powf_scalar(&a, 2.0).unwrap();
- let expected = Float64Array::from(vec![1.0, 4.0, 9.0]);
- assert_eq!(expected, actual);
- let a = Float64Array::from(vec![Some(1.0), None, Some(3.0)]);
- let actual = powf_scalar(&a, 2.0).unwrap();
- let expected = Float64Array::from(vec![Some(1.0), None, Some(9.0)]);
- assert_eq!(expected, actual);
- }
-}
diff --git a/arrow/src/compute/kernels/arity.rs b/arrow/src/compute/kernels/arity.rs
deleted file mode 100644
index 4aa7f3d..0000000
--- a/arrow/src/compute/kernels/arity.rs
+++ /dev/null
@@ -1,74 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-//! Defines kernels suitable to perform operations to primitive arrays.
-
-use crate::array::{Array, ArrayData, PrimitiveArray};
-use crate::buffer::Buffer;
-use crate::datatypes::ArrowPrimitiveType;
-
-#[inline]
-fn into_primitive_array_data<I: ArrowPrimitiveType, O: ArrowPrimitiveType>(
- array: &PrimitiveArray<I>,
- buffer: Buffer,
-) -> ArrayData {
- ArrayData::new(
- O::DATA_TYPE,
- array.len(),
- None,
- array.data_ref().null_buffer().cloned(),
- 0,
- vec![buffer],
- vec![],
- )
-}
-
-/// Applies an unary and infalible function to a primitive array.
-/// This is the fastest way to perform an operation on a primitive array when
-/// the benefits of a vectorized operation outweights the cost of branching nulls and non-nulls.
-/// # Implementation
-/// This will apply the function for all values, including those on null slots.
-/// This implies that the operation must be infalible for any value of the corresponding type
-/// or this function may panic.
-/// # Example
-/// ```rust
-/// # use arrow::array::Int32Array;
-/// # use arrow::datatypes::Int32Type;
-/// # use arrow::compute::kernels::arity::unary;
-/// # fn main() {
-/// let array = Int32Array::from(vec![Some(5), Some(7), None]);
-/// let c = unary::<_, _, Int32Type>(&array, |x| x * 2 + 1);
-/// assert_eq!(c, Int32Array::from(vec![Some(11), Some(15), None]));
-/// # }
-/// ```
-pub fn unary<I, F, O>(array: &PrimitiveArray<I>, op: F) -> PrimitiveArray<O>
-where
- I: ArrowPrimitiveType,
- O: ArrowPrimitiveType,
- F: Fn(I::Native) -> O::Native,
-{
- let values = array.values().iter().map(|v| op(*v));
- // JUSTIFICATION
- // Benefit
- // ~60% speedup
- // Soundness
- // `values` is an iterator with a known size because arrays are sized.
- let buffer = unsafe { Buffer::from_trusted_len_iter(values) };
-
- let data = into_primitive_array_data::<_, O>(array, buffer);
- PrimitiveArray::<O>::from(data)
-}
diff --git a/arrow/src/compute/kernels/boolean.rs b/arrow/src/compute/kernels/boolean.rs
deleted file mode 100644
index c9c0bc3..0000000
--- a/arrow/src/compute/kernels/boolean.rs
+++ /dev/null
@@ -1,1159 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-//! Defines boolean kernels on Arrow `BooleanArray`'s, e.g. `AND`, `OR` and `NOT`.
-//!
-//! These kernels can leverage SIMD if available on your system. Currently no runtime
-//! detection is provided, you should enable the specific SIMD intrinsics using
-//! `RUSTFLAGS="-C target-feature=+avx2"` for example. See the documentation
-//! [here](https://doc.rust-lang.org/stable/core/arch/) for more information.
-
-use std::ops::Not;
-
-use crate::array::{Array, ArrayData, BooleanArray, PrimitiveArray};
-use crate::buffer::{
- buffer_bin_and, buffer_bin_or, buffer_unary_not, Buffer, MutableBuffer,
-};
-use crate::compute::util::combine_option_bitmap;
-use crate::datatypes::{ArrowNumericType, DataType};
-use crate::error::{ArrowError, Result};
-use crate::util::bit_util::{ceil, round_upto_multiple_of_64};
-use core::iter;
-use lexical_core::Integer;
-
-fn binary_boolean_kleene_kernel<F>(
- left: &BooleanArray,
- right: &BooleanArray,
- op: F,
-) -> Result<BooleanArray>
-where
- F: Fn(u64, u64, u64, u64) -> (u64, u64),
-{
- if left.len() != right.len() {
- return Err(ArrowError::ComputeError(
- "Cannot perform bitwise operation on arrays of different length".to_string(),
- ));
- }
-
- // length and offset of boolean array is measured in bits
- let len = left.len();
-
- // result length measured in bytes (incl. remainder)
- let mut result_len = round_upto_multiple_of_64(len) / 8;
- // The iterator that applies the kleene_op closure always chains an additional iteration
- // for the remainder chunk, even without a remainder. If the remainder is absent
- // (length % 64 == 0), kleene_op would resize the result buffers (value_buffer and
- // valid_buffer) to store 8 additional bytes, because result_len wouldn't include a remainder
- // chunk. The resizing is unnecessary and expensive. We can prevent it by adding 8 bytes to
- // result_len here. Nonetheless, all bits of these 8 bytes will be 0.
- if len % 64 == 0 {
- result_len += 8;
- }
-
- let mut value_buffer = MutableBuffer::new(result_len);
- let mut valid_buffer = MutableBuffer::new(result_len);
-
- let kleene_op = |((left_data, left_valid), (right_data, right_valid)): (
- (u64, u64),
- (u64, u64),
- )| {
- let left_true = left_valid & left_data;
- let left_false = left_valid & !left_data;
-
- let right_true = right_valid & right_data;
- let right_false = right_valid & !right_data;
-
- let (value, valid) = op(left_true, left_false, right_true, right_false);
-
- value_buffer.extend_from_slice(&[value]);
- valid_buffer.extend_from_slice(&[valid]);
- };
-
- let left_offset = left.offset();
- let right_offset = right.offset();
-
- let left_buffer = left.values();
- let right_buffer = right.values();
-
- let left_chunks = left_buffer.bit_chunks(left_offset, len);
- let right_chunks = right_buffer.bit_chunks(right_offset, len);
-
- let left_rem = left_chunks.remainder_bits();
- let right_rem = right_chunks.remainder_bits();
-
- let opt_left_valid_chunks_and_rem = left
- .data_ref()
- .null_buffer()
- .map(|b| b.bit_chunks(left_offset, len))
- .map(|chunks| (chunks.iter(), chunks.remainder_bits()));
- let opt_right_valid_chunks_and_rem = right
- .data_ref()
- .null_buffer()
- .map(|b| b.bit_chunks(right_offset, len))
- .map(|chunks| (chunks.iter(), chunks.remainder_bits()));
-
- match (
- opt_left_valid_chunks_and_rem,
- opt_right_valid_chunks_and_rem,
- ) {
- (
- Some((left_valid_chunks, left_valid_rem)),
- Some((right_valid_chunks, right_valid_rem)),
- ) => {
- left_chunks
- .iter()
- .zip(left_valid_chunks)
- .zip(right_chunks.iter().zip(right_valid_chunks))
- .chain(iter::once((
- (left_rem, left_valid_rem),
- (right_rem, right_valid_rem),
- )))
- .for_each(kleene_op);
- }
- (Some((left_valid_chunks, left_valid_rem)), None) => {
- left_chunks
- .iter()
- .zip(left_valid_chunks)
- .zip(right_chunks.iter().zip(iter::repeat(u64::MAX)))
- .chain(iter::once((
- (left_rem, left_valid_rem),
- (right_rem, u64::MAX),
- )))
- .for_each(kleene_op);
- }
- (None, Some((right_valid_chunks, right_valid_rem))) => {
- left_chunks
- .iter()
- .zip(iter::repeat(u64::MAX))
- .zip(right_chunks.iter().zip(right_valid_chunks))
- .chain(iter::once((
- (left_rem, u64::MAX),
- (right_rem, right_valid_rem),
- )))
- .for_each(kleene_op);
- }
- (None, None) => {
- left_chunks
- .iter()
- .zip(iter::repeat(u64::MAX))
- .zip(right_chunks.iter().zip(iter::repeat(u64::MAX)))
- .chain(iter::once(((left_rem, u64::MAX), (right_rem, u64::MAX))))
- .for_each(kleene_op);
- }
- };
-
- let bool_buffer: Buffer = value_buffer.into();
- let bool_valid_buffer: Buffer = valid_buffer.into();
-
- let array_data = ArrayData::new(
- DataType::Boolean,
- len,
- None,
- Some(bool_valid_buffer),
- left_offset,
- vec![bool_buffer],
- vec![],
- );
-
- Ok(BooleanArray::from(array_data))
-}
-
-/// Helper function to implement binary kernels
-fn binary_boolean_kernel<F>(
- left: &BooleanArray,
- right: &BooleanArray,
- op: F,
-) -> Result<BooleanArray>
-where
- F: Fn(&Buffer, usize, &Buffer, usize, usize) -> Buffer,
-{
- if left.len() != right.len() {
- return Err(ArrowError::ComputeError(
- "Cannot perform bitwise operation on arrays of different length".to_string(),
- ));
- }
-
- let len = left.len();
-
- let left_data = left.data_ref();
- let right_data = right.data_ref();
- let null_bit_buffer = combine_option_bitmap(&left_data, &right_data, len)?;
-
- let left_buffer = &left_data.buffers()[0];
- let right_buffer = &right_data.buffers()[0];
- let left_offset = left.offset();
- let right_offset = right.offset();
-
- let values = op(&left_buffer, left_offset, &right_buffer, right_offset, len);
-
- let data = ArrayData::new(
- DataType::Boolean,
- len,
- None,
- null_bit_buffer,
- 0,
- vec![values],
- vec![],
- );
- Ok(BooleanArray::from(data))
-}
-
-/// Performs `AND` operation on two arrays. If either left or right value is null then the
-/// result is also null.
-/// # Error
-/// This function errors when the arrays have different lengths.
-/// # Example
-/// ```rust
-/// use arrow::array::BooleanArray;
-/// use arrow::error::Result;
-/// use arrow::compute::kernels::boolean::and;
-/// # fn main() -> Result<()> {
-/// let a = BooleanArray::from(vec![Some(false), Some(true), None]);
-/// let b = BooleanArray::from(vec![Some(true), Some(true), Some(false)]);
-/// let and_ab = and(&a, &b)?;
-/// assert_eq!(and_ab, BooleanArray::from(vec![Some(false), Some(true), None]));
-/// # Ok(())
-/// # }
-/// ```
-pub fn and(left: &BooleanArray, right: &BooleanArray) -> Result<BooleanArray> {
- binary_boolean_kernel(&left, &right, buffer_bin_and)
-}
-
-/// Logical 'and' boolean values with Kleene logic
-///
-/// # Behavior
-///
-/// This function behaves as follows with nulls:
-///
-/// * `true` and `null` = `null`
-/// * `null` and `true` = `null`
-/// * `false` and `null` = `false`
-/// * `null` and `false` = `false`
-/// * `null` and `null` = `null`
-///
-/// In other words, in this context a null value really means \"unknown\",
-/// and an unknown value 'and' false is always false.
-/// For a different null behavior, see function \"and\".
-///
-/// # Example
-///
-/// ```rust
-/// use arrow::array::BooleanArray;
-/// use arrow::error::Result;
-/// use arrow::compute::kernels::boolean::and_kleene;
-/// # fn main() -> Result<()> {
-/// let a = BooleanArray::from(vec![Some(true), Some(false), None]);
-/// let b = BooleanArray::from(vec![None, None, None]);
-/// let and_ab = and_kleene(&a, &b)?;
-/// assert_eq!(and_ab, BooleanArray::from(vec![None, Some(false), None]));
-/// # Ok(())
-/// # }
-/// ```
-///
-/// # Fails
-///
-/// If the operands have different lengths
-pub fn and_kleene(left: &BooleanArray, right: &BooleanArray) -> Result<BooleanArray> {
- if left.null_count().is_zero() && right.null_count().is_zero() {
- return and(left, right);
- }
-
- let op = |left_true, left_false, right_true, right_false| {
- (
- left_true & right_true,
- left_false | right_false | (left_true & right_true),
- )
- };
-
- binary_boolean_kleene_kernel(left, right, op)
-}
-
-/// Performs `OR` operation on two arrays. If either left or right value is null then the
-/// result is also null.
-/// # Error
-/// This function errors when the arrays have different lengths.
-/// # Example
-/// ```rust
-/// use arrow::array::BooleanArray;
-/// use arrow::error::Result;
-/// use arrow::compute::kernels::boolean::or;
-/// # fn main() -> Result<()> {
-/// let a = BooleanArray::from(vec![Some(false), Some(true), None]);
-/// let b = BooleanArray::from(vec![Some(true), Some(true), Some(false)]);
-/// let or_ab = or(&a, &b)?;
-/// assert_eq!(or_ab, BooleanArray::from(vec![Some(true), Some(true), None]));
-/// # Ok(())
-/// # }
-/// ```
-pub fn or(left: &BooleanArray, right: &BooleanArray) -> Result<BooleanArray> {
- binary_boolean_kernel(&left, &right, buffer_bin_or)
-}
-
-/// Logical 'or' boolean values with Kleene logic
-///
-/// # Behavior
-///
-/// This function behaves as follows with nulls:
-///
-/// * `true` or `null` = `true`
-/// * `null` or `true` = `true`
-/// * `false` or `null` = `null`
-/// * `null` or `false` = `null`
-/// * `null` or `null` = `null`
-///
-/// In other words, in this context a null value really means \"unknown\",
-/// and an unknown value 'or' true is always true.
-/// For a different null behavior, see function \"or\".
-///
-/// # Example
-///
-/// ```rust
-/// use arrow::array::BooleanArray;
-/// use arrow::error::Result;
-/// use arrow::compute::kernels::boolean::or_kleene;
-/// # fn main() -> Result<()> {
-/// let a = BooleanArray::from(vec![Some(true), Some(false), None]);
-/// let b = BooleanArray::from(vec![None, None, None]);
-/// let or_ab = or_kleene(&a, &b)?;
-/// assert_eq!(or_ab, BooleanArray::from(vec![Some(true), None, None]));
-/// # Ok(())
-/// # }
-/// ```
-///
-/// # Fails
-///
-/// If the operands have different lengths
-pub fn or_kleene(left: &BooleanArray, right: &BooleanArray) -> Result<BooleanArray> {
- if left.null_count().is_zero() && right.null_count().is_zero() {
- return or(left, right);
- }
-
- let op = |left_true, left_false, right_true, right_false| {
- (
- left_true | right_true,
- left_true | right_true | (left_false & right_false),
- )
- };
-
- binary_boolean_kleene_kernel(left, right, op)
-}
-
-/// Performs unary `NOT` operation on an arrays. If value is null then the result is also
-/// null.
-/// # Error
-/// This function never errors. It returns an error for consistency.
-/// # Example
-/// ```rust
-/// use arrow::array::BooleanArray;
-/// use arrow::error::Result;
-/// use arrow::compute::kernels::boolean::not;
-/// # fn main() -> Result<()> {
-/// let a = BooleanArray::from(vec![Some(false), Some(true), None]);
-/// let not_a = not(&a)?;
-/// assert_eq!(not_a, BooleanArray::from(vec![Some(true), Some(false), None]));
-/// # Ok(())
-/// # }
-/// ```
-pub fn not(left: &BooleanArray) -> Result<BooleanArray> {
- let left_offset = left.offset();
- let len = left.len();
-
- let data = left.data_ref();
- let null_bit_buffer = data
- .null_bitmap()
- .as_ref()
- .map(|b| b.bits.bit_slice(left_offset, len));
-
- let values = buffer_unary_not(&data.buffers()[0], left_offset, len);
-
- let data = ArrayData::new(
- DataType::Boolean,
- len,
- None,
- null_bit_buffer,
- 0,
- vec![values],
- vec![],
- );
- Ok(BooleanArray::from(data))
-}
-
-/// Returns a non-null [BooleanArray] with whether each value of the array is null.
-/// # Error
-/// This function never errors.
-/// # Example
-/// ```rust
-/// # use arrow::error::Result;
-/// use arrow::array::BooleanArray;
-/// use arrow::compute::kernels::boolean::is_null;
-/// # fn main() -> Result<()> {
-/// let a = BooleanArray::from(vec![Some(false), Some(true), None]);
-/// let a_is_null = is_null(&a)?;
-/// assert_eq!(a_is_null, BooleanArray::from(vec![false, false, true]));
-/// # Ok(())
-/// # }
-/// ```
-pub fn is_null(input: &Array) -> Result<BooleanArray> {
- let len = input.len();
-
- let output = match input.data_ref().null_buffer() {
- None => {
- let len_bytes = ceil(len, 8);
- MutableBuffer::from_len_zeroed(len_bytes).into()
- }
- Some(buffer) => buffer_unary_not(buffer, input.offset(), len),
- };
-
- let data =
- ArrayData::new(DataType::Boolean, len, None, None, 0, vec![output], vec![]);
-
- Ok(BooleanArray::from(data))
-}
-
-/// Returns a non-null [BooleanArray] with whether each value of the array is not null.
-/// # Error
-/// This function never errors.
-/// # Example
-/// ```rust
-/// # use arrow::error::Result;
-/// use arrow::array::BooleanArray;
-/// use arrow::compute::kernels::boolean::is_not_null;
-/// # fn main() -> Result<()> {
-/// let a = BooleanArray::from(vec![Some(false), Some(true), None]);
-/// let a_is_not_null = is_not_null(&a)?;
-/// assert_eq!(a_is_not_null, BooleanArray::from(vec![true, true, false]));
-/// # Ok(())
-/// # }
-/// ```
-pub fn is_not_null(input: &Array) -> Result<BooleanArray> {
- let len = input.len();
-
- let output = match input.data_ref().null_buffer() {
- None => {
- let len_bytes = ceil(len, 8);
- MutableBuffer::new(len_bytes)
- .with_bitset(len_bytes, true)
- .into()
- }
- Some(buffer) => buffer.bit_slice(input.offset(), len),
- };
-
- let data =
- ArrayData::new(DataType::Boolean, len, None, None, 0, vec![output], vec![]);
-
- Ok(BooleanArray::from(data))
-}
-
-/// Copies original array, setting null bit to true if a secondary comparison boolean array is set to true.
-/// Typically used to implement NULLIF.
-// NOTE: For now this only supports Primitive Arrays. Although the code could be made generic, the issue
-// is that currently the bitmap operations result in a final bitmap which is aligned to bit 0, and thus
-// the left array's data needs to be sliced to a new offset, and for non-primitive arrays shifting the
-// data might be too complicated. In the future, to avoid shifting left array's data, we could instead
-// shift the final bitbuffer to the right, prepending with 0's instead.
-pub fn nullif<T>(
- left: &PrimitiveArray<T>,
- right: &BooleanArray,
-) -> Result<PrimitiveArray<T>>
-where
- T: ArrowNumericType,
-{
- if left.len() != right.len() {
- return Err(ArrowError::ComputeError(
- "Cannot perform comparison operation on arrays of different length"
- .to_string(),
- ));
- }
- let left_data = left.data();
- let right_data = right.data();
-
- // If left has no bitmap, create a new one with all values set for nullity op later
- // left=0 (null) right=null output bitmap=null
- // left=0 right=1 output bitmap=null
- // left=1 (set) right=null output bitmap=set (passthrough)
- // left=1 right=1 & comp=true output bitmap=null
- // left=1 right=1 & comp=false output bitmap=set
- //
- // Thus: result = left null bitmap & (!right_values | !right_bitmap)
- // OR left null bitmap & !(right_values & right_bitmap)
- //
- // Do the right expression !(right_values & right_bitmap) first since there are two steps
- // TRICK: convert BooleanArray buffer as a bitmap for faster operation
- let right_combo_buffer = match right.data().null_bitmap() {
- Some(right_bitmap) => {
- // NOTE: right values and bitmaps are combined and stay at bit offset right.offset()
- (right.values() & &right_bitmap.bits).ok().map(|b| b.not())
- }
- None => Some(!right.values()),
- };
-
- // AND of original left null bitmap with right expression
- // Here we take care of the possible offsets of the left and right arrays all at once.
- let modified_null_buffer = match left_data.null_bitmap() {
- Some(left_null_bitmap) => match right_combo_buffer {
- Some(rcb) => Some(buffer_bin_and(
- &left_null_bitmap.bits,
- left_data.offset(),
- &rcb,
- right_data.offset(),
- left_data.len(),
- )),
- None => Some(
- left_null_bitmap
- .bits
- .bit_slice(left_data.offset(), left.len()),
- ),
- },
- None => right_combo_buffer
- .map(|rcb| rcb.bit_slice(right_data.offset(), right_data.len())),
- };
-
- // Align/shift left data on offset as needed, since new bitmaps are shifted and aligned to 0 already
- // NOTE: this probably only works for primitive arrays.
- let data_buffers = if left.offset() == 0 {
- left_data.buffers().to_vec()
- } else {
- // Shift each data buffer by type's bit_width * offset.
- left_data
- .buffers()
- .iter()
- .map(|buf| buf.slice(left.offset() * T::get_byte_width()))
- .collect::<Vec<_>>()
- };
-
- // Construct new array with same values but modified null bitmap
- // TODO: shift data buffer as needed
- let data = ArrayData::new(
- T::DATA_TYPE,
- left.len(),
- None, // force new to compute the number of null bits
- modified_null_buffer,
- 0, // No need for offset since left data has been shifted
- data_buffers,
- left_data.child_data().to_vec(),
- );
- Ok(PrimitiveArray::<T>::from(data))
-}
-
-#[cfg(test)]
-mod tests {
- use super::*;
- use crate::array::{ArrayRef, Int32Array};
- use std::sync::Arc;
-
- #[test]
- fn test_bool_array_and() {
- let a = BooleanArray::from(vec![false, false, true, true]);
- let b = BooleanArray::from(vec![false, true, false, true]);
- let c = and(&a, &b).unwrap();
-
- let expected = BooleanArray::from(vec![false, false, false, true]);
-
- assert_eq!(c, expected);
- }
-
- #[test]
- fn test_bool_array_or() {
- let a = BooleanArray::from(vec![false, false, true, true]);
- let b = BooleanArray::from(vec![false, true, false, true]);
- let c = or(&a, &b).unwrap();
-
- let expected = BooleanArray::from(vec![false, true, true, true]);
-
- assert_eq!(c, expected);
- }
-
- #[test]
- fn test_bool_array_or_nulls() {
- let a = BooleanArray::from(vec![
- None,
- None,
- None,
- Some(false),
- Some(false),
- Some(false),
- Some(true),
- Some(true),
- Some(true),
- ]);
- let b = BooleanArray::from(vec![
- None,
- Some(false),
- Some(true),
- None,
- Some(false),
- Some(true),
- None,
- Some(false),
- Some(true),
- ]);
- let c = or(&a, &b).unwrap();
-
- let expected = BooleanArray::from(vec![
- None,
- None,
- None,
- None,
- Some(false),
- Some(true),
- None,
- Some(true),
- Some(true),
- ]);
-
- assert_eq!(c, expected);
- }
-
- #[test]
- fn test_binary_boolean_kleene_kernel() {
- // the kleene kernel is based on chunking and we want to also create
- // cases, where the number of values is not a multiple of 64
- for &value in [true, false].iter() {
- for &is_valid in [true, false].iter() {
- for &n in [0usize, 1, 63, 64, 65, 127, 128].iter() {
- let a = BooleanArray::from(vec![Some(true); n]);
- let b = BooleanArray::from(vec![None; n]);
-
- let result = binary_boolean_kleene_kernel(&a, &b, |_, _, _, _| {
- let tmp_value = if value { u64::MAX } else { 0 };
- let tmp_is_valid = if is_valid { u64::MAX } else { 0 };
- (tmp_value, tmp_is_valid)
- })
- .unwrap();
-
- assert_eq!(result.len(), n);
- (0..n).for_each(|idx| {
- assert_eq!(value, result.value(idx));
- assert_eq!(is_valid, result.is_valid(idx));
- });
- }
- }
- }
- }
-
- #[test]
- fn test_boolean_array_kleene_no_remainder() {
- let n = 1024;
- let a = BooleanArray::from(vec![true; n]);
- let b = BooleanArray::from(vec![None; n]);
- let result = or_kleene(&a, &b).unwrap();
-
- assert_eq!(result, a);
- }
-
- #[test]
- fn test_bool_array_and_kleene_nulls() {
- let a = BooleanArray::from(vec![
- None,
- None,
- None,
- Some(false),
- Some(false),
- Some(false),
- Some(true),
- Some(true),
- Some(true),
- ]);
- let b = BooleanArray::from(vec![
- None,
- Some(false),
- Some(true),
- None,
- Some(false),
- Some(true),
- None,
- Some(false),
- Some(true),
- ]);
- let c = and_kleene(&a, &b).unwrap();
-
- let expected = BooleanArray::from(vec![
- None,
- Some(false),
- None,
- Some(false),
- Some(false),
- Some(false),
- None,
- Some(false),
- Some(true),
- ]);
-
- assert_eq!(c, expected);
- }
-
- #[test]
- fn test_bool_array_or_kleene_nulls() {
- let a = BooleanArray::from(vec![
- None,
- None,
- None,
- Some(false),
- Some(false),
- Some(false),
- Some(true),
- Some(true),
- Some(true),
- ]);
- let b = BooleanArray::from(vec![
- None,
- Some(false),
- Some(true),
- None,
- Some(false),
- Some(true),
- None,
- Some(false),
- Some(true),
- ]);
- let c = or_kleene(&a, &b).unwrap();
-
- let expected = BooleanArray::from(vec![
- None,
- None,
- Some(true),
- None,
- Some(false),
- Some(true),
- Some(true),
- Some(true),
- Some(true),
- ]);
-
- assert_eq!(c, expected);
- }
-
- #[test]
- fn test_bool_array_or_kleene_right_sided_nulls() {
- let a = BooleanArray::from(vec![false, false, false, true, true, true]);
-
- // ensure null bitmap of a is absent
- assert!(a.data_ref().null_bitmap().is_none());
-
- let b = BooleanArray::from(vec![
- Some(true),
- Some(false),
- None,
- Some(true),
- Some(false),
- None,
- ]);
-
- // ensure null bitmap of b is present
- assert!(b.data_ref().null_bitmap().is_some());
-
- let c = or_kleene(&a, &b).unwrap();
-
- let expected = BooleanArray::from(vec![
- Some(true),
- Some(false),
- None,
- Some(true),
- Some(true),
- Some(true),
- ]);
-
- assert_eq!(c, expected);
- }
-
- #[test]
- fn test_bool_array_or_kleene_left_sided_nulls() {
- let a = BooleanArray::from(vec![
- Some(true),
- Some(false),
- None,
- Some(true),
- Some(false),
- None,
- ]);
-
- // ensure null bitmap of b is absent
- assert!(a.data_ref().null_bitmap().is_some());
-
- let b = BooleanArray::from(vec![false, false, false, true, true, true]);
-
- // ensure null bitmap of a is present
- assert!(b.data_ref().null_bitmap().is_none());
-
- let c = or_kleene(&a, &b).unwrap();
-
- let expected = BooleanArray::from(vec![
- Some(true),
- Some(false),
- None,
- Some(true),
- Some(true),
- Some(true),
- ]);
-
- assert_eq!(c, expected);
- }
-
- #[test]
- fn test_bool_array_not() {
- let a = BooleanArray::from(vec![false, true]);
- let c = not(&a).unwrap();
-
- let expected = BooleanArray::from(vec![true, false]);
-
- assert_eq!(c, expected);
- }
-
- #[test]
- fn test_bool_array_not_sliced() {
- let a = BooleanArray::from(vec![None, Some(true), Some(false), None, Some(true)]);
- let a = a.slice(1, 4);
- let a = a.as_any().downcast_ref::<BooleanArray>().unwrap();
- let c = not(&a).unwrap();
-
- let expected =
- BooleanArray::from(vec![Some(false), Some(true), None, Some(false)]);
-
- assert_eq!(c, expected);
- }
-
- #[test]
- fn test_bool_array_and_nulls() {
- let a = BooleanArray::from(vec![
- None,
- None,
- None,
- Some(false),
- Some(false),
- Some(false),
- Some(true),
- Some(true),
- Some(true),
- ]);
- let b = BooleanArray::from(vec![
- None,
- Some(false),
- Some(true),
- None,
- Some(false),
- Some(true),
- None,
- Some(false),
- Some(true),
- ]);
- let c = and(&a, &b).unwrap();
-
- let expected = BooleanArray::from(vec![
- None,
- None,
- None,
- None,
- Some(false),
- Some(false),
- None,
- Some(false),
- Some(true),
- ]);
-
- assert_eq!(c, expected);
- }
-
- #[test]
- fn test_bool_array_and_sliced_same_offset() {
- let a = BooleanArray::from(vec![
- false, false, false, false, false, false, false, false, false, false, true,
- true,
- ]);
- let b = BooleanArray::from(vec![
- false, false, false, false, false, false, false, false, false, true, false,
- true,
- ]);
-
- let a = a.slice(8, 4);
- let a = a.as_any().downcast_ref::<BooleanArray>().unwrap();
- let b = b.slice(8, 4);
- let b = b.as_any().downcast_ref::<BooleanArray>().unwrap();
-
- let c = and(&a, &b).unwrap();
-
- let expected = BooleanArray::from(vec![false, false, false, true]);
-
- assert_eq!(expected, c);
- }
-
- #[test]
- fn test_bool_array_and_sliced_same_offset_mod8() {
- let a = BooleanArray::from(vec![
- false, false, true, true, false, false, false, false, false, false, false,
- false,
- ]);
- let b = BooleanArray::from(vec![
- false, false, false, false, false, false, false, false, false, true, false,
- true,
- ]);
-
- let a = a.slice(0, 4);
- let a = a.as_any().downcast_ref::<BooleanArray>().unwrap();
- let b = b.slice(8, 4);
- let b = b.as_any().downcast_ref::<BooleanArray>().unwrap();
-
- let c = and(&a, &b).unwrap();
-
- let expected = BooleanArray::from(vec![false, false, false, true]);
-
- assert_eq!(expected, c);
- }
-
- #[test]
- fn test_bool_array_and_sliced_offset1() {
- let a = BooleanArray::from(vec![
- false, false, false, false, false, false, false, false, false, false, true,
- true,
- ]);
- let b = BooleanArray::from(vec![false, true, false, true]);
-
- let a = a.slice(8, 4);
- let a = a.as_any().downcast_ref::<BooleanArray>().unwrap();
-
- let c = and(&a, &b).unwrap();
-
- let expected = BooleanArray::from(vec![false, false, false, true]);
-
- assert_eq!(expected, c);
- }
-
- #[test]
- fn test_bool_array_and_sliced_offset2() {
- let a = BooleanArray::from(vec![false, false, true, true]);
- let b = BooleanArray::from(vec![
- false, false, false, false, false, false, false, false, false, true, false,
- true,
- ]);
-
- let b = b.slice(8, 4);
- let b = b.as_any().downcast_ref::<BooleanArray>().unwrap();
-
- let c = and(&a, &b).unwrap();
-
- let expected = BooleanArray::from(vec![false, false, false, true]);
-
- assert_eq!(expected, c);
- }
-
- #[test]
- fn test_bool_array_and_nulls_offset() {
- let a = BooleanArray::from(vec![None, Some(false), Some(true), None, Some(true)]);
- let a = a.slice(1, 4);
- let a = a.as_any().downcast_ref::<BooleanArray>().unwrap();
-
- let b = BooleanArray::from(vec![
- None,
- None,
- Some(true),
- Some(false),
- Some(true),
- Some(true),
- ]);
-
- let b = b.slice(2, 4);
- let b = b.as_any().downcast_ref::<BooleanArray>().unwrap();
-
- let c = and(&a, &b).unwrap();
-
- let expected =
- BooleanArray::from(vec![Some(false), Some(false), None, Some(true)]);
-
- assert_eq!(expected, c);
- }
-
- #[test]
- fn test_nonnull_array_is_null() {
- let a: ArrayRef = Arc::new(Int32Array::from(vec![1, 2, 3, 4]));
-
- let res = is_null(a.as_ref()).unwrap();
-
- let expected = BooleanArray::from(vec![false, false, false, false]);
-
- assert_eq!(expected, res);
- assert_eq!(&None, res.data_ref().null_bitmap());
- }
-
- #[test]
- fn test_nonnull_array_with_offset_is_null() {
- let a = Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 7, 6, 5, 4, 3, 2, 1]);
- let a = a.slice(8, 4);
-
- let res = is_null(a.as_ref()).unwrap();
-
- let expected = BooleanArray::from(vec![false, false, false, false]);
-
- assert_eq!(expected, res);
- assert_eq!(&None, res.data_ref().null_bitmap());
- }
-
- #[test]
- fn test_nonnull_array_is_not_null() {
- let a = Int32Array::from(vec![1, 2, 3, 4]);
-
- let res = is_not_null(&a).unwrap();
-
- let expected = BooleanArray::from(vec![true, true, true, true]);
-
- assert_eq!(expected, res);
- assert_eq!(&None, res.data_ref().null_bitmap());
- }
-
- #[test]
- fn test_nonnull_array_with_offset_is_not_null() {
- let a = Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 7, 6, 5, 4, 3, 2, 1]);
- let a = a.slice(8, 4);
-
- let res = is_not_null(a.as_ref()).unwrap();
-
- let expected = BooleanArray::from(vec![true, true, true, true]);
-
- assert_eq!(expected, res);
- assert_eq!(&None, res.data_ref().null_bitmap());
- }
-
- #[test]
- fn test_nullable_array_is_null() {
- let a = Int32Array::from(vec![Some(1), None, Some(3), None]);
-
- let res = is_null(&a).unwrap();
-
- let expected = BooleanArray::from(vec![false, true, false, true]);
-
- assert_eq!(expected, res);
- assert_eq!(&None, res.data_ref().null_bitmap());
- }
-
- #[test]
- fn test_nullable_array_with_offset_is_null() {
- let a = Int32Array::from(vec![
- None,
- None,
- None,
- None,
- None,
- None,
- None,
- None,
- // offset 8, previous None values are skipped by the slice
- Some(1),
- None,
- Some(2),
- None,
- Some(3),
- Some(4),
- None,
- None,
- ]);
- let a = a.slice(8, 4);
-
- let res = is_null(a.as_ref()).unwrap();
-
- let expected = BooleanArray::from(vec![false, true, false, true]);
-
- assert_eq!(expected, res);
- assert_eq!(&None, res.data_ref().null_bitmap());
- }
-
- #[test]
- fn test_nullable_array_is_not_null() {
- let a = Int32Array::from(vec![Some(1), None, Some(3), None]);
-
- let res = is_not_null(&a).unwrap();
-
- let expected = BooleanArray::from(vec![true, false, true, false]);
-
- assert_eq!(expected, res);
- assert_eq!(&None, res.data_ref().null_bitmap());
- }
-
- #[test]
- fn test_nullable_array_with_offset_is_not_null() {
- let a = Int32Array::from(vec![
- None,
- None,
- None,
- None,
- None,
- None,
- None,
- None,
- // offset 8, previous None values are skipped by the slice
- Some(1),
- None,
- Some(2),
- None,
- Some(3),
- Some(4),
- None,
- None,
- ]);
- let a = a.slice(8, 4);
-
- let res = is_not_null(a.as_ref()).unwrap();
-
- let expected = BooleanArray::from(vec![true, false, true, false]);
-
- assert_eq!(expected, res);
- assert_eq!(&None, res.data_ref().null_bitmap());
- }
-
- #[test]
- fn test_nullif_int_array() {
- let a = Int32Array::from(vec![Some(15), None, Some(8), Some(1), Some(9)]);
- let comp =
- BooleanArray::from(vec![Some(false), None, Some(true), Some(false), None]);
- let res = nullif(&a, &comp).unwrap();
-
- let expected = Int32Array::from(vec![
- Some(15),
- None,
- None, // comp true, slot 2 turned into null
- Some(1),
- // Even though comp array / right is null, should still pass through original value
- // comp true, slot 2 turned into null
- Some(9),
- ]);
-
- assert_eq!(expected, res);
- }
-
- #[test]
- fn test_nullif_int_array_offset() {
- let a = Int32Array::from(vec![None, Some(15), Some(8), Some(1), Some(9)]);
- let a = a.slice(1, 3); // Some(15), Some(8), Some(1)
- let a = a.as_any().downcast_ref::<Int32Array>().unwrap();
- let comp = BooleanArray::from(vec![
- Some(false),
- Some(false),
- Some(false),
- None,
- Some(true),
- Some(false),
- None,
- ]);
- let comp = comp.slice(2, 3); // Some(false), None, Some(true)
- let comp = comp.as_any().downcast_ref::<BooleanArray>().unwrap();
- let res = nullif(&a, &comp).unwrap();
-
- let expected = Int32Array::from(vec![
- Some(15), // False => keep it
- Some(8), // None => keep it
- None, // true => None
- ]);
- assert_eq!(&expected, &res)
- }
-}
diff --git a/arrow/src/compute/kernels/cast.rs b/arrow/src/compute/kernels/cast.rs
deleted file mode 100644
index d755e1a..0000000
--- a/arrow/src/compute/kernels/cast.rs
+++ /dev/null
@@ -1,3871 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-//! Defines cast kernels for `ArrayRef`, to convert `Array`s between
-//! supported datatypes.
-//!
-//! Example:
-//!
-//! ```
-//! use arrow::array::*;
-//! use arrow::compute::cast;
-//! use arrow::datatypes::DataType;
-//! use std::sync::Arc;
-//!
-//! let a = Int32Array::from(vec![5, 6, 7]);
-//! let array = Arc::new(a) as ArrayRef;
-//! let b = cast(&array, &DataType::Float64).unwrap();
-//! let c = b.as_any().downcast_ref::<Float64Array>().unwrap();
-//! assert_eq!(5.0, c.value(0));
-//! assert_eq!(6.0, c.value(1));
-//! assert_eq!(7.0, c.value(2));
-//! ```
-
-use std::str;
-use std::sync::Arc;
-
-use crate::buffer::MutableBuffer;
-use crate::compute::kernels::arithmetic::{divide, multiply};
-use crate::compute::kernels::arity::unary;
-use crate::compute::kernels::cast_utils::string_to_timestamp_nanos;
-use crate::datatypes::*;
-use crate::error::{ArrowError, Result};
-use crate::{array::*, compute::take};
-use crate::{buffer::Buffer, util::serialization::lexical_to_string};
-use num::{NumCast, ToPrimitive};
-
-/// CastOptions provides a way to override the default cast behaviors
-#[derive(Debug)]
-pub struct CastOptions {
- /// how to handle cast failures, either return NULL (safe=true) or return ERR (safe=false)
- pub safe: bool,
-}
-
-pub const DEFAULT_CAST_OPTIONS: CastOptions = CastOptions { safe: true };
-
-/// Return true if a value of type `from_type` can be cast into a
-/// value of `to_type`. Note that such as cast may be lossy.
-///
-/// If this function returns true to stay consistent with the `cast` kernel below.
-pub fn can_cast_types(from_type: &DataType, to_type: &DataType) -> bool {
- use self::DataType::*;
- if from_type == to_type {
- return true;
- }
-
- match (from_type, to_type) {
- (Struct(_), _) => false,
- (_, Struct(_)) => false,
- (LargeList(list_from), LargeList(list_to)) => {
- can_cast_types(list_from.data_type(), list_to.data_type())
- }
- (List(list_from), List(list_to)) => {
- can_cast_types(list_from.data_type(), list_to.data_type())
- }
- (List(list_from), LargeList(list_to)) => {
- list_from.data_type() == list_to.data_type()
- }
- (List(_), _) => false,
- (_, List(list_to)) => can_cast_types(from_type, list_to.data_type()),
- (_, LargeList(list_to)) => can_cast_types(from_type, list_to.data_type()),
- (Dictionary(_, from_value_type), Dictionary(_, to_value_type)) => {
- can_cast_types(from_value_type, to_value_type)
- }
- (Dictionary(_, value_type), _) => can_cast_types(value_type, to_type),
- (_, Dictionary(_, value_type)) => can_cast_types(from_type, value_type),
-
- (_, Boolean) => DataType::is_numeric(from_type),
- (Boolean, _) => DataType::is_numeric(to_type) || to_type == &Utf8,
-
- (Utf8, LargeUtf8) => true,
- (LargeUtf8, Utf8) => true,
- (Utf8, Date32) => true,
- (Utf8, Date64) => true,
- (Utf8, Timestamp(TimeUnit::Nanosecond, None)) => true,
- (Utf8, _) => DataType::is_numeric(to_type),
- (LargeUtf8, Date32) => true,
- (LargeUtf8, Date64) => true,
- (LargeUtf8, Timestamp(TimeUnit::Nanosecond, None)) => true,
- (LargeUtf8, _) => DataType::is_numeric(to_type),
- (_, Utf8) | (_, LargeUtf8) => {
- DataType::is_numeric(from_type) || from_type == &Binary
- }
-
- // start numeric casts
- (UInt8, UInt16) => true,
- (UInt8, UInt32) => true,
- (UInt8, UInt64) => true,
- (UInt8, Int8) => true,
- (UInt8, Int16) => true,
- (UInt8, Int32) => true,
- (UInt8, Int64) => true,
- (UInt8, Float32) => true,
- (UInt8, Float64) => true,
-
- (UInt16, UInt8) => true,
- (UInt16, UInt32) => true,
- (UInt16, UInt64) => true,
- (UInt16, Int8) => true,
- (UInt16, Int16) => true,
- (UInt16, Int32) => true,
- (UInt16, Int64) => true,
- (UInt16, Float32) => true,
- (UInt16, Float64) => true,
-
- (UInt32, UInt8) => true,
- (UInt32, UInt16) => true,
- (UInt32, UInt64) => true,
- (UInt32, Int8) => true,
- (UInt32, Int16) => true,
- (UInt32, Int32) => true,
- (UInt32, Int64) => true,
- (UInt32, Float32) => true,
- (UInt32, Float64) => true,
-
- (UInt64, UInt8) => true,
- (UInt64, UInt16) => true,
- (UInt64, UInt32) => true,
- (UInt64, Int8) => true,
- (UInt64, Int16) => true,
- (UInt64, Int32) => true,
- (UInt64, Int64) => true,
- (UInt64, Float32) => true,
- (UInt64, Float64) => true,
-
- (Int8, UInt8) => true,
- (Int8, UInt16) => true,
- (Int8, UInt32) => true,
- (Int8, UInt64) => true,
- (Int8, Int16) => true,
- (Int8, Int32) => true,
- (Int8, Int64) => true,
- (Int8, Float32) => true,
- (Int8, Float64) => true,
-
- (Int16, UInt8) => true,
- (Int16, UInt16) => true,
- (Int16, UInt32) => true,
- (Int16, UInt64) => true,
- (Int16, Int8) => true,
- (Int16, Int32) => true,
- (Int16, Int64) => true,
- (Int16, Float32) => true,
- (Int16, Float64) => true,
-
- (Int32, UInt8) => true,
- (Int32, UInt16) => true,
- (Int32, UInt32) => true,
- (Int32, UInt64) => true,
- (Int32, Int8) => true,
- (Int32, Int16) => true,
- (Int32, Int64) => true,
- (Int32, Float32) => true,
- (Int32, Float64) => true,
-
- (Int64, UInt8) => true,
- (Int64, UInt16) => true,
- (Int64, UInt32) => true,
- (Int64, UInt64) => true,
- (Int64, Int8) => true,
- (Int64, Int16) => true,
- (Int64, Int32) => true,
- (Int64, Float32) => true,
- (Int64, Float64) => true,
-
- (Float32, UInt8) => true,
- (Float32, UInt16) => true,
- (Float32, UInt32) => true,
- (Float32, UInt64) => true,
- (Float32, Int8) => true,
- (Float32, Int16) => true,
- (Float32, Int32) => true,
- (Float32, Int64) => true,
- (Float32, Float64) => true,
-
- (Float64, UInt8) => true,
- (Float64, UInt16) => true,
- (Float64, UInt32) => true,
- (Float64, UInt64) => true,
- (Float64, Int8) => true,
- (Float64, Int16) => true,
- (Float64, Int32) => true,
- (Float64, Int64) => true,
- (Float64, Float32) => true,
- // end numeric casts
-
- // temporal casts
- (Int32, Date32) => true,
- (Int32, Date64) => true,
- (Int32, Time32(_)) => true,
- (Date32, Int32) => true,
- (Date32, Int64) => true,
- (Time32(_), Int32) => true,
- (Int64, Date64) => true,
- (Int64, Date32) => true,
- (Int64, Time64(_)) => true,
- (Date64, Int64) => true,
- (Date64, Int32) => true,
- (Time64(_), Int64) => true,
- (Date32, Date64) => true,
- (Date64, Date32) => true,
- (Time32(TimeUnit::Second), Time32(TimeUnit::Millisecond)) => true,
- (Time32(TimeUnit::Millisecond), Time32(TimeUnit::Second)) => true,
- (Time32(_), Time64(_)) => true,
- (Time64(TimeUnit::Microsecond), Time64(TimeUnit::Nanosecond)) => true,
- (Time64(TimeUnit::Nanosecond), Time64(TimeUnit::Microsecond)) => true,
- (Time64(_), Time32(to_unit)) => {
- matches!(to_unit, TimeUnit::Second | TimeUnit::Millisecond)
- }
- (Timestamp(_, _), Int64) => true,
- (Int64, Timestamp(_, _)) => true,
- (Timestamp(_, _), Timestamp(_, _)) => true,
- (Timestamp(_, _), Date32) => true,
- (Timestamp(_, _), Date64) => true,
- // date64 to timestamp might not make sense,
- (Int64, Duration(_)) => true,
- (Null, Int32) => true,
- (_, _) => false,
- }
-}
-
-/// Cast `array` to the provided data type and return a new Array with
-/// type `to_type`, if possible.
-///
-/// Behavior:
-/// * Boolean to Utf8: `true` => '1', `false` => `0`
-/// * Utf8 to numeric: strings that can't be parsed to numbers return null, float strings
-/// in integer casts return null
-/// * Numeric to boolean: 0 returns `false`, any other value returns `true`
-/// * List to List: the underlying data type is cast
-/// * Primitive to List: a list array with 1 value per slot is created
-/// * Date32 and Date64: precision lost when going to higher interval
-/// * Time32 and Time64: precision lost when going to higher interval
-/// * Timestamp and Date{32|64}: precision lost when going to higher interval
-/// * Temporal to/from backing primitive: zero-copy with data type change
-///
-/// Unsupported Casts
-/// * To or from `StructArray`
-/// * List to primitive
-/// * Utf8 to boolean
-/// * Interval and duration
-pub fn cast(array: &ArrayRef, to_type: &DataType) -> Result<ArrayRef> {
- cast_with_options(array, to_type, &DEFAULT_CAST_OPTIONS)
-}
-
-/// Cast `array` to the provided data type and return a new Array with
-/// type `to_type`, if possible. It accepts `CastOptions` to allow consumers
-/// to configure cast behavior.
-///
-/// Behavior:
-/// * Boolean to Utf8: `true` => '1', `false` => `0`
-/// * Utf8 to numeric: strings that can't be parsed to numbers return null, float strings
-/// in integer casts return null
-/// * Numeric to boolean: 0 returns `false`, any other value returns `true`
-/// * List to List: the underlying data type is cast
-/// * Primitive to List: a list array with 1 value per slot is created
-/// * Date32 and Date64: precision lost when going to higher interval
-/// * Time32 and Time64: precision lost when going to higher interval
-/// * Timestamp and Date{32|64}: precision lost when going to higher interval
-/// * Temporal to/from backing primitive: zero-copy with data type change
-///
-/// Unsupported Casts
-/// * To or from `StructArray`
-/// * List to primitive
-/// * Utf8 to boolean
-/// * Interval and duration
-pub fn cast_with_options(
- array: &ArrayRef,
- to_type: &DataType,
- cast_options: &CastOptions,
-) -> Result<ArrayRef> {
- use DataType::*;
- let from_type = array.data_type();
-
- // clone array if types are the same
- if from_type == to_type {
- return Ok(array.clone());
- }
- match (from_type, to_type) {
- (Struct(_), _) => Err(ArrowError::CastError(
- "Cannot cast from struct to other types".to_string(),
- )),
- (_, Struct(_)) => Err(ArrowError::CastError(
- "Cannot cast to struct from other types".to_string(),
- )),
- (List(_), List(ref to)) => {
- cast_list_inner::<i32>(array, to, to_type, cast_options)
- }
- (LargeList(_), LargeList(ref to)) => {
- cast_list_inner::<i64>(array, to, to_type, cast_options)
- }
- (List(list_from), LargeList(list_to)) => {
- if list_to.data_type() != list_from.data_type() {
- Err(ArrowError::CastError(
- "cannot cast list to large-list with different child data".into(),
- ))
- } else {
- cast_list_container::<i32, i64>(&**array, cast_options)
- }
- }
- (LargeList(list_from), List(list_to)) => {
- if list_to.data_type() != list_from.data_type() {
- Err(ArrowError::CastError(
- "cannot cast large-list to list with different child data".into(),
- ))
- } else {
- cast_list_container::<i64, i32>(&**array, cast_options)
- }
- }
- (List(_), _) => Err(ArrowError::CastError(
- "Cannot cast list to non-list data types".to_string(),
- )),
- (_, List(ref to)) => {
- cast_primitive_to_list::<i32>(array, to, to_type, cast_options)
- }
- (_, LargeList(ref to)) => {
- cast_primitive_to_list::<i64>(array, to, to_type, cast_options)
- }
- (Dictionary(index_type, _), _) => match **index_type {
- DataType::Int8 => dictionary_cast::<Int8Type>(array, to_type, cast_options),
- DataType::Int16 => dictionary_cast::<Int16Type>(array, to_type, cast_options),
- DataType::Int32 => dictionary_cast::<Int32Type>(array, to_type, cast_options),
- DataType::Int64 => dictionary_cast::<Int64Type>(array, to_type, cast_options),
- DataType::UInt8 => dictionary_cast::<UInt8Type>(array, to_type, cast_options),
- DataType::UInt16 => {
- dictionary_cast::<UInt16Type>(array, to_type, cast_options)
- }
- DataType::UInt32 => {
- dictionary_cast::<UInt32Type>(array, to_type, cast_options)
- }
- DataType::UInt64 => {
- dictionary_cast::<UInt64Type>(array, to_type, cast_options)
- }
- _ => Err(ArrowError::CastError(format!(
- "Casting from dictionary type {:?} to {:?} not supported",
- from_type, to_type,
- ))),
- },
- (_, Dictionary(index_type, value_type)) => match **index_type {
- DataType::Int8 => {
- cast_to_dictionary::<Int8Type>(array, value_type, cast_options)
- }
- DataType::Int16 => {
- cast_to_dictionary::<Int16Type>(array, value_type, cast_options)
- }
- DataType::Int32 => {
- cast_to_dictionary::<Int32Type>(array, value_type, cast_options)
- }
- DataType::Int64 => {
- cast_to_dictionary::<Int64Type>(array, value_type, cast_options)
- }
- DataType::UInt8 => {
- cast_to_dictionary::<UInt8Type>(array, value_type, cast_options)
- }
- DataType::UInt16 => {
- cast_to_dictionary::<UInt16Type>(array, value_type, cast_options)
- }
- DataType::UInt32 => {
- cast_to_dictionary::<UInt32Type>(array, value_type, cast_options)
- }
- DataType::UInt64 => {
- cast_to_dictionary::<UInt64Type>(array, value_type, cast_options)
- }
- _ => Err(ArrowError::CastError(format!(
- "Casting from type {:?} to dictionary type {:?} not supported",
- from_type, to_type,
- ))),
- },
- (_, Boolean) => match from_type {
- UInt8 => cast_numeric_to_bool::<UInt8Type>(array),
- UInt16 => cast_numeric_to_bool::<UInt16Type>(array),
- UInt32 => cast_numeric_to_bool::<UInt32Type>(array),
- UInt64 => cast_numeric_to_bool::<UInt64Type>(array),
- Int8 => cast_numeric_to_bool::<Int8Type>(array),
- Int16 => cast_numeric_to_bool::<Int16Type>(array),
- Int32 => cast_numeric_to_bool::<Int32Type>(array),
- Int64 => cast_numeric_to_bool::<Int64Type>(array),
- Float32 => cast_numeric_to_bool::<Float32Type>(array),
- Float64 => cast_numeric_to_bool::<Float64Type>(array),
- Utf8 => Err(ArrowError::CastError(format!(
- "Casting from {:?} to {:?} not supported",
- from_type, to_type,
- ))),
- _ => Err(ArrowError::CastError(format!(
- "Casting from {:?} to {:?} not supported",
- from_type, to_type,
- ))),
- },
- (Boolean, _) => match to_type {
- UInt8 => cast_bool_to_numeric::<UInt8Type>(array, cast_options),
- UInt16 => cast_bool_to_numeric::<UInt16Type>(array, cast_options),
- UInt32 => cast_bool_to_numeric::<UInt32Type>(array, cast_options),
- UInt64 => cast_bool_to_numeric::<UInt64Type>(array, cast_options),
- Int8 => cast_bool_to_numeric::<Int8Type>(array, cast_options),
- Int16 => cast_bool_to_numeric::<Int16Type>(array, cast_options),
- Int32 => cast_bool_to_numeric::<Int32Type>(array, cast_options),
- Int64 => cast_bool_to_numeric::<Int64Type>(array, cast_options),
- Float32 => cast_bool_to_numeric::<Float32Type>(array, cast_options),
- Float64 => cast_bool_to_numeric::<Float64Type>(array, cast_options),
- Utf8 => {
- let array = array.as_any().downcast_ref::<BooleanArray>().unwrap();
- Ok(Arc::new(
- array
- .iter()
- .map(|value| value.map(|value| if value { "1" } else { "0" }))
- .collect::<StringArray>(),
- ))
- }
- _ => Err(ArrowError::CastError(format!(
- "Casting from {:?} to {:?} not supported",
- from_type, to_type,
- ))),
- },
- (Utf8, _) => match to_type {
- LargeUtf8 => cast_str_container::<i32, i64>(&**array),
- UInt8 => cast_string_to_numeric::<UInt8Type, i32>(array, cast_options),
- UInt16 => cast_string_to_numeric::<UInt16Type, i32>(array, cast_options),
- UInt32 => cast_string_to_numeric::<UInt32Type, i32>(array, cast_options),
- UInt64 => cast_string_to_numeric::<UInt64Type, i32>(array, cast_options),
- Int8 => cast_string_to_numeric::<Int8Type, i32>(array, cast_options),
- Int16 => cast_string_to_numeric::<Int16Type, i32>(array, cast_options),
- Int32 => cast_string_to_numeric::<Int32Type, i32>(array, cast_options),
- Int64 => cast_string_to_numeric::<Int64Type, i32>(array, cast_options),
- Float32 => cast_string_to_numeric::<Float32Type, i32>(array, cast_options),
- Float64 => cast_string_to_numeric::<Float64Type, i32>(array, cast_options),
- Date32 => cast_string_to_date32::<i32>(&**array, cast_options),
- Date64 => cast_string_to_date64::<i32>(&**array, cast_options),
- Timestamp(TimeUnit::Nanosecond, None) => {
- cast_string_to_timestamp_ns::<i32>(&**array, cast_options)
- }
- _ => Err(ArrowError::CastError(format!(
- "Casting from {:?} to {:?} not supported",
- from_type, to_type,
- ))),
- },
- (_, Utf8) => match from_type {
- LargeUtf8 => cast_str_container::<i64, i32>(&**array),
- UInt8 => cast_numeric_to_string::<UInt8Type, i32>(array),
- UInt16 => cast_numeric_to_string::<UInt16Type, i32>(array),
- UInt32 => cast_numeric_to_string::<UInt32Type, i32>(array),
- UInt64 => cast_numeric_to_string::<UInt64Type, i32>(array),
- Int8 => cast_numeric_to_string::<Int8Type, i32>(array),
- Int16 => cast_numeric_to_string::<Int16Type, i32>(array),
- Int32 => cast_numeric_to_string::<Int32Type, i32>(array),
- Int64 => cast_numeric_to_string::<Int64Type, i32>(array),
- Float32 => cast_numeric_to_string::<Float32Type, i32>(array),
- Float64 => cast_numeric_to_string::<Float64Type, i32>(array),
- Binary => {
- let array = array.as_any().downcast_ref::<BinaryArray>().unwrap();
- Ok(Arc::new(
- array
- .iter()
- .map(|maybe_value| match maybe_value {
- Some(value) => {
- let result = str::from_utf8(value);
- if cast_options.safe {
- Ok(result.ok())
- } else {
- Some(result.map_err(|_| {
- ArrowError::CastError(
- "Cannot cast binary to string".to_string(),
- )
- }))
- .transpose()
- }
- }
- None => Ok(None),
- })
- .collect::<Result<StringArray>>()?,
- ))
- }
- _ => Err(ArrowError::CastError(format!(
- "Casting from {:?} to {:?} not supported",
- from_type, to_type,
- ))),
- },
- (_, LargeUtf8) => match from_type {
- UInt8 => cast_numeric_to_string::<UInt8Type, i64>(array),
- UInt16 => cast_numeric_to_string::<UInt16Type, i64>(array),
- UInt32 => cast_numeric_to_string::<UInt32Type, i64>(array),
- UInt64 => cast_numeric_to_string::<UInt64Type, i64>(array),
- Int8 => cast_numeric_to_string::<Int8Type, i64>(array),
- Int16 => cast_numeric_to_string::<Int16Type, i64>(array),
- Int32 => cast_numeric_to_string::<Int32Type, i64>(array),
- Int64 => cast_numeric_to_string::<Int64Type, i64>(array),
- Float32 => cast_numeric_to_string::<Float32Type, i64>(array),
- Float64 => cast_numeric_to_string::<Float64Type, i64>(array),
- Binary => {
- let array = array.as_any().downcast_ref::<BinaryArray>().unwrap();
- Ok(Arc::new(
- array
- .iter()
- .map(|maybe_value| match maybe_value {
- Some(value) => {
- let result = str::from_utf8(value);
- if cast_options.safe {
- Ok(result.ok())
- } else {
- Some(result.map_err(|_| {
- ArrowError::CastError(
- "Cannot cast binary to string".to_string(),
- )
- }))
- .transpose()
- }
- }
- None => Ok(None),
- })
- .collect::<Result<LargeStringArray>>()?,
- ))
- }
- _ => Err(ArrowError::CastError(format!(
- "Casting from {:?} to {:?} not supported",
- from_type, to_type,
- ))),
- },
- (LargeUtf8, _) => match to_type {
- UInt8 => cast_string_to_numeric::<UInt8Type, i64>(array, cast_options),
- UInt16 => cast_string_to_numeric::<UInt16Type, i64>(array, cast_options),
- UInt32 => cast_string_to_numeric::<UInt32Type, i64>(array, cast_options),
- UInt64 => cast_string_to_numeric::<UInt64Type, i64>(array, cast_options),
- Int8 => cast_string_to_numeric::<Int8Type, i64>(array, cast_options),
- Int16 => cast_string_to_numeric::<Int16Type, i64>(array, cast_options),
- Int32 => cast_string_to_numeric::<Int32Type, i64>(array, cast_options),
- Int64 => cast_string_to_numeric::<Int64Type, i64>(array, cast_options),
- Float32 => cast_string_to_numeric::<Float32Type, i64>(array, cast_options),
- Float64 => cast_string_to_numeric::<Float64Type, i64>(array, cast_options),
- Date32 => cast_string_to_date32::<i64>(&**array, cast_options),
- Date64 => cast_string_to_date64::<i64>(&**array, cast_options),
- Timestamp(TimeUnit::Nanosecond, None) => {
- cast_string_to_timestamp_ns::<i64>(&**array, cast_options)
- }
- _ => Err(ArrowError::CastError(format!(
- "Casting from {:?} to {:?} not supported",
- from_type, to_type,
- ))),
- },
-
- // start numeric casts
- (UInt8, UInt16) => cast_numeric_arrays::<UInt8Type, UInt16Type>(array),
- (UInt8, UInt32) => cast_numeric_arrays::<UInt8Type, UInt32Type>(array),
- (UInt8, UInt64) => cast_numeric_arrays::<UInt8Type, UInt64Type>(array),
- (UInt8, Int8) => cast_numeric_arrays::<UInt8Type, Int8Type>(array),
- (UInt8, Int16) => cast_numeric_arrays::<UInt8Type, Int16Type>(array),
- (UInt8, Int32) => cast_numeric_arrays::<UInt8Type, Int32Type>(array),
- (UInt8, Int64) => cast_numeric_arrays::<UInt8Type, Int64Type>(array),
- (UInt8, Float32) => cast_numeric_arrays::<UInt8Type, Float32Type>(array),
- (UInt8, Float64) => cast_numeric_arrays::<UInt8Type, Float64Type>(array),
-
- (UInt16, UInt8) => cast_numeric_arrays::<UInt16Type, UInt8Type>(array),
- (UInt16, UInt32) => cast_numeric_arrays::<UInt16Type, UInt32Type>(array),
- (UInt16, UInt64) => cast_numeric_arrays::<UInt16Type, UInt64Type>(array),
- (UInt16, Int8) => cast_numeric_arrays::<UInt16Type, Int8Type>(array),
- (UInt16, Int16) => cast_numeric_arrays::<UInt16Type, Int16Type>(array),
- (UInt16, Int32) => cast_numeric_arrays::<UInt16Type, Int32Type>(array),
- (UInt16, Int64) => cast_numeric_arrays::<UInt16Type, Int64Type>(array),
- (UInt16, Float32) => cast_numeric_arrays::<UInt16Type, Float32Type>(array),
- (UInt16, Float64) => cast_numeric_arrays::<UInt16Type, Float64Type>(array),
-
- (UInt32, UInt8) => cast_numeric_arrays::<UInt32Type, UInt8Type>(array),
- (UInt32, UInt16) => cast_numeric_arrays::<UInt32Type, UInt16Type>(array),
- (UInt32, UInt64) => cast_numeric_arrays::<UInt32Type, UInt64Type>(array),
- (UInt32, Int8) => cast_numeric_arrays::<UInt32Type, Int8Type>(array),
- (UInt32, Int16) => cast_numeric_arrays::<UInt32Type, Int16Type>(array),
- (UInt32, Int32) => cast_numeric_arrays::<UInt32Type, Int32Type>(array),
- (UInt32, Int64) => cast_numeric_arrays::<UInt32Type, Int64Type>(array),
- (UInt32, Float32) => cast_numeric_arrays::<UInt32Type, Float32Type>(array),
- (UInt32, Float64) => cast_numeric_arrays::<UInt32Type, Float64Type>(array),
-
- (UInt64, UInt8) => cast_numeric_arrays::<UInt64Type, UInt8Type>(array),
- (UInt64, UInt16) => cast_numeric_arrays::<UInt64Type, UInt16Type>(array),
- (UInt64, UInt32) => cast_numeric_arrays::<UInt64Type, UInt32Type>(array),
- (UInt64, Int8) => cast_numeric_arrays::<UInt64Type, Int8Type>(array),
- (UInt64, Int16) => cast_numeric_arrays::<UInt64Type, Int16Type>(array),
- (UInt64, Int32) => cast_numeric_arrays::<UInt64Type, Int32Type>(array),
- (UInt64, Int64) => cast_numeric_arrays::<UInt64Type, Int64Type>(array),
- (UInt64, Float32) => cast_numeric_arrays::<UInt64Type, Float32Type>(array),
- (UInt64, Float64) => cast_numeric_arrays::<UInt64Type, Float64Type>(array),
-
- (Int8, UInt8) => cast_numeric_arrays::<Int8Type, UInt8Type>(array),
- (Int8, UInt16) => cast_numeric_arrays::<Int8Type, UInt16Type>(array),
- (Int8, UInt32) => cast_numeric_arrays::<Int8Type, UInt32Type>(array),
- (Int8, UInt64) => cast_numeric_arrays::<Int8Type, UInt64Type>(array),
- (Int8, Int16) => cast_numeric_arrays::<Int8Type, Int16Type>(array),
- (Int8, Int32) => cast_numeric_arrays::<Int8Type, Int32Type>(array),
- (Int8, Int64) => cast_numeric_arrays::<Int8Type, Int64Type>(array),
- (Int8, Float32) => cast_numeric_arrays::<Int8Type, Float32Type>(array),
- (Int8, Float64) => cast_numeric_arrays::<Int8Type, Float64Type>(array),
-
- (Int16, UInt8) => cast_numeric_arrays::<Int16Type, UInt8Type>(array),
- (Int16, UInt16) => cast_numeric_arrays::<Int16Type, UInt16Type>(array),
- (Int16, UInt32) => cast_numeric_arrays::<Int16Type, UInt32Type>(array),
- (Int16, UInt64) => cast_numeric_arrays::<Int16Type, UInt64Type>(array),
- (Int16, Int8) => cast_numeric_arrays::<Int16Type, Int8Type>(array),
- (Int16, Int32) => cast_numeric_arrays::<Int16Type, Int32Type>(array),
- (Int16, Int64) => cast_numeric_arrays::<Int16Type, Int64Type>(array),
- (Int16, Float32) => cast_numeric_arrays::<Int16Type, Float32Type>(array),
- (Int16, Float64) => cast_numeric_arrays::<Int16Type, Float64Type>(array),
-
- (Int32, UInt8) => cast_numeric_arrays::<Int32Type, UInt8Type>(array),
- (Int32, UInt16) => cast_numeric_arrays::<Int32Type, UInt16Type>(array),
- (Int32, UInt32) => cast_numeric_arrays::<Int32Type, UInt32Type>(array),
- (Int32, UInt64) => cast_numeric_arrays::<Int32Type, UInt64Type>(array),
- (Int32, Int8) => cast_numeric_arrays::<Int32Type, Int8Type>(array),
- (Int32, Int16) => cast_numeric_arrays::<Int32Type, Int16Type>(array),
- (Int32, Int64) => cast_numeric_arrays::<Int32Type, Int64Type>(array),
- (Int32, Float32) => cast_numeric_arrays::<Int32Type, Float32Type>(array),
- (Int32, Float64) => cast_numeric_arrays::<Int32Type, Float64Type>(array),
-
- (Int64, UInt8) => cast_numeric_arrays::<Int64Type, UInt8Type>(array),
- (Int64, UInt16) => cast_numeric_arrays::<Int64Type, UInt16Type>(array),
- (Int64, UInt32) => cast_numeric_arrays::<Int64Type, UInt32Type>(array),
- (Int64, UInt64) => cast_numeric_arrays::<Int64Type, UInt64Type>(array),
- (Int64, Int8) => cast_numeric_arrays::<Int64Type, Int8Type>(array),
- (Int64, Int16) => cast_numeric_arrays::<Int64Type, Int16Type>(array),
- (Int64, Int32) => cast_numeric_arrays::<Int64Type, Int32Type>(array),
- (Int64, Float32) => cast_numeric_arrays::<Int64Type, Float32Type>(array),
- (Int64, Float64) => cast_numeric_arrays::<Int64Type, Float64Type>(array),
-
- (Float32, UInt8) => cast_numeric_arrays::<Float32Type, UInt8Type>(array),
- (Float32, UInt16) => cast_numeric_arrays::<Float32Type, UInt16Type>(array),
- (Float32, UInt32) => cast_numeric_arrays::<Float32Type, UInt32Type>(array),
- (Float32, UInt64) => cast_numeric_arrays::<Float32Type, UInt64Type>(array),
- (Float32, Int8) => cast_numeric_arrays::<Float32Type, Int8Type>(array),
- (Float32, Int16) => cast_numeric_arrays::<Float32Type, Int16Type>(array),
- (Float32, Int32) => cast_numeric_arrays::<Float32Type, Int32Type>(array),
- (Float32, Int64) => cast_numeric_arrays::<Float32Type, Int64Type>(array),
- (Float32, Float64) => cast_numeric_arrays::<Float32Type, Float64Type>(array),
-
- (Float64, UInt8) => cast_numeric_arrays::<Float64Type, UInt8Type>(array),
- (Float64, UInt16) => cast_numeric_arrays::<Float64Type, UInt16Type>(array),
- (Float64, UInt32) => cast_numeric_arrays::<Float64Type, UInt32Type>(array),
- (Float64, UInt64) => cast_numeric_arrays::<Float64Type, UInt64Type>(array),
- (Float64, Int8) => cast_numeric_arrays::<Float64Type, Int8Type>(array),
- (Float64, Int16) => cast_numeric_arrays::<Float64Type, Int16Type>(array),
- (Float64, Int32) => cast_numeric_arrays::<Float64Type, Int32Type>(array),
- (Float64, Int64) => cast_numeric_arrays::<Float64Type, Int64Type>(array),
- (Float64, Float32) => cast_numeric_arrays::<Float64Type, Float32Type>(array),
- // end numeric casts
-
- // temporal casts
- (Int32, Date32) => cast_array_data::<Date32Type>(array, to_type.clone()),
- (Int32, Date64) => cast_with_options(
- &cast_with_options(array, &DataType::Date32, &cast_options)?,
- &DataType::Date64,
- &cast_options,
- ),
- (Int32, Time32(TimeUnit::Second)) => {
- cast_array_data::<Time32SecondType>(array, to_type.clone())
- }
- (Int32, Time32(TimeUnit::Millisecond)) => {
- cast_array_data::<Time32MillisecondType>(array, to_type.clone())
- }
- // No support for microsecond/nanosecond with i32
- (Date32, Int32) => cast_array_data::<Int32Type>(array, to_type.clone()),
- (Date32, Int64) => cast_with_options(
- &cast_with_options(array, &DataType::Int32, cast_options)?,
- &DataType::Int64,
- &cast_options,
- ),
- (Time32(_), Int32) => cast_array_data::<Int32Type>(array, to_type.clone()),
- (Int64, Date64) => cast_array_data::<Date64Type>(array, to_type.clone()),
- (Int64, Date32) => cast_with_options(
- &cast_with_options(array, &DataType::Int32, &cast_options)?,
- &DataType::Date32,
- &cast_options,
- ),
- // No support for second/milliseconds with i64
- (Int64, Time64(TimeUnit::Microsecond)) => {
- cast_array_data::<Time64MicrosecondType>(array, to_type.clone())
- }
- (Int64, Time64(TimeUnit::Nanosecond)) => {
- cast_array_data::<Time64NanosecondType>(array, to_type.clone())
- }
-
- (Date64, Int64) => cast_array_data::<Int64Type>(array, to_type.clone()),
- (Date64, Int32) => cast_with_options(
- &cast_with_options(array, &DataType::Int64, &cast_options)?,
- &DataType::Int32,
- &cast_options,
- ),
- (Time64(_), Int64) => cast_array_data::<Int64Type>(array, to_type.clone()),
- (Date32, Date64) => {
- let date_array = array.as_any().downcast_ref::<Date32Array>().unwrap();
-
- let values =
- unary::<_, _, Date64Type>(date_array, |x| x as i64 * MILLISECONDS_IN_DAY);
-
- Ok(Arc::new(values) as ArrayRef)
- }
- (Date64, Date32) => {
- let date_array = array.as_any().downcast_ref::<Date64Array>().unwrap();
-
- let values = unary::<_, _, Date32Type>(date_array, |x| {
- (x / MILLISECONDS_IN_DAY) as i32
- });
-
- Ok(Arc::new(values) as ArrayRef)
- }
- (Time32(TimeUnit::Second), Time32(TimeUnit::Millisecond)) => {
- let time_array = array.as_any().downcast_ref::<Time32SecondArray>().unwrap();
-
- let values = unary::<_, _, Time32MillisecondType>(time_array, |x| {
- x * MILLISECONDS as i32
- });
-
- Ok(Arc::new(values) as ArrayRef)
- }
- (Time32(TimeUnit::Millisecond), Time32(TimeUnit::Second)) => {
- let time_array = array
- .as_any()
- .downcast_ref::<Time32MillisecondArray>()
- .unwrap();
-
- let values = unary::<_, _, Time32SecondType>(time_array, |x| {
- x / (MILLISECONDS as i32)
- });
-
- Ok(Arc::new(values) as ArrayRef)
- }
- //(Time32(TimeUnit::Second), Time64(_)) => {},
- (Time32(from_unit), Time64(to_unit)) => {
- let time_array = Int32Array::from(array.data().clone());
- // note: (numeric_cast + SIMD multiply) is faster than (cast & multiply)
- let c: Int64Array = numeric_cast(&time_array);
- let from_size = time_unit_multiple(&from_unit);
- let to_size = time_unit_multiple(&to_unit);
- // from is only smaller than to if 64milli/64second don't exist
- let mult = Int64Array::from(vec![to_size / from_size; array.len()]);
- let converted = multiply(&c, &mult)?;
- let array_ref = Arc::new(converted) as ArrayRef;
- use TimeUnit::*;
- match to_unit {
- Microsecond => cast_array_data::<TimestampMicrosecondType>(
- &array_ref,
- to_type.clone(),
- ),
- Nanosecond => cast_array_data::<TimestampNanosecondType>(
- &array_ref,
- to_type.clone(),
- ),
- _ => unreachable!("array type not supported"),
- }
- }
- (Time64(TimeUnit::Microsecond), Time64(TimeUnit::Nanosecond)) => {
- let time_array = array
- .as_any()
- .downcast_ref::<Time64MicrosecondArray>()
- .unwrap();
-
- let values =
- unary::<_, _, Time64NanosecondType>(time_array, |x| x * MILLISECONDS);
- Ok(Arc::new(values) as ArrayRef)
- }
- (Time64(TimeUnit::Nanosecond), Time64(TimeUnit::Microsecond)) => {
- let time_array = array
- .as_any()
- .downcast_ref::<Time64NanosecondArray>()
- .unwrap();
-
- let values =
- unary::<_, _, Time64MicrosecondType>(time_array, |x| x / MILLISECONDS);
- Ok(Arc::new(values) as ArrayRef)
- }
- (Time64(from_unit), Time32(to_unit)) => {
- let time_array = Int64Array::from(array.data().clone());
- let from_size = time_unit_multiple(&from_unit);
- let to_size = time_unit_multiple(&to_unit);
- let divisor = from_size / to_size;
- match to_unit {
- TimeUnit::Second => {
- let values = unary::<_, _, Time32SecondType>(&time_array, |x| {
- (x as i64 / divisor) as i32
- });
- Ok(Arc::new(values) as ArrayRef)
- }
- TimeUnit::Millisecond => {
- let values = unary::<_, _, Time32MillisecondType>(&time_array, |x| {
- (x as i64 / divisor) as i32
- });
- Ok(Arc::new(values) as ArrayRef)
- }
- _ => unreachable!("array type not supported"),
- }
- }
- (Timestamp(_, _), Int64) => cast_array_data::<Int64Type>(array, to_type.clone()),
- (Int64, Timestamp(to_unit, _)) => {
- use TimeUnit::*;
- match to_unit {
- Second => cast_array_data::<TimestampSecondType>(array, to_type.clone()),
- Millisecond => {
- cast_array_data::<TimestampMillisecondType>(array, to_type.clone())
- }
- Microsecond => {
- cast_array_data::<TimestampMicrosecondType>(array, to_type.clone())
- }
- Nanosecond => {
- cast_array_data::<TimestampNanosecondType>(array, to_type.clone())
- }
- }
- }
- (Timestamp(from_unit, _), Timestamp(to_unit, _)) => {
- let time_array = Int64Array::from(array.data().clone());
- let from_size = time_unit_multiple(&from_unit);
- let to_size = time_unit_multiple(&to_unit);
- // we either divide or multiply, depending on size of each unit
- // units are never the same when the types are the same
- let converted = if from_size >= to_size {
- divide(
- &time_array,
- &Int64Array::from(vec![from_size / to_size; array.len()]),
- )?
- } else {
- multiply(
- &time_array,
- &Int64Array::from(vec![to_size / from_size; array.len()]),
- )?
- };
- let array_ref = Arc::new(converted) as ArrayRef;
- use TimeUnit::*;
- match to_unit {
- Second => {
- cast_array_data::<TimestampSecondType>(&array_ref, to_type.clone())
- }
- Millisecond => cast_array_data::<TimestampMillisecondType>(
- &array_ref,
- to_type.clone(),
- ),
- Microsecond => cast_array_data::<TimestampMicrosecondType>(
- &array_ref,
- to_type.clone(),
- ),
- Nanosecond => cast_array_data::<TimestampNanosecondType>(
- &array_ref,
- to_type.clone(),
- ),
- }
- }
- (Timestamp(from_unit, _), Date32) => {
- let time_array = Int64Array::from(array.data().clone());
- let from_size = time_unit_multiple(&from_unit) * SECONDS_IN_DAY;
- let mut b = Date32Builder::new(array.len());
- for i in 0..array.len() {
- if array.is_null(i) {
- b.append_null()?;
- } else {
- b.append_value((time_array.value(i) / from_size) as i32)?;
- }
- }
-
- Ok(Arc::new(b.finish()) as ArrayRef)
- }
- (Timestamp(from_unit, _), Date64) => {
- let from_size = time_unit_multiple(&from_unit);
- let to_size = MILLISECONDS;
-
- // Scale time_array by (to_size / from_size) using a
- // single integer operation, but need to avoid integer
- // math rounding down to zero
-
- match to_size.cmp(&from_size) {
- std::cmp::Ordering::Less => {
- let time_array = Date64Array::from(array.data().clone());
- Ok(Arc::new(divide(
- &time_array,
- &Date64Array::from(vec![from_size / to_size; array.len()]),
- )?) as ArrayRef)
- }
- std::cmp::Ordering::Equal => {
- cast_array_data::<Date64Type>(array, to_type.clone())
- }
- std::cmp::Ordering::Greater => {
- let time_array = Date64Array::from(array.data().clone());
- Ok(Arc::new(multiply(
- &time_array,
- &Date64Array::from(vec![to_size / from_size; array.len()]),
- )?) as ArrayRef)
- }
- }
- }
- // date64 to timestamp might not make sense,
- (Int64, Duration(to_unit)) => {
- use TimeUnit::*;
- match to_unit {
- Second => cast_array_data::<DurationSecondType>(array, to_type.clone()),
- Millisecond => {
- cast_array_data::<DurationMillisecondType>(array, to_type.clone())
- }
- Microsecond => {
- cast_array_data::<DurationMicrosecondType>(array, to_type.clone())
- }
- Nanosecond => {
- cast_array_data::<DurationNanosecondType>(array, to_type.clone())
- }
- }
- }
-
- // null to primitive/flat types
- (Null, Int32) => Ok(Arc::new(Int32Array::from(vec![None; array.len()]))),
-
- (_, _) => Err(ArrowError::CastError(format!(
- "Casting from {:?} to {:?} not supported",
- from_type, to_type,
- ))),
- }
-}
-
-/// Get the time unit as a multiple of a second
-const fn time_unit_multiple(unit: &TimeUnit) -> i64 {
- match unit {
- TimeUnit::Second => 1,
- TimeUnit::Millisecond => MILLISECONDS,
- TimeUnit::Microsecond => MICROSECONDS,
- TimeUnit::Nanosecond => NANOSECONDS,
- }
-}
-
-/// Number of seconds in a day
-const SECONDS_IN_DAY: i64 = 86_400;
-/// Number of milliseconds in a second
-const MILLISECONDS: i64 = 1_000;
-/// Number of microseconds in a second
-const MICROSECONDS: i64 = 1_000_000;
-/// Number of nanoseconds in a second
-const NANOSECONDS: i64 = 1_000_000_000;
-/// Number of milliseconds in a day
-const MILLISECONDS_IN_DAY: i64 = SECONDS_IN_DAY * MILLISECONDS;
-/// Number of days between 0001-01-01 and 1970-01-01
-const EPOCH_DAYS_FROM_CE: i32 = 719_163;
-
-/// Cast an array by changing its array_data type to the desired type
-///
-/// Arrays should have the same primitive data type, otherwise this should fail.
-/// We do not perform this check on primitive data types as we only use this
-/// function internally, where it is guaranteed to be infallible.
-fn cast_array_data<TO>(array: &ArrayRef, to_type: DataType) -> Result<ArrayRef>
-where
- TO: ArrowNumericType,
-{
- let data = ArrayData::new(
- to_type,
- array.len(),
- Some(array.null_count()),
- array.data().null_bitmap().clone().map(|bitmap| bitmap.bits),
- array.data().offset(),
- array.data().buffers().to_vec(),
- vec![],
- );
- Ok(Arc::new(PrimitiveArray::<TO>::from(data)) as ArrayRef)
-}
-
-/// Convert Array into a PrimitiveArray of type, and apply numeric cast
-fn cast_numeric_arrays<FROM, TO>(from: &ArrayRef) -> Result<ArrayRef>
-where
- FROM: ArrowNumericType,
- TO: ArrowNumericType,
- FROM::Native: num::NumCast,
- TO::Native: num::NumCast,
-{
- Ok(Arc::new(numeric_cast::<FROM, TO>(
- from.as_any()
- .downcast_ref::<PrimitiveArray<FROM>>()
- .unwrap(),
- )))
-}
-
-/// Natural cast between numeric types
-fn numeric_cast<T, R>(from: &PrimitiveArray<T>) -> PrimitiveArray<R>
-where
- T: ArrowNumericType,
- R: ArrowNumericType,
- T::Native: num::NumCast,
- R::Native: num::NumCast,
-{
- let iter = from
- .iter()
- .map(|v| v.and_then(num::cast::cast::<T::Native, R::Native>));
- // Soundness:
- // The iterator is trustedLen because it comes from an `PrimitiveArray`.
- unsafe { PrimitiveArray::<R>::from_trusted_len_iter(iter) }
-}
-
-/// Cast numeric types to Utf8
-fn cast_numeric_to_string<FROM, OffsetSize>(array: &ArrayRef) -> Result<ArrayRef>
-where
- FROM: ArrowNumericType,
- FROM::Native: lexical_core::ToLexical,
- OffsetSize: StringOffsetSizeTrait,
-{
- Ok(Arc::new(numeric_to_string_cast::<FROM, OffsetSize>(
- array
- .as_any()
- .downcast_ref::<PrimitiveArray<FROM>>()
- .unwrap(),
- )))
-}
-
-fn numeric_to_string_cast<T, OffsetSize>(
- from: &PrimitiveArray<T>,
-) -> GenericStringArray<OffsetSize>
-where
- T: ArrowPrimitiveType + ArrowNumericType,
- T::Native: lexical_core::ToLexical,
- OffsetSize: StringOffsetSizeTrait,
-{
- from.iter()
- .map(|maybe_value| maybe_value.map(lexical_to_string))
- .collect()
-}
-
-/// Cast numeric types to Utf8
-fn cast_string_to_numeric<T, Offset: StringOffsetSizeTrait>(
- from: &ArrayRef,
- cast_options: &CastOptions,
-) -> Result<ArrayRef>
-where
- T: ArrowNumericType,
- <T as ArrowPrimitiveType>::Native: lexical_core::FromLexical,
-{
- Ok(Arc::new(string_to_numeric_cast::<T, Offset>(
- from.as_any()
- .downcast_ref::<GenericStringArray<Offset>>()
- .unwrap(),
- cast_options,
- )?))
-}
-
-fn string_to_numeric_cast<T, Offset: StringOffsetSizeTrait>(
- from: &GenericStringArray<Offset>,
- cast_options: &CastOptions,
-) -> Result<PrimitiveArray<T>>
-where
- T: ArrowNumericType,
- <T as ArrowPrimitiveType>::Native: lexical_core::FromLexical,
-{
- if cast_options.safe {
- let iter = (0..from.len()).map(|i| {
- if from.is_null(i) {
- None
- } else {
- lexical_core::parse(from.value(i).as_bytes()).ok()
- }
- });
- // Benefit:
- // 20% performance improvement
- // Soundness:
- // The iterator is trustedLen because it comes from an `StringArray`.
- Ok(unsafe { PrimitiveArray::<T>::from_trusted_len_iter(iter) })
- } else {
- let vec = (0..from.len())
- .map(|i| {
- if from.is_null(i) {
- Ok(None)
- } else {
- let string = from.value(i);
- let result = lexical_core::parse(string.as_bytes());
- Some(result.map_err(|_| {
- ArrowError::CastError(format!(
- "Cannot cast string '{}' to value of {} type",
- string,
- std::any::type_name::<T>()
- ))
- }))
- .transpose()
- }
- })
- .collect::<Result<Vec<_>>>()?;
- // Benefit:
- // 20% performance improvement
- // Soundness:
- // The iterator is trustedLen because it comes from an `StringArray`.
- Ok(unsafe { PrimitiveArray::<T>::from_trusted_len_iter(vec.iter()) })
- }
-}
-
-/// Casts generic string arrays to Date32Array
-fn cast_string_to_date32<Offset: StringOffsetSizeTrait>(
- array: &dyn Array,
- cast_options: &CastOptions,
-) -> Result<ArrayRef> {
- use chrono::Datelike;
- let string_array = array
- .as_any()
- .downcast_ref::<GenericStringArray<Offset>>()
- .unwrap();
-
- let array = if cast_options.safe {
- let iter = (0..string_array.len()).map(|i| {
- if string_array.is_null(i) {
- None
- } else {
- string_array
- .value(i)
- .parse::<chrono::NaiveDate>()
- .map(|date| date.num_days_from_ce() - EPOCH_DAYS_FROM_CE)
- .ok()
- }
- });
-
- // Benefit:
- // 20% performance improvement
- // Soundness:
- // The iterator is trustedLen because it comes from an `StringArray`.
- unsafe { Date32Array::from_trusted_len_iter(iter) }
- } else {
- let vec = (0..string_array.len())
- .map(|i| {
- if string_array.is_null(i) {
- Ok(None)
- } else {
- let string = string_array
- .value(i);
-
- let result = string
- .parse::<chrono::NaiveDate>()
- .map(|date| date.num_days_from_ce() - EPOCH_DAYS_FROM_CE);
-
- Some(result.map_err(|_| {
- ArrowError::CastError(
- format!("Cannot cast string '{}' to value of arrow::datatypes::types::Date32Type type", string),
- )
- }))
- .transpose()
- }
- })
- .collect::<Result<Vec<Option<i32>>>>()?;
-
- // Benefit:
- // 20% performance improvement
- // Soundness:
- // The iterator is trustedLen because it comes from an `StringArray`.
- unsafe { Date32Array::from_trusted_len_iter(vec.iter()) }
- };
-
- Ok(Arc::new(array) as ArrayRef)
-}
-
-/// Casts generic string arrays to Date64Array
-fn cast_string_to_date64<Offset: StringOffsetSizeTrait>(
- array: &dyn Array,
- cast_options: &CastOptions,
-) -> Result<ArrayRef> {
- let string_array = array
- .as_any()
- .downcast_ref::<GenericStringArray<Offset>>()
- .unwrap();
-
- let array = if cast_options.safe {
- let iter = (0..string_array.len()).map(|i| {
- if string_array.is_null(i) {
- None
- } else {
- string_array
- .value(i)
- .parse::<chrono::NaiveDateTime>()
- .map(|datetime| datetime.timestamp_millis())
- .ok()
- }
- });
-
- // Benefit:
- // 20% performance improvement
- // Soundness:
- // The iterator is trustedLen because it comes from an `StringArray`.
- unsafe { Date64Array::from_trusted_len_iter(iter) }
- } else {
- let vec = (0..string_array.len())
- .map(|i| {
- if string_array.is_null(i) {
- Ok(None)
- } else {
- let string = string_array
- .value(i);
-
- let result = string
- .parse::<chrono::NaiveDateTime>()
- .map(|datetime| datetime.timestamp_millis());
-
- Some(result.map_err(|_| {
- ArrowError::CastError(
- format!("Cannot cast string '{}' to value of arrow::datatypes::types::Date64Type type", string),
- )
- }))
- .transpose()
- }
- })
- .collect::<Result<Vec<Option<i64>>>>()?;
-
- // Benefit:
- // 20% performance improvement
- // Soundness:
- // The iterator is trustedLen because it comes from an `StringArray`.
- unsafe { Date64Array::from_trusted_len_iter(vec.iter()) }
- };
-
- Ok(Arc::new(array) as ArrayRef)
-}
-
-/// Casts generic string arrays to TimeStampNanosecondArray
-fn cast_string_to_timestamp_ns<Offset: StringOffsetSizeTrait>(
- array: &dyn Array,
- cast_options: &CastOptions,
-) -> Result<ArrayRef> {
- let string_array = array
- .as_any()
- .downcast_ref::<GenericStringArray<Offset>>()
- .unwrap();
-
- let array = if cast_options.safe {
- let iter = (0..string_array.len()).map(|i| {
- if string_array.is_null(i) {
- None
- } else {
- string_to_timestamp_nanos(string_array.value(i)).ok()
- }
- });
- // Benefit:
- // 20% performance improvement
- // Soundness:
- // The iterator is trustedLen because it comes from an `StringArray`.
- unsafe { TimestampNanosecondArray::from_trusted_len_iter(iter) }
- } else {
- let vec = (0..string_array.len())
- .map(|i| {
- if string_array.is_null(i) {
- Ok(None)
- } else {
- let result = string_to_timestamp_nanos(string_array.value(i));
- Some(result).transpose()
- }
- })
- .collect::<Result<Vec<Option<i64>>>>()?;
-
- // Benefit:
- // 20% performance improvement
- // Soundness:
- // The iterator is trustedLen because it comes from an `StringArray`.
- unsafe { TimestampNanosecondArray::from_trusted_len_iter(vec.iter()) }
- };
-
- Ok(Arc::new(array) as ArrayRef)
-}
-
-/// Cast numeric types to Boolean
-///
-/// Any zero value returns `false` while non-zero returns `true`
-fn cast_numeric_to_bool<FROM>(from: &ArrayRef) -> Result<ArrayRef>
-where
- FROM: ArrowNumericType,
-{
- numeric_to_bool_cast::<FROM>(
- from.as_any()
- .downcast_ref::<PrimitiveArray<FROM>>()
- .unwrap(),
- )
- .map(|to| Arc::new(to) as ArrayRef)
-}
-
-fn numeric_to_bool_cast<T>(from: &PrimitiveArray<T>) -> Result<BooleanArray>
-where
- T: ArrowPrimitiveType + ArrowNumericType,
-{
- let mut b = BooleanBuilder::new(from.len());
-
- for i in 0..from.len() {
- if from.is_null(i) {
- b.append_null()?;
- } else if from.value(i) != T::default_value() {
- b.append_value(true)?;
- } else {
- b.append_value(false)?;
- }
- }
-
- Ok(b.finish())
-}
-
-/// Cast Boolean types to numeric
-///
-/// `false` returns 0 while `true` returns 1
-fn cast_bool_to_numeric<TO>(
- from: &ArrayRef,
- cast_options: &CastOptions,
-) -> Result<ArrayRef>
-where
- TO: ArrowNumericType,
- TO::Native: num::cast::NumCast,
-{
- Ok(Arc::new(bool_to_numeric_cast::<TO>(
- from.as_any().downcast_ref::<BooleanArray>().unwrap(),
- cast_options,
- )))
-}
-
-fn bool_to_numeric_cast<T>(
- from: &BooleanArray,
- _cast_options: &CastOptions,
-) -> PrimitiveArray<T>
-where
- T: ArrowNumericType,
- T::Native: num::NumCast,
-{
- let iter = (0..from.len()).map(|i| {
- if from.is_null(i) {
- None
- } else if from.value(i) {
- // a workaround to cast a primitive to T::Native, infallible
- num::cast::cast(1)
- } else {
- Some(T::default_value())
- }
- });
- // Benefit:
- // 20% performance improvement
- // Soundness:
- // The iterator is trustedLen because it comes from a Range
- unsafe { PrimitiveArray::<T>::from_trusted_len_iter(iter) }
-}
-
-/// Attempts to cast an `ArrayDictionary` with index type K into
-/// `to_type` for supported types.
-///
-/// K is the key type
-fn dictionary_cast<K: ArrowDictionaryKeyType>(
- array: &ArrayRef,
- to_type: &DataType,
- cast_options: &CastOptions,
-) -> Result<ArrayRef> {
- use DataType::*;
-
- match to_type {
- Dictionary(to_index_type, to_value_type) => {
- let dict_array = array
- .as_any()
- .downcast_ref::<DictionaryArray<K>>()
- .ok_or_else(|| {
- ArrowError::ComputeError(
- "Internal Error: Cannot cast dictionary to DictionaryArray of expected type".to_string(),
- )
- })?;
-
- let keys_array: ArrayRef =
- Arc::new(PrimitiveArray::<K>::from(dict_array.keys().data().clone()));
- let values_array = dict_array.values();
- let cast_keys = cast_with_options(&keys_array, to_index_type, &cast_options)?;
- let cast_values =
- cast_with_options(values_array, to_value_type, &cast_options)?;
-
- // Failure to cast keys (because they don't fit in the
- // target type) results in NULL values;
- if cast_keys.null_count() > keys_array.null_count() {
- return Err(ArrowError::ComputeError(format!(
- "Could not convert {} dictionary indexes from {:?} to {:?}",
- cast_keys.null_count() - keys_array.null_count(),
- keys_array.data_type(),
- to_index_type
- )));
- }
-
- // keys are data, child_data is values (dictionary)
- let data = ArrayData::new(
- to_type.clone(),
- cast_keys.len(),
- Some(cast_keys.null_count()),
- cast_keys
- .data()
- .null_bitmap()
- .clone()
- .map(|bitmap| bitmap.bits),
- cast_keys.data().offset(),
- cast_keys.data().buffers().to_vec(),
- vec![cast_values.data().clone()],
- );
-
- // create the appropriate array type
- let new_array: ArrayRef = match **to_index_type {
- Int8 => Arc::new(DictionaryArray::<Int8Type>::from(data)),
- Int16 => Arc::new(DictionaryArray::<Int16Type>::from(data)),
- Int32 => Arc::new(DictionaryArray::<Int32Type>::from(data)),
- Int64 => Arc::new(DictionaryArray::<Int64Type>::from(data)),
- UInt8 => Arc::new(DictionaryArray::<UInt8Type>::from(data)),
- UInt16 => Arc::new(DictionaryArray::<UInt16Type>::from(data)),
- UInt32 => Arc::new(DictionaryArray::<UInt32Type>::from(data)),
- UInt64 => Arc::new(DictionaryArray::<UInt64Type>::from(data)),
- _ => {
- return Err(ArrowError::CastError(format!(
- "Unsupported type {:?} for dictionary index",
- to_index_type
- )))
- }
- };
-
- Ok(new_array)
- }
- _ => unpack_dictionary::<K>(array, to_type, cast_options),
- }
-}
-
-// Unpack a dictionary where the keys are of type <K> into a flattened array of type to_type
-fn unpack_dictionary<K>(
- array: &ArrayRef,
- to_type: &DataType,
- cast_options: &CastOptions,
-) -> Result<ArrayRef>
-where
- K: ArrowDictionaryKeyType,
-{
- let dict_array = array
- .as_any()
- .downcast_ref::<DictionaryArray<K>>()
- .ok_or_else(|| {
- ArrowError::ComputeError(
- "Internal Error: Cannot cast dictionary to DictionaryArray of expected type".to_string(),
- )
- })?;
-
- // attempt to cast the dict values to the target type
- // use the take kernel to expand out the dictionary
- let cast_dict_values =
- cast_with_options(&dict_array.values(), to_type, cast_options)?;
-
- // Note take requires first casting the indices to u32
- let keys_array: ArrayRef =
- Arc::new(PrimitiveArray::<K>::from(dict_array.keys().data().clone()));
- let indicies = cast_with_options(&keys_array, &DataType::UInt32, cast_options)?;
- let u32_indicies =
- indicies
- .as_any()
- .downcast_ref::<UInt32Array>()
- .ok_or_else(|| {
- ArrowError::ComputeError(
- "Internal Error: Cannot cast dict indices to UInt32".to_string(),
- )
- })?;
-
- take(cast_dict_values.as_ref(), u32_indicies, None)
-}
-
-/// Attempts to encode an array into an `ArrayDictionary` with index
-/// type K and value (dictionary) type value_type
-///
-/// K is the key type
-fn cast_to_dictionary<K: ArrowDictionaryKeyType>(
- array: &ArrayRef,
- dict_value_type: &DataType,
- cast_options: &CastOptions,
-) -> Result<ArrayRef> {
- use DataType::*;
-
- match *dict_value_type {
- Int8 => pack_numeric_to_dictionary::<K, Int8Type>(
- array,
- dict_value_type,
- cast_options,
- ),
- Int16 => pack_numeric_to_dictionary::<K, Int16Type>(
- array,
- dict_value_type,
- cast_options,
- ),
- Int32 => pack_numeric_to_dictionary::<K, Int32Type>(
- array,
- dict_value_type,
- cast_options,
- ),
- Int64 => pack_numeric_to_dictionary::<K, Int64Type>(
- array,
- dict_value_type,
- cast_options,
- ),
- UInt8 => pack_numeric_to_dictionary::<K, UInt8Type>(
- array,
- dict_value_type,
- cast_options,
- ),
- UInt16 => pack_numeric_to_dictionary::<K, UInt16Type>(
- array,
- dict_value_type,
- cast_options,
- ),
- UInt32 => pack_numeric_to_dictionary::<K, UInt32Type>(
- array,
- dict_value_type,
- cast_options,
- ),
- UInt64 => pack_numeric_to_dictionary::<K, UInt64Type>(
- array,
- dict_value_type,
- cast_options,
- ),
- Utf8 => pack_string_to_dictionary::<K>(array, cast_options),
- _ => Err(ArrowError::CastError(format!(
- "Unsupported output type for dictionary packing: {:?}",
- dict_value_type
- ))),
- }
-}
-
-// Packs the data from the primitive array of type <V> to a
-// DictionaryArray with keys of type K and values of value_type V
-fn pack_numeric_to_dictionary<K, V>(
- array: &ArrayRef,
- dict_value_type: &DataType,
- cast_options: &CastOptions,
-) -> Result<ArrayRef>
-where
- K: ArrowDictionaryKeyType,
- V: ArrowNumericType,
-{
- // attempt to cast the source array values to the target value type (the dictionary values type)
- let cast_values = cast_with_options(array, &dict_value_type, cast_options)?;
- let values = cast_values
- .as_any()
- .downcast_ref::<PrimitiveArray<V>>()
- .unwrap();
-
- let keys_builder = PrimitiveBuilder::<K>::new(values.len());
- let values_builder = PrimitiveBuilder::<V>::new(values.len());
- let mut b = PrimitiveDictionaryBuilder::new(keys_builder, values_builder);
-
- // copy each element one at a time
- for i in 0..values.len() {
- if values.is_null(i) {
- b.append_null()?;
- } else {
- b.append(values.value(i))?;
- }
- }
- Ok(Arc::new(b.finish()))
-}
-
-// Packs the data as a StringDictionaryArray, if possible, with the
-// key types of K
-fn pack_string_to_dictionary<K>(
- array: &ArrayRef,
- cast_options: &CastOptions,
-) -> Result<ArrayRef>
-where
- K: ArrowDictionaryKeyType,
-{
- let cast_values = cast_with_options(array, &DataType::Utf8, cast_options)?;
- let values = cast_values.as_any().downcast_ref::<StringArray>().unwrap();
-
- let keys_builder = PrimitiveBuilder::<K>::new(values.len());
- let values_builder = StringBuilder::new(values.len());
- let mut b = StringDictionaryBuilder::new(keys_builder, values_builder);
-
- // copy each element one at a time
- for i in 0..values.len() {
- if values.is_null(i) {
- b.append_null()?;
- } else {
- b.append(values.value(i))?;
- }
- }
- Ok(Arc::new(b.finish()))
-}
-
-/// Helper function that takes a primitive array and casts to a (generic) list array.
-fn cast_primitive_to_list<OffsetSize: OffsetSizeTrait + NumCast>(
- array: &ArrayRef,
- to: &Field,
- to_type: &DataType,
- cast_options: &CastOptions,
-) -> Result<ArrayRef> {
- // cast primitive to list's primitive
- let cast_array = cast_with_options(array, to.data_type(), cast_options)?;
- // create offsets, where if array.len() = 2, we have [0,1,2]
- // Safety:
- // Length of range can be trusted.
- // Note: could not yet create a generic range in stable Rust.
- let offsets = unsafe {
- MutableBuffer::from_trusted_len_iter(
- (0..=array.len()).map(|i| OffsetSize::from(i).expect("integer")),
- )
- };
-
- let list_data = ArrayData::new(
- to_type.clone(),
- array.len(),
- Some(cast_array.null_count()),
- cast_array
- .data()
- .null_bitmap()
- .clone()
- .map(|bitmap| bitmap.bits),
- 0,
- vec![offsets.into()],
- vec![cast_array.data().clone()],
- );
- let list_array =
- Arc::new(GenericListArray::<OffsetSize>::from(list_data)) as ArrayRef;
-
- Ok(list_array)
-}
-
-/// Helper function that takes an Generic list container and casts the inner datatype.
-fn cast_list_inner<OffsetSize: OffsetSizeTrait>(
- array: &Arc<dyn Array>,
- to: &Field,
- to_type: &DataType,
- cast_options: &CastOptions,
-) -> Result<ArrayRef> {
- let data = array.data_ref();
- let underlying_array = make_array(data.child_data()[0].clone());
- let cast_array = cast_with_options(&underlying_array, to.data_type(), cast_options)?;
- let array_data = ArrayData::new(
- to_type.clone(),
- array.len(),
- Some(cast_array.null_count()),
- cast_array
- .data()
- .null_bitmap()
- .clone()
- .map(|bitmap| bitmap.bits),
- array.offset(),
- // reuse offset buffer
- data.buffers().to_vec(),
- vec![cast_array.data().clone()],
- );
- let list = GenericListArray::<OffsetSize>::from(array_data);
- Ok(Arc::new(list) as ArrayRef)
-}
-
-/// Helper function to cast from `Utf8` to `LargeUtf8` and vice versa. If the `LargeUtf8` is too large for
-/// a `Utf8` array it will return an Error.
-fn cast_str_container<OffsetSizeFrom, OffsetSizeTo>(array: &dyn Array) -> Result<ArrayRef>
-where
- OffsetSizeFrom: StringOffsetSizeTrait + ToPrimitive,
- OffsetSizeTo: StringOffsetSizeTrait + NumCast + ArrowNativeType,
-{
- let str_array = array
- .as_any()
- .downcast_ref::<GenericStringArray<OffsetSizeFrom>>()
- .unwrap();
- let list_data = array.data();
- let str_values_buf = str_array.value_data();
-
- let offsets = unsafe { list_data.buffers()[0].typed_data::<OffsetSizeFrom>() };
-
- let mut offset_builder = BufferBuilder::<OffsetSizeTo>::new(offsets.len());
- offsets.iter().try_for_each::<_, Result<_>>(|offset| {
- let offset = OffsetSizeTo::from(*offset).ok_or_else(|| {
- ArrowError::ComputeError(
- "large-utf8 array too large to cast to utf8-array".into(),
- )
- })?;
- offset_builder.append(offset);
- Ok(())
- })?;
-
- let offset_buffer = offset_builder.finish();
-
- let dtype = if matches!(std::mem::size_of::<OffsetSizeTo>(), 8) {
- DataType::LargeUtf8
- } else {
- DataType::Utf8
- };
-
- let mut builder = ArrayData::builder(dtype)
- .offset(array.offset())
- .len(array.len())
- .add_buffer(offset_buffer)
- .add_buffer(str_values_buf);
-
- if let Some(buf) = list_data.null_buffer() {
- builder = builder.null_bit_buffer(buf.clone())
- }
- let data = builder.build();
- Ok(Arc::new(GenericStringArray::<OffsetSizeTo>::from(data)))
-}
-
-/// Cast the container type of List/Largelist array but not the inner types.
-/// This function can leave the value data intact and only has to cast the offset dtypes.
-fn cast_list_container<OffsetSizeFrom, OffsetSizeTo>(
- array: &dyn Array,
- _cast_options: &CastOptions,
-) -> Result<ArrayRef>
-where
- OffsetSizeFrom: OffsetSizeTrait + ToPrimitive,
- OffsetSizeTo: OffsetSizeTrait + NumCast,
-{
- let data = array.data_ref();
- // the value data stored by the list
- let value_data = data.child_data()[0].clone();
-
- let out_dtype = match array.data_type() {
- DataType::List(value_type) => {
- assert_eq!(
- std::mem::size_of::<OffsetSizeFrom>(),
- std::mem::size_of::<i32>()
- );
- assert_eq!(
- std::mem::size_of::<OffsetSizeTo>(),
- std::mem::size_of::<i64>()
- );
- DataType::LargeList(value_type.clone())
- }
- DataType::LargeList(value_type) => {
- assert_eq!(
- std::mem::size_of::<OffsetSizeFrom>(),
- std::mem::size_of::<i64>()
- );
- assert_eq!(
- std::mem::size_of::<OffsetSizeTo>(),
- std::mem::size_of::<i32>()
- );
- if value_data.len() > i32::MAX as usize {
- return Err(ArrowError::ComputeError(
- "LargeList too large to cast to List".into(),
- ));
- }
- DataType::List(value_type.clone())
- }
- // implementation error
- _ => unreachable!(),
- };
-
- // Safety:
- // The first buffer is the offsets and they are aligned to OffSetSizeFrom: (i64 or i32)
- // Justification:
- // The safe variant data.buffer::<OffsetSizeFrom> take the offset into account and we
- // cannot create a list array with offsets starting at non zero.
- let offsets = unsafe { data.buffers()[0].as_slice().align_to::<OffsetSizeFrom>() }.1;
-
- let iter = offsets.iter().map(|idx| {
- let idx: OffsetSizeTo = NumCast::from(*idx).unwrap();
- idx
- });
-
- // SAFETY
- // A slice produces a trusted length iterator
- let offset_buffer = unsafe { Buffer::from_trusted_len_iter(iter) };
-
- // wrap up
- let mut builder = ArrayData::builder(out_dtype)
- .offset(array.offset())
- .len(array.len())
- .add_buffer(offset_buffer)
- .add_child_data(value_data);
-
- if let Some(buf) = data.null_buffer() {
- builder = builder.null_bit_buffer(buf.clone())
- }
- let data = builder.build();
- Ok(make_array(data))
-}
-
-#[cfg(test)]
-mod tests {
- use super::*;
- use crate::{buffer::Buffer, util::display::array_value_to_string};
-
- #[test]
- fn test_cast_i32_to_f64() {
- let a = Int32Array::from(vec![5, 6, 7, 8, 9]);
- let array = Arc::new(a) as ArrayRef;
- let b = cast(&array, &DataType::Float64).unwrap();
- let c = b.as_any().downcast_ref::<Float64Array>().unwrap();
- assert!(5.0 - c.value(0) < f64::EPSILON);
- assert!(6.0 - c.value(1) < f64::EPSILON);
- assert!(7.0 - c.value(2) < f64::EPSILON);
- assert!(8.0 - c.value(3) < f64::EPSILON);
- assert!(9.0 - c.value(4) < f64::EPSILON);
- }
-
- #[test]
- fn test_cast_i32_to_u8() {
- let a = Int32Array::from(vec![-5, 6, -7, 8, 100000000]);
- let array = Arc::new(a) as ArrayRef;
- let b = cast(&array, &DataType::UInt8).unwrap();
- let c = b.as_any().downcast_ref::<UInt8Array>().unwrap();
- assert_eq!(false, c.is_valid(0));
- assert_eq!(6, c.value(1));
- assert_eq!(false, c.is_valid(2));
- assert_eq!(8, c.value(3));
- // overflows return None
- assert_eq!(false, c.is_valid(4));
- }
-
- #[test]
- fn test_cast_i32_to_u8_sliced() {
- let a = Int32Array::from(vec![-5, 6, -7, 8, 100000000]);
- let array = Arc::new(a) as ArrayRef;
- assert_eq!(0, array.offset());
- let array = array.slice(2, 3);
- assert_eq!(2, array.offset());
- let b = cast(&array, &DataType::UInt8).unwrap();
- assert_eq!(3, b.len());
- assert_eq!(0, b.offset());
- let c = b.as_any().downcast_ref::<UInt8Array>().unwrap();
- assert_eq!(false, c.is_valid(0));
- assert_eq!(8, c.value(1));
- // overflows return None
- assert_eq!(false, c.is_valid(2));
- }
-
- #[test]
- fn test_cast_i32_to_i32() {
- let a = Int32Array::from(vec![5, 6, 7, 8, 9]);
- let array = Arc::new(a) as ArrayRef;
- let b = cast(&array, &DataType::Int32).unwrap();
- let c = b.as_any().downcast_ref::<Int32Array>().unwrap();
- assert_eq!(5, c.value(0));
- assert_eq!(6, c.value(1));
- assert_eq!(7, c.value(2));
- assert_eq!(8, c.value(3));
- assert_eq!(9, c.value(4));
- }
-
- #[test]
- fn test_cast_i32_to_list_i32() {
- let a = Int32Array::from(vec![5, 6, 7, 8, 9]);
- let array = Arc::new(a) as ArrayRef;
- let b = cast(
- &array,
- &DataType::List(Box::new(Field::new("item", DataType::Int32, true))),
- )
- .unwrap();
- assert_eq!(5, b.len());
- let arr = b.as_any().downcast_ref::<ListArray>().unwrap();
- assert_eq!(&[0, 1, 2, 3, 4, 5], arr.value_offsets());
- assert_eq!(1, arr.value_length(0));
- assert_eq!(1, arr.value_length(1));
- assert_eq!(1, arr.value_length(2));
- assert_eq!(1, arr.value_length(3));
- assert_eq!(1, arr.value_length(4));
- let values = arr.values();
- let c = values.as_any().downcast_ref::<Int32Array>().unwrap();
- assert_eq!(5, c.value(0));
- assert_eq!(6, c.value(1));
- assert_eq!(7, c.value(2));
- assert_eq!(8, c.value(3));
- assert_eq!(9, c.value(4));
- }
-
- #[test]
- fn test_cast_i32_to_list_i32_nullable() {
- let a = Int32Array::from(vec![Some(5), None, Some(7), Some(8), Some(9)]);
- let array = Arc::new(a) as ArrayRef;
- let b = cast(
- &array,
- &DataType::List(Box::new(Field::new("item", DataType::Int32, true))),
- )
- .unwrap();
- assert_eq!(5, b.len());
- assert_eq!(1, b.null_count());
- let arr = b.as_any().downcast_ref::<ListArray>().unwrap();
- assert_eq!(&[0, 1, 2, 3, 4, 5], arr.value_offsets());
- assert_eq!(1, arr.value_length(0));
- assert_eq!(1, arr.value_length(1));
- assert_eq!(1, arr.value_length(2));
- assert_eq!(1, arr.value_length(3));
- assert_eq!(1, arr.value_length(4));
- let values = arr.values();
- let c = values.as_any().downcast_ref::<Int32Array>().unwrap();
- assert_eq!(1, c.null_count());
- assert_eq!(5, c.value(0));
- assert_eq!(false, c.is_valid(1));
- assert_eq!(7, c.value(2));
- assert_eq!(8, c.value(3));
- assert_eq!(9, c.value(4));
- }
-
- #[test]
- fn test_cast_i32_to_list_f64_nullable_sliced() {
- let a = Int32Array::from(vec![Some(5), None, Some(7), Some(8), None, Some(10)]);
- let array = Arc::new(a) as ArrayRef;
- let array = array.slice(2, 4);
- let b = cast(
- &array,
- &DataType::List(Box::new(Field::new("item", DataType::Float64, true))),
- )
- .unwrap();
- assert_eq!(4, b.len());
- assert_eq!(1, b.null_count());
- let arr = b.as_any().downcast_ref::<ListArray>().unwrap();
- assert_eq!(&[0, 1, 2, 3, 4], arr.value_offsets());
- assert_eq!(1, arr.value_length(0));
- assert_eq!(1, arr.value_length(1));
- assert_eq!(1, arr.value_length(2));
- assert_eq!(1, arr.value_length(3));
- let values = arr.values();
- let c = values.as_any().downcast_ref::<Float64Array>().unwrap();
- assert_eq!(1, c.null_count());
- assert!(7.0 - c.value(0) < f64::EPSILON);
- assert!(8.0 - c.value(1) < f64::EPSILON);
- assert_eq!(false, c.is_valid(2));
- assert!(10.0 - c.value(3) < f64::EPSILON);
- }
-
- #[test]
- fn test_cast_utf8_to_i32() {
- let a = StringArray::from(vec!["5", "6", "seven", "8", "9.1"]);
- let array = Arc::new(a) as ArrayRef;
- let b = cast(&array, &DataType::Int32).unwrap();
- let c = b.as_any().downcast_ref::<Int32Array>().unwrap();
- assert_eq!(5, c.value(0));
- assert_eq!(6, c.value(1));
- assert_eq!(false, c.is_valid(2));
- assert_eq!(8, c.value(3));
- assert_eq!(false, c.is_valid(4));
- }
-
- #[test]
- fn test_cast_with_options_utf8_to_i32() {
- let a = StringArray::from(vec!["5", "6", "seven", "8", "9.1"]);
- let array = Arc::new(a) as ArrayRef;
- let result =
- cast_with_options(&array, &DataType::Int32, &CastOptions { safe: false });
- match result {
- Ok(_) => panic!("expected error"),
- Err(e) => {
- assert!(e.to_string().contains(
- "Cast error: Cannot cast string 'seven' to value of arrow::datatypes::types::Int32Type type"
- ))
- }
- }
- }
-
- #[test]
- fn test_cast_bool_to_i32() {
- let a = BooleanArray::from(vec![Some(true), Some(false), None]);
- let array = Arc::new(a) as ArrayRef;
- let b = cast(&array, &DataType::Int32).unwrap();
- let c = b.as_any().downcast_ref::<Int32Array>().unwrap();
- assert_eq!(1, c.value(0));
- assert_eq!(0, c.value(1));
- assert_eq!(false, c.is_valid(2));
- }
-
- #[test]
- fn test_cast_bool_to_f64() {
- let a = BooleanArray::from(vec![Some(true), Some(false), None]);
- let array = Arc::new(a) as ArrayRef;
- let b = cast(&array, &DataType::Float64).unwrap();
- let c = b.as_any().downcast_ref::<Float64Array>().unwrap();
- assert!(1.0 - c.value(0) < f64::EPSILON);
- assert!(0.0 - c.value(1) < f64::EPSILON);
- assert_eq!(false, c.is_valid(2));
- }
-
- #[test]
- #[should_panic(
- expected = "Casting from Int32 to Timestamp(Microsecond, None) not supported"
- )]
- fn test_cast_int32_to_timestamp() {
- let a = Int32Array::from(vec![Some(2), Some(10), None]);
- let array = Arc::new(a) as ArrayRef;
- cast(&array, &DataType::Timestamp(TimeUnit::Microsecond, None)).unwrap();
- }
-
- #[test]
- fn test_cast_list_i32_to_list_u16() {
- // Construct a value array
- let value_data = Int32Array::from(vec![0, 0, 0, -1, -2, -1, 2, 100000000])
- .data()
- .clone();
-
- let value_offsets = Buffer::from_slice_ref(&[0, 3, 6, 8]);
-
- // Construct a list array from the above two
- let list_data_type =
- DataType::List(Box::new(Field::new("item", DataType::Int32, true)));
- let list_data = ArrayData::builder(list_data_type)
- .len(3)
- .add_buffer(value_offsets)
- .add_child_data(value_data)
- .build();
- let list_array = Arc::new(ListArray::from(list_data)) as ArrayRef;
-
- let cast_array = cast(
- &list_array,
- &DataType::List(Box::new(Field::new("item", DataType::UInt16, true))),
- )
- .unwrap();
- // 3 negative values should get lost when casting to unsigned,
- // 1 value should overflow
- assert_eq!(4, cast_array.null_count());
- // offsets should be the same
- assert_eq!(
- list_array.data().buffers().to_vec(),
- cast_array.data().buffers().to_vec()
- );
- let array = cast_array
- .as_ref()
- .as_any()
- .downcast_ref::<ListArray>()
- .unwrap();
- assert_eq!(DataType::UInt16, array.value_type());
- assert_eq!(4, array.values().null_count());
- assert_eq!(3, array.value_length(0));
- assert_eq!(3, array.value_length(1));
- assert_eq!(2, array.value_length(2));
- let values = array.values();
- let u16arr = values.as_any().downcast_ref::<UInt16Array>().unwrap();
- assert_eq!(8, u16arr.len());
- assert_eq!(4, u16arr.null_count());
-
- assert_eq!(0, u16arr.value(0));
- assert_eq!(0, u16arr.value(1));
- assert_eq!(0, u16arr.value(2));
- assert_eq!(false, u16arr.is_valid(3));
- assert_eq!(false, u16arr.is_valid(4));
- assert_eq!(false, u16arr.is_valid(5));
- assert_eq!(2, u16arr.value(6));
- assert_eq!(false, u16arr.is_valid(7));
- }
-
- #[test]
- #[should_panic(
- expected = "Casting from Int32 to Timestamp(Microsecond, None) not supported"
- )]
- fn test_cast_list_i32_to_list_timestamp() {
- // Construct a value array
- let value_data = Int32Array::from(vec![0, 0, 0, -1, -2, -1, 2, 8, 100000000])
- .data()
- .clone();
-
- let value_offsets = Buffer::from_slice_ref(&[0, 3, 6, 9]);
-
- // Construct a list array from the above two
- let list_data_type =
- DataType::List(Box::new(Field::new("item", DataType::Int32, true)));
- let list_data = ArrayData::builder(list_data_type)
- .len(3)
- .add_buffer(value_offsets)
- .add_child_data(value_data)
- .build();
- let list_array = Arc::new(ListArray::from(list_data)) as ArrayRef;
-
- cast(
- &list_array,
- &DataType::List(Box::new(Field::new(
- "item",
- DataType::Timestamp(TimeUnit::Microsecond, None),
- true,
- ))),
- )
- .unwrap();
- }
-
- #[test]
- fn test_cast_date32_to_date64() {
- let a = Date32Array::from(vec![10000, 17890]);
- let array = Arc::new(a) as ArrayRef;
- let b = cast(&array, &DataType::Date64).unwrap();
- let c = b.as_any().downcast_ref::<Date64Array>().unwrap();
- assert_eq!(864000000000, c.value(0));
- assert_eq!(1545696000000, c.value(1));
- }
-
- #[test]
- fn test_cast_date64_to_date32() {
- let a = Date64Array::from(vec![Some(864000000005), Some(1545696000001), None]);
- let array = Arc::new(a) as ArrayRef;
- let b = cast(&array, &DataType::Date32).unwrap();
- let c = b.as_any().downcast_ref::<Date32Array>().unwrap();
- assert_eq!(10000, c.value(0));
- assert_eq!(17890, c.value(1));
- assert!(c.is_null(2));
- }
-
- #[test]
- fn test_cast_string_to_timestamp() {
- let a1 = Arc::new(StringArray::from(vec![
- Some("2020-09-08T12:00:00+00:00"),
- Some("Not a valid date"),
- None,
- ])) as ArrayRef;
- let a2 = Arc::new(LargeStringArray::from(vec![
- Some("2020-09-08T12:00:00+00:00"),
- Some("Not a valid date"),
- None,
- ])) as ArrayRef;
- for array in &[a1, a2] {
- let b =
- cast(array, &DataType::Timestamp(TimeUnit::Nanosecond, None)).unwrap();
- let c = b
- .as_any()
- .downcast_ref::<TimestampNanosecondArray>()
- .unwrap();
- assert_eq!(1599566400000000000, c.value(0));
- assert!(c.is_null(1));
- assert!(c.is_null(2));
- }
- }
-
- #[test]
- fn test_cast_date32_to_int32() {
- let a = Date32Array::from(vec![10000, 17890]);
- let array = Arc::new(a) as ArrayRef;
- let b = cast(&array, &DataType::Int32).unwrap();
- let c = b.as_any().downcast_ref::<Int32Array>().unwrap();
- assert_eq!(10000, c.value(0));
- assert_eq!(17890, c.value(1));
- }
-
- #[test]
- fn test_cast_int32_to_date32() {
- let a = Int32Array::from(vec![10000, 17890]);
- let array = Arc::new(a) as ArrayRef;
- let b = cast(&array, &DataType::Date32).unwrap();
- let c = b.as_any().downcast_ref::<Date32Array>().unwrap();
- assert_eq!(10000, c.value(0));
- assert_eq!(17890, c.value(1));
- }
-
- #[test]
- fn test_cast_timestamp_to_date32() {
- let a = TimestampMillisecondArray::from_opt_vec(
- vec![Some(864000000005), Some(1545696000001), None],
- Some(String::from("UTC")),
- );
- let array = Arc::new(a) as ArrayRef;
- let b = cast(&array, &DataType::Date32).unwrap();
- let c = b.as_any().downcast_ref::<Date32Array>().unwrap();
- assert_eq!(10000, c.value(0));
- assert_eq!(17890, c.value(1));
- assert!(c.is_null(2));
- }
-
- #[test]
- fn test_cast_timestamp_to_date64() {
- let a = TimestampMillisecondArray::from_opt_vec(
- vec![Some(864000000005), Some(1545696000001), None],
- None,
- );
- let array = Arc::new(a) as ArrayRef;
- let b = cast(&array, &DataType::Date64).unwrap();
- let c = b.as_any().downcast_ref::<Date64Array>().unwrap();
- assert_eq!(864000000005, c.value(0));
- assert_eq!(1545696000001, c.value(1));
- assert!(c.is_null(2));
- }
-
- #[test]
- fn test_cast_timestamp_to_i64() {
- let a = TimestampMillisecondArray::from_opt_vec(
- vec![Some(864000000005), Some(1545696000001), None],
- Some("UTC".to_string()),
- );
- let array = Arc::new(a) as ArrayRef;
- let b = cast(&array, &DataType::Int64).unwrap();
- let c = b.as_any().downcast_ref::<Int64Array>().unwrap();
- assert_eq!(&DataType::Int64, c.data_type());
- assert_eq!(864000000005, c.value(0));
- assert_eq!(1545696000001, c.value(1));
- assert!(c.is_null(2));
- }
-
- #[test]
- fn test_cast_between_timestamps() {
- let a = TimestampMillisecondArray::from_opt_vec(
- vec![Some(864000003005), Some(1545696002001), None],
- None,
- );
- let array = Arc::new(a) as ArrayRef;
- let b = cast(&array, &DataType::Timestamp(TimeUnit::Second, None)).unwrap();
- let c = b.as_any().downcast_ref::<TimestampSecondArray>().unwrap();
- assert_eq!(864000003, c.value(0));
- assert_eq!(1545696002, c.value(1));
- assert!(c.is_null(2));
- }
-
- #[test]
- fn test_cast_to_strings() {
- let a = Arc::new(Int32Array::from(vec![1, 2, 3])) as ArrayRef;
- let out = cast(&a, &DataType::Utf8).unwrap();
- let out = out
- .as_any()
- .downcast_ref::<StringArray>()
- .unwrap()
- .into_iter()
- .collect::<Vec<_>>();
- assert_eq!(out, vec![Some("1"), Some("2"), Some("3")]);
- let out = cast(&a, &DataType::LargeUtf8).unwrap();
- let out = out
- .as_any()
- .downcast_ref::<LargeStringArray>()
- .unwrap()
- .into_iter()
- .collect::<Vec<_>>();
- assert_eq!(out, vec![Some("1"), Some("2"), Some("3")]);
- }
-
- #[test]
- fn test_str_to_str_casts() {
- for data in vec![
- vec![Some("foo"), Some("bar"), Some("ham")],
- vec![Some("foo"), None, Some("bar")],
- ] {
- let a = Arc::new(LargeStringArray::from(data.clone())) as ArrayRef;
- let to = cast(&a, &DataType::Utf8).unwrap();
- let expect = a
- .as_any()
- .downcast_ref::<LargeStringArray>()
- .unwrap()
- .into_iter()
- .collect::<Vec<_>>();
- let out = to
- .as_any()
- .downcast_ref::<StringArray>()
- .unwrap()
- .into_iter()
- .collect::<Vec<_>>();
- assert_eq!(expect, out);
-
- let a = Arc::new(StringArray::from(data)) as ArrayRef;
- let to = cast(&a, &DataType::LargeUtf8).unwrap();
- let expect = a
- .as_any()
- .downcast_ref::<StringArray>()
- .unwrap()
- .into_iter()
- .collect::<Vec<_>>();
- let out = to
- .as_any()
- .downcast_ref::<LargeStringArray>()
- .unwrap()
- .into_iter()
- .collect::<Vec<_>>();
- assert_eq!(expect, out);
- }
- }
-
- #[test]
- fn test_cast_from_f64() {
- let f64_values: Vec<f64> = vec![
- std::i64::MIN as f64,
- std::i32::MIN as f64,
- std::i16::MIN as f64,
- std::i8::MIN as f64,
- 0_f64,
- std::u8::MAX as f64,
- std::u16::MAX as f64,
- std::u32::MAX as f64,
- std::u64::MAX as f64,
- ];
- let f64_array: ArrayRef = Arc::new(Float64Array::from(f64_values));
-
- let f64_expected = vec![
- "-9223372036854776000.0",
- "-2147483648.0",
- "-32768.0",
- "-128.0",
- "0.0",
- "255.0",
- "65535.0",
- "4294967295.0",
- "18446744073709552000.0",
- ];
- assert_eq!(
- f64_expected,
- get_cast_values::<Float64Type>(&f64_array, &DataType::Float64)
- );
-
- let f32_expected = vec![
- "-9223372000000000000.0",
- "-2147483600.0",
- "-32768.0",
- "-128.0",
- "0.0",
- "255.0",
- "65535.0",
- "4294967300.0",
- "18446744000000000000.0",
- ];
- assert_eq!(
- f32_expected,
- get_cast_values::<Float32Type>(&f64_array, &DataType::Float32)
- );
-
- let i64_expected = vec![
- "-9223372036854775808",
- "-2147483648",
- "-32768",
- "-128",
- "0",
- "255",
- "65535",
- "4294967295",
- "null",
- ];
- assert_eq!(
- i64_expected,
- get_cast_values::<Int64Type>(&f64_array, &DataType::Int64)
- );
-
- let i32_expected = vec![
- "null",
- "-2147483648",
- "-32768",
- "-128",
- "0",
- "255",
- "65535",
- "null",
- "null",
- ];
- assert_eq!(
- i32_expected,
- get_cast_values::<Int32Type>(&f64_array, &DataType::Int32)
- );
-
- let i16_expected = vec![
- "null", "null", "-32768", "-128", "0", "255", "null", "null", "null",
- ];
- assert_eq!(
- i16_expected,
- get_cast_values::<Int16Type>(&f64_array, &DataType::Int16)
- );
-
- let i8_expected = vec![
- "null", "null", "null", "-128", "0", "null", "null", "null", "null",
- ];
- assert_eq!(
- i8_expected,
- get_cast_values::<Int8Type>(&f64_array, &DataType::Int8)
- );
-
- let u64_expected = vec![
- "null",
- "null",
- "null",
- "null",
- "0",
- "255",
- "65535",
- "4294967295",
- "null",
- ];
- assert_eq!(
- u64_expected,
- get_cast_values::<UInt64Type>(&f64_array, &DataType::UInt64)
- );
-
- let u32_expected = vec![
- "null",
- "null",
- "null",
- "null",
- "0",
- "255",
- "65535",
- "4294967295",
- "null",
- ];
- assert_eq!(
- u32_expected,
- get_cast_values::<UInt32Type>(&f64_array, &DataType::UInt32)
- );
-
- let u16_expected = vec![
- "null", "null", "null", "null", "0", "255", "65535", "null", "null",
- ];
- assert_eq!(
- u16_expected,
- get_cast_values::<UInt16Type>(&f64_array, &DataType::UInt16)
- );
-
- let u8_expected = vec![
- "null", "null", "null", "null", "0", "255", "null", "null", "null",
- ];
- assert_eq!(
- u8_expected,
- get_cast_values::<UInt8Type>(&f64_array, &DataType::UInt8)
- );
- }
-
- #[test]
- fn test_cast_from_f32() {
- let f32_values: Vec<f32> = vec![
- std::i32::MIN as f32,
- std::i32::MIN as f32,
- std::i16::MIN as f32,
- std::i8::MIN as f32,
- 0_f32,
- std::u8::MAX as f32,
- std::u16::MAX as f32,
- std::u32::MAX as f32,
- std::u32::MAX as f32,
- ];
- let f32_array: ArrayRef = Arc::new(Float32Array::from(f32_values));
-
- let f64_expected = vec![
- "-2147483648.0",
- "-2147483648.0",
- "-32768.0",
- "-128.0",
- "0.0",
- "255.0",
- "65535.0",
- "4294967296.0",
- "4294967296.0",
- ];
- assert_eq!(
- f64_expected,
- get_cast_values::<Float64Type>(&f32_array, &DataType::Float64)
- );
-
- let f32_expected = vec![
- "-2147483600.0",
- "-2147483600.0",
- "-32768.0",
- "-128.0",
- "0.0",
- "255.0",
- "65535.0",
- "4294967300.0",
- "4294967300.0",
- ];
- assert_eq!(
- f32_expected,
- get_cast_values::<Float32Type>(&f32_array, &DataType::Float32)
- );
-
- let i64_expected = vec![
- "-2147483648",
- "-2147483648",
- "-32768",
- "-128",
- "0",
- "255",
- "65535",
- "4294967296",
- "4294967296",
- ];
- assert_eq!(
- i64_expected,
- get_cast_values::<Int64Type>(&f32_array, &DataType::Int64)
- );
-
- let i32_expected = vec![
- "-2147483648",
- "-2147483648",
- "-32768",
- "-128",
- "0",
- "255",
- "65535",
- "null",
- "null",
- ];
- assert_eq!(
- i32_expected,
- get_cast_values::<Int32Type>(&f32_array, &DataType::Int32)
- );
-
- let i16_expected = vec![
- "null", "null", "-32768", "-128", "0", "255", "null", "null", "null",
- ];
- assert_eq!(
- i16_expected,
- get_cast_values::<Int16Type>(&f32_array, &DataType::Int16)
- );
-
- let i8_expected = vec![
- "null", "null", "null", "-128", "0", "null", "null", "null", "null",
- ];
- assert_eq!(
- i8_expected,
- get_cast_values::<Int8Type>(&f32_array, &DataType::Int8)
- );
-
- let u64_expected = vec![
- "null",
- "null",
- "null",
- "null",
- "0",
- "255",
- "65535",
- "4294967296",
- "4294967296",
- ];
- assert_eq!(
- u64_expected,
- get_cast_values::<UInt64Type>(&f32_array, &DataType::UInt64)
- );
-
- let u32_expected = vec![
- "null", "null", "null", "null", "0", "255", "65535", "null", "null",
- ];
- assert_eq!(
- u32_expected,
- get_cast_values::<UInt32Type>(&f32_array, &DataType::UInt32)
- );
-
- let u16_expected = vec![
- "null", "null", "null", "null", "0", "255", "65535", "null", "null",
- ];
- assert_eq!(
- u16_expected,
- get_cast_values::<UInt16Type>(&f32_array, &DataType::UInt16)
- );
-
- let u8_expected = vec![
- "null", "null", "null", "null", "0", "255", "null", "null", "null",
- ];
- assert_eq!(
- u8_expected,
- get_cast_values::<UInt8Type>(&f32_array, &DataType::UInt8)
- );
- }
-
- #[test]
- fn test_cast_from_uint64() {
- let u64_values: Vec<u64> = vec![
- 0,
- std::u8::MAX as u64,
- std::u16::MAX as u64,
- std::u32::MAX as u64,
- std::u64::MAX,
- ];
- let u64_array: ArrayRef = Arc::new(UInt64Array::from(u64_values));
-
- let f64_expected = vec![
- "0.0",
- "255.0",
- "65535.0",
- "4294967295.0",
- "18446744073709552000.0",
- ];
- assert_eq!(
- f64_expected,
- get_cast_values::<Float64Type>(&u64_array, &DataType::Float64)
- );
-
- let f32_expected = vec![
- "0.0",
- "255.0",
- "65535.0",
- "4294967300.0",
- "18446744000000000000.0",
- ];
- assert_eq!(
- f32_expected,
- get_cast_values::<Float32Type>(&u64_array, &DataType::Float32)
- );
-
- let i64_expected = vec!["0", "255", "65535", "4294967295", "null"];
- assert_eq!(
- i64_expected,
- get_cast_values::<Int64Type>(&u64_array, &DataType::Int64)
- );
-
- let i32_expected = vec!["0", "255", "65535", "null", "null"];
- assert_eq!(
- i32_expected,
- get_cast_values::<Int32Type>(&u64_array, &DataType::Int32)
- );
-
- let i16_expected = vec!["0", "255", "null", "null", "null"];
- assert_eq!(
- i16_expected,
- get_cast_values::<Int16Type>(&u64_array, &DataType::Int16)
- );
-
- let i8_expected = vec!["0", "null", "null", "null", "null"];
- assert_eq!(
- i8_expected,
- get_cast_values::<Int8Type>(&u64_array, &DataType::Int8)
- );
-
- let u64_expected =
- vec!["0", "255", "65535", "4294967295", "18446744073709551615"];
- assert_eq!(
- u64_expected,
- get_cast_values::<UInt64Type>(&u64_array, &DataType::UInt64)
- );
-
- let u32_expected = vec!["0", "255", "65535", "4294967295", "null"];
- assert_eq!(
- u32_expected,
- get_cast_values::<UInt32Type>(&u64_array, &DataType::UInt32)
- );
-
- let u16_expected = vec!["0", "255", "65535", "null", "null"];
- assert_eq!(
- u16_expected,
- get_cast_values::<UInt16Type>(&u64_array, &DataType::UInt16)
- );
-
- let u8_expected = vec!["0", "255", "null", "null", "null"];
- assert_eq!(
- u8_expected,
- get_cast_values::<UInt8Type>(&u64_array, &DataType::UInt8)
- );
- }
-
- #[test]
- fn test_cast_from_uint32() {
- let u32_values: Vec<u32> = vec![
- 0,
- std::u8::MAX as u32,
- std::u16::MAX as u32,
- std::u32::MAX as u32,
- ];
- let u32_array: ArrayRef = Arc::new(UInt32Array::from(u32_values));
-
- let f64_expected = vec!["0.0", "255.0", "65535.0", "4294967295.0"];
- assert_eq!(
- f64_expected,
- get_cast_values::<Float64Type>(&u32_array, &DataType::Float64)
- );
-
- let f32_expected = vec!["0.0", "255.0", "65535.0", "4294967300.0"];
- assert_eq!(
- f32_expected,
- get_cast_values::<Float32Type>(&u32_array, &DataType::Float32)
- );
-
- let i64_expected = vec!["0", "255", "65535", "4294967295"];
- assert_eq!(
- i64_expected,
- get_cast_values::<Int64Type>(&u32_array, &DataType::Int64)
- );
-
- let i32_expected = vec!["0", "255", "65535", "null"];
- assert_eq!(
- i32_expected,
- get_cast_values::<Int32Type>(&u32_array, &DataType::Int32)
- );
-
- let i16_expected = vec!["0", "255", "null", "null"];
- assert_eq!(
- i16_expected,
- get_cast_values::<Int16Type>(&u32_array, &DataType::Int16)
- );
-
- let i8_expected = vec!["0", "null", "null", "null"];
- assert_eq!(
- i8_expected,
- get_cast_values::<Int8Type>(&u32_array, &DataType::Int8)
- );
-
- let u64_expected = vec!["0", "255", "65535", "4294967295"];
- assert_eq!(
- u64_expected,
- get_cast_values::<UInt64Type>(&u32_array, &DataType::UInt64)
- );
-
- let u32_expected = vec!["0", "255", "65535", "4294967295"];
- assert_eq!(
- u32_expected,
- get_cast_values::<UInt32Type>(&u32_array, &DataType::UInt32)
- );
-
- let u16_expected = vec!["0", "255", "65535", "null"];
- assert_eq!(
- u16_expected,
- get_cast_values::<UInt16Type>(&u32_array, &DataType::UInt16)
- );
-
- let u8_expected = vec!["0", "255", "null", "null"];
- assert_eq!(
- u8_expected,
- get_cast_values::<UInt8Type>(&u32_array, &DataType::UInt8)
- );
- }
-
- #[test]
- fn test_cast_from_uint16() {
- let u16_values: Vec<u16> = vec![0, std::u8::MAX as u16, std::u16::MAX as u16];
- let u16_array: ArrayRef = Arc::new(UInt16Array::from(u16_values));
-
- let f64_expected = vec!["0.0", "255.0", "65535.0"];
- assert_eq!(
- f64_expected,
- get_cast_values::<Float64Type>(&u16_array, &DataType::Float64)
- );
-
- let f32_expected = vec!["0.0", "255.0", "65535.0"];
- assert_eq!(
- f32_expected,
- get_cast_values::<Float32Type>(&u16_array, &DataType::Float32)
- );
-
- let i64_expected = vec!["0", "255", "65535"];
- assert_eq!(
- i64_expected,
- get_cast_values::<Int64Type>(&u16_array, &DataType::Int64)
- );
-
- let i32_expected = vec!["0", "255", "65535"];
- assert_eq!(
- i32_expected,
- get_cast_values::<Int32Type>(&u16_array, &DataType::Int32)
- );
-
- let i16_expected = vec!["0", "255", "null"];
- assert_eq!(
- i16_expected,
- get_cast_values::<Int16Type>(&u16_array, &DataType::Int16)
- );
-
- let i8_expected = vec!["0", "null", "null"];
- assert_eq!(
- i8_expected,
- get_cast_values::<Int8Type>(&u16_array, &DataType::Int8)
- );
-
- let u64_expected = vec!["0", "255", "65535"];
- assert_eq!(
- u64_expected,
- get_cast_values::<UInt64Type>(&u16_array, &DataType::UInt64)
- );
-
- let u32_expected = vec!["0", "255", "65535"];
- assert_eq!(
- u32_expected,
- get_cast_values::<UInt32Type>(&u16_array, &DataType::UInt32)
- );
-
- let u16_expected = vec!["0", "255", "65535"];
- assert_eq!(
- u16_expected,
- get_cast_values::<UInt16Type>(&u16_array, &DataType::UInt16)
- );
-
- let u8_expected = vec!["0", "255", "null"];
- assert_eq!(
- u8_expected,
- get_cast_values::<UInt8Type>(&u16_array, &DataType::UInt8)
- );
- }
-
- #[test]
- fn test_cast_from_uint8() {
- let u8_values: Vec<u8> = vec![0, std::u8::MAX];
- let u8_array: ArrayRef = Arc::new(UInt8Array::from(u8_values));
-
- let f64_expected = vec!["0.0", "255.0"];
- assert_eq!(
- f64_expected,
- get_cast_values::<Float64Type>(&u8_array, &DataType::Float64)
- );
-
- let f32_expected = vec!["0.0", "255.0"];
- assert_eq!(
- f32_expected,
- get_cast_values::<Float32Type>(&u8_array, &DataType::Float32)
- );
-
- let i64_expected = vec!["0", "255"];
- assert_eq!(
- i64_expected,
- get_cast_values::<Int64Type>(&u8_array, &DataType::Int64)
- );
-
- let i32_expected = vec!["0", "255"];
- assert_eq!(
- i32_expected,
- get_cast_values::<Int32Type>(&u8_array, &DataType::Int32)
- );
-
- let i16_expected = vec!["0", "255"];
- assert_eq!(
- i16_expected,
- get_cast_values::<Int16Type>(&u8_array, &DataType::Int16)
- );
-
- let i8_expected = vec!["0", "null"];
- assert_eq!(
- i8_expected,
- get_cast_values::<Int8Type>(&u8_array, &DataType::Int8)
- );
-
- let u64_expected = vec!["0", "255"];
- assert_eq!(
- u64_expected,
- get_cast_values::<UInt64Type>(&u8_array, &DataType::UInt64)
- );
-
- let u32_expected = vec!["0", "255"];
- assert_eq!(
- u32_expected,
- get_cast_values::<UInt32Type>(&u8_array, &DataType::UInt32)
- );
-
- let u16_expected = vec!["0", "255"];
- assert_eq!(
- u16_expected,
- get_cast_values::<UInt16Type>(&u8_array, &DataType::UInt16)
- );
-
- let u8_expected = vec!["0", "255"];
- assert_eq!(
- u8_expected,
- get_cast_values::<UInt8Type>(&u8_array, &DataType::UInt8)
- );
- }
-
- #[test]
- fn test_cast_from_int64() {
- let i64_values: Vec<i64> = vec![
- std::i64::MIN,
- std::i32::MIN as i64,
- std::i16::MIN as i64,
- std::i8::MIN as i64,
- 0,
- std::i8::MAX as i64,
- std::i16::MAX as i64,
- std::i32::MAX as i64,
- std::i64::MAX,
- ];
- let i64_array: ArrayRef = Arc::new(Int64Array::from(i64_values));
-
- let f64_expected = vec![
- "-9223372036854776000.0",
- "-2147483648.0",
- "-32768.0",
- "-128.0",
- "0.0",
- "127.0",
- "32767.0",
- "2147483647.0",
- "9223372036854776000.0",
- ];
- assert_eq!(
- f64_expected,
- get_cast_values::<Float64Type>(&i64_array, &DataType::Float64)
- );
-
- let f32_expected = vec![
- "-9223372000000000000.0",
- "-2147483600.0",
- "-32768.0",
- "-128.0",
- "0.0",
- "127.0",
- "32767.0",
- "2147483600.0",
- "9223372000000000000.0",
- ];
- assert_eq!(
- f32_expected,
- get_cast_values::<Float32Type>(&i64_array, &DataType::Float32)
- );
-
- let i64_expected = vec![
- "-9223372036854775808",
- "-2147483648",
- "-32768",
- "-128",
- "0",
- "127",
- "32767",
- "2147483647",
- "9223372036854775807",
- ];
- assert_eq!(
- i64_expected,
- get_cast_values::<Int64Type>(&i64_array, &DataType::Int64)
- );
-
- let i32_expected = vec![
- "null",
- "-2147483648",
- "-32768",
- "-128",
- "0",
- "127",
- "32767",
- "2147483647",
- "null",
- ];
- assert_eq!(
- i32_expected,
- get_cast_values::<Int32Type>(&i64_array, &DataType::Int32)
- );
-
- assert_eq!(
- i32_expected,
- get_cast_values::<Date32Type>(&i64_array, &DataType::Date32)
- );
-
- let i16_expected = vec![
- "null", "null", "-32768", "-128", "0", "127", "32767", "null", "null",
- ];
- assert_eq!(
- i16_expected,
- get_cast_values::<Int16Type>(&i64_array, &DataType::Int16)
- );
-
- let i8_expected = vec![
- "null", "null", "null", "-128", "0", "127", "null", "null", "null",
- ];
- assert_eq!(
- i8_expected,
- get_cast_values::<Int8Type>(&i64_array, &DataType::Int8)
- );
-
- let u64_expected = vec![
- "null",
- "null",
- "null",
- "null",
- "0",
- "127",
- "32767",
- "2147483647",
- "9223372036854775807",
- ];
- assert_eq!(
- u64_expected,
- get_cast_values::<UInt64Type>(&i64_array, &DataType::UInt64)
- );
-
- let u32_expected = vec![
- "null",
- "null",
- "null",
- "null",
- "0",
- "127",
- "32767",
- "2147483647",
- "null",
- ];
- assert_eq!(
- u32_expected,
- get_cast_values::<UInt32Type>(&i64_array, &DataType::UInt32)
- );
-
- let u16_expected = vec![
- "null", "null", "null", "null", "0", "127", "32767", "null", "null",
- ];
- assert_eq!(
- u16_expected,
- get_cast_values::<UInt16Type>(&i64_array, &DataType::UInt16)
- );
-
- let u8_expected = vec![
- "null", "null", "null", "null", "0", "127", "null", "null", "null",
- ];
- assert_eq!(
- u8_expected,
- get_cast_values::<UInt8Type>(&i64_array, &DataType::UInt8)
- );
- }
-
- #[test]
- fn test_cast_from_int32() {
- let i32_values: Vec<i32> = vec![
- std::i32::MIN as i32,
- std::i16::MIN as i32,
- std::i8::MIN as i32,
- 0,
- std::i8::MAX as i32,
- std::i16::MAX as i32,
- std::i32::MAX as i32,
- ];
- let i32_array: ArrayRef = Arc::new(Int32Array::from(i32_values));
-
- let f64_expected = vec![
- "-2147483648.0",
- "-32768.0",
- "-128.0",
- "0.0",
- "127.0",
- "32767.0",
- "2147483647.0",
- ];
- assert_eq!(
- f64_expected,
- get_cast_values::<Float64Type>(&i32_array, &DataType::Float64)
- );
-
- let f32_expected = vec![
- "-2147483600.0",
- "-32768.0",
- "-128.0",
- "0.0",
- "127.0",
- "32767.0",
- "2147483600.0",
- ];
- assert_eq!(
- f32_expected,
- get_cast_values::<Float32Type>(&i32_array, &DataType::Float32)
- );
-
- let i16_expected = vec!["null", "-32768", "-128", "0", "127", "32767", "null"];
- assert_eq!(
- i16_expected,
- get_cast_values::<Int16Type>(&i32_array, &DataType::Int16)
- );
-
- let i8_expected = vec!["null", "null", "-128", "0", "127", "null", "null"];
- assert_eq!(
- i8_expected,
- get_cast_values::<Int8Type>(&i32_array, &DataType::Int8)
- );
-
- let u64_expected =
- vec!["null", "null", "null", "0", "127", "32767", "2147483647"];
- assert_eq!(
- u64_expected,
- get_cast_values::<UInt64Type>(&i32_array, &DataType::UInt64)
- );
-
- let u32_expected =
- vec!["null", "null", "null", "0", "127", "32767", "2147483647"];
- assert_eq!(
- u32_expected,
- get_cast_values::<UInt32Type>(&i32_array, &DataType::UInt32)
- );
-
- let u16_expected = vec!["null", "null", "null", "0", "127", "32767", "null"];
- assert_eq!(
- u16_expected,
- get_cast_values::<UInt16Type>(&i32_array, &DataType::UInt16)
- );
-
- let u8_expected = vec!["null", "null", "null", "0", "127", "null", "null"];
- assert_eq!(
- u8_expected,
- get_cast_values::<UInt8Type>(&i32_array, &DataType::UInt8)
- );
-
- // The date32 to date64 cast increases the numerical values in order to keep the same dates.
- let i64_expected = vec![
- "-185542587187200000",
- "-2831155200000",
- "-11059200000",
- "0",
- "10972800000",
- "2831068800000",
- "185542587100800000",
- ];
- assert_eq!(
- i64_expected,
- get_cast_values::<Date64Type>(&i32_array, &DataType::Date64)
- );
- }
-
- #[test]
- fn test_cast_from_int16() {
- let i16_values: Vec<i16> = vec![
- std::i16::MIN,
- std::i8::MIN as i16,
- 0,
- std::i8::MAX as i16,
- std::i16::MAX,
- ];
- let i16_array: ArrayRef = Arc::new(Int16Array::from(i16_values));
-
- let f64_expected = vec!["-32768.0", "-128.0", "0.0", "127.0", "32767.0"];
- assert_eq!(
- f64_expected,
- get_cast_values::<Float64Type>(&i16_array, &DataType::Float64)
- );
-
- let f32_expected = vec!["-32768.0", "-128.0", "0.0", "127.0", "32767.0"];
- assert_eq!(
- f32_expected,
- get_cast_values::<Float32Type>(&i16_array, &DataType::Float32)
- );
-
- let i64_expected = vec!["-32768", "-128", "0", "127", "32767"];
- assert_eq!(
- i64_expected,
- get_cast_values::<Int64Type>(&i16_array, &DataType::Int64)
- );
-
- let i32_expected = vec!["-32768", "-128", "0", "127", "32767"];
- assert_eq!(
- i32_expected,
- get_cast_values::<Int32Type>(&i16_array, &DataType::Int32)
- );
-
- let i16_expected = vec!["-32768", "-128", "0", "127", "32767"];
- assert_eq!(
- i16_expected,
- get_cast_values::<Int16Type>(&i16_array, &DataType::Int16)
- );
-
- let i8_expected = vec!["null", "-128", "0", "127", "null"];
- assert_eq!(
- i8_expected,
- get_cast_values::<Int8Type>(&i16_array, &DataType::Int8)
- );
-
- let u64_expected = vec!["null", "null", "0", "127", "32767"];
- assert_eq!(
- u64_expected,
- get_cast_values::<UInt64Type>(&i16_array, &DataType::UInt64)
- );
-
- let u32_expected = vec!["null", "null", "0", "127", "32767"];
- assert_eq!(
- u32_expected,
- get_cast_values::<UInt32Type>(&i16_array, &DataType::UInt32)
- );
-
- let u16_expected = vec!["null", "null", "0", "127", "32767"];
- assert_eq!(
- u16_expected,
- get_cast_values::<UInt16Type>(&i16_array, &DataType::UInt16)
- );
-
- let u8_expected = vec!["null", "null", "0", "127", "null"];
- assert_eq!(
- u8_expected,
- get_cast_values::<UInt8Type>(&i16_array, &DataType::UInt8)
- );
- }
-
- #[test]
- fn test_cast_from_date32() {
- let i32_values: Vec<i32> = vec![
- std::i32::MIN as i32,
- std::i16::MIN as i32,
- std::i8::MIN as i32,
- 0,
- std::i8::MAX as i32,
- std::i16::MAX as i32,
- std::i32::MAX as i32,
- ];
- let date32_array: ArrayRef = Arc::new(Date32Array::from(i32_values));
-
- let i64_expected = vec![
- "-2147483648",
- "-32768",
- "-128",
- "0",
- "127",
- "32767",
- "2147483647",
- ];
- assert_eq!(
- i64_expected,
- get_cast_values::<Int64Type>(&date32_array, &DataType::Int64)
- );
- }
-
- #[test]
- fn test_cast_from_int8() {
- let i8_values: Vec<i8> = vec![std::i8::MIN, 0, std::i8::MAX];
- let i8_array: ArrayRef = Arc::new(Int8Array::from(i8_values));
-
- let f64_expected = vec!["-128.0", "0.0", "127.0"];
- assert_eq!(
- f64_expected,
- get_cast_values::<Float64Type>(&i8_array, &DataType::Float64)
- );
-
- let f32_expected = vec!["-128.0", "0.0", "127.0"];
- assert_eq!(
- f32_expected,
- get_cast_values::<Float32Type>(&i8_array, &DataType::Float32)
- );
-
- let i64_expected = vec!["-128", "0", "127"];
- assert_eq!(
- i64_expected,
- get_cast_values::<Int64Type>(&i8_array, &DataType::Int64)
- );
-
- let i32_expected = vec!["-128", "0", "127"];
- assert_eq!(
- i32_expected,
- get_cast_values::<Int32Type>(&i8_array, &DataType::Int32)
- );
-
- let i16_expected = vec!["-128", "0", "127"];
- assert_eq!(
- i16_expected,
- get_cast_values::<Int16Type>(&i8_array, &DataType::Int16)
- );
-
- let i8_expected = vec!["-128", "0", "127"];
- assert_eq!(
- i8_expected,
- get_cast_values::<Int8Type>(&i8_array, &DataType::Int8)
- );
-
- let u64_expected = vec!["null", "0", "127"];
- assert_eq!(
- u64_expected,
- get_cast_values::<UInt64Type>(&i8_array, &DataType::UInt64)
- );
-
- let u32_expected = vec!["null", "0", "127"];
- assert_eq!(
- u32_expected,
- get_cast_values::<UInt32Type>(&i8_array, &DataType::UInt32)
- );
-
- let u16_expected = vec!["null", "0", "127"];
- assert_eq!(
- u16_expected,
- get_cast_values::<UInt16Type>(&i8_array, &DataType::UInt16)
- );
-
- let u8_expected = vec!["null", "0", "127"];
- assert_eq!(
- u8_expected,
- get_cast_values::<UInt8Type>(&i8_array, &DataType::UInt8)
- );
- }
-
- /// Convert `array` into a vector of strings by casting to data type dt
- fn get_cast_values<T>(array: &ArrayRef, dt: &DataType) -> Vec<String>
- where
- T: ArrowNumericType,
- {
- let c = cast(&array, dt).unwrap();
- let a = c.as_any().downcast_ref::<PrimitiveArray<T>>().unwrap();
- let mut v: Vec<String> = vec![];
- for i in 0..array.len() {
- if a.is_null(i) {
- v.push("null".to_string())
- } else {
- v.push(format!("{:?}", a.value(i)));
- }
- }
- v
- }
-
- #[test]
- fn test_cast_utf8_dict() {
- // FROM a dictionary with of Utf8 values
- use DataType::*;
-
- let keys_builder = PrimitiveBuilder::<Int8Type>::new(10);
- let values_builder = StringBuilder::new(10);
- let mut builder = StringDictionaryBuilder::new(keys_builder, values_builder);
- builder.append("one").unwrap();
- builder.append_null().unwrap();
- builder.append("three").unwrap();
- let array: ArrayRef = Arc::new(builder.finish());
-
- let expected = vec!["one", "null", "three"];
-
- // Test casting TO StringArray
- let cast_type = Utf8;
- let cast_array = cast(&array, &cast_type).expect("cast to UTF-8 failed");
- assert_eq!(cast_array.data_type(), &cast_type);
- assert_eq!(array_to_strings(&cast_array), expected);
-
- // Test casting TO Dictionary (with different index sizes)
-
- let cast_type = Dictionary(Box::new(Int16), Box::new(Utf8));
- let cast_array = cast(&array, &cast_type).expect("cast failed");
- assert_eq!(cast_array.data_type(), &cast_type);
- assert_eq!(array_to_strings(&cast_array), expected);
-
- let cast_type = Dictionary(Box::new(Int32), Box::new(Utf8));
- let cast_array = cast(&array, &cast_type).expect("cast failed");
- assert_eq!(cast_array.data_type(), &cast_type);
- assert_eq!(array_to_strings(&cast_array), expected);
-
- let cast_type = Dictionary(Box::new(Int64), Box::new(Utf8));
- let cast_array = cast(&array, &cast_type).expect("cast failed");
- assert_eq!(cast_array.data_type(), &cast_type);
- assert_eq!(array_to_strings(&cast_array), expected);
-
- let cast_type = Dictionary(Box::new(UInt8), Box::new(Utf8));
- let cast_array = cast(&array, &cast_type).expect("cast failed");
- assert_eq!(cast_array.data_type(), &cast_type);
- assert_eq!(array_to_strings(&cast_array), expected);
-
- let cast_type = Dictionary(Box::new(UInt16), Box::new(Utf8));
- let cast_array = cast(&array, &cast_type).expect("cast failed");
- assert_eq!(cast_array.data_type(), &cast_type);
- assert_eq!(array_to_strings(&cast_array), expected);
-
- let cast_type = Dictionary(Box::new(UInt32), Box::new(Utf8));
- let cast_array = cast(&array, &cast_type).expect("cast failed");
- assert_eq!(cast_array.data_type(), &cast_type);
- assert_eq!(array_to_strings(&cast_array), expected);
-
- let cast_type = Dictionary(Box::new(UInt64), Box::new(Utf8));
- let cast_array = cast(&array, &cast_type).expect("cast failed");
- assert_eq!(cast_array.data_type(), &cast_type);
- assert_eq!(array_to_strings(&cast_array), expected);
- }
-
- #[test]
- fn test_cast_dict_to_dict_bad_index_value_primitive() {
- use DataType::*;
- // test converting from an array that has indexes of a type
- // that are out of bounds for a particular other kind of
- // index.
-
- let keys_builder = PrimitiveBuilder::<Int32Type>::new(10);
- let values_builder = PrimitiveBuilder::<Int64Type>::new(10);
- let mut builder = PrimitiveDictionaryBuilder::new(keys_builder, values_builder);
-
- // add 200 distinct values (which can be stored by a
- // dictionary indexed by int32, but not a dictionary indexed
- // with int8)
- for i in 0..200 {
- builder.append(i).unwrap();
- }
- let array: ArrayRef = Arc::new(builder.finish());
-
- let cast_type = Dictionary(Box::new(Int8), Box::new(Utf8));
- let res = cast(&array, &cast_type);
- assert!(res.is_err());
- let actual_error = format!("{:?}", res);
- let expected_error = "Could not convert 72 dictionary indexes from Int32 to Int8";
- assert!(
- actual_error.contains(expected_error),
- "did not find expected error '{}' in actual error '{}'",
- actual_error,
- expected_error
- );
- }
-
- #[test]
- fn test_cast_dict_to_dict_bad_index_value_utf8() {
- use DataType::*;
- // Same test as test_cast_dict_to_dict_bad_index_value but use
- // string values (and encode the expected behavior here);
-
- let keys_builder = PrimitiveBuilder::<Int32Type>::new(10);
- let values_builder = StringBuilder::new(10);
- let mut builder = StringDictionaryBuilder::new(keys_builder, values_builder);
-
- // add 200 distinct values (which can be stored by a
- // dictionary indexed by int32, but not a dictionary indexed
- // with int8)
- for i in 0..200 {
- let val = format!("val{}", i);
- builder.append(&val).unwrap();
- }
- let array: ArrayRef = Arc::new(builder.finish());
-
- let cast_type = Dictionary(Box::new(Int8), Box::new(Utf8));
- let res = cast(&array, &cast_type);
- assert!(res.is_err());
- let actual_error = format!("{:?}", res);
- let expected_error = "Could not convert 72 dictionary indexes from Int32 to Int8";
- assert!(
- actual_error.contains(expected_error),
- "did not find expected error '{}' in actual error '{}'",
- actual_error,
- expected_error
- );
- }
-
- #[test]
- fn test_cast_primitive_dict() {
- // FROM a dictionary with of INT32 values
- use DataType::*;
-
- let keys_builder = PrimitiveBuilder::<Int8Type>::new(10);
- let values_builder = PrimitiveBuilder::<Int32Type>::new(10);
- let mut builder = PrimitiveDictionaryBuilder::new(keys_builder, values_builder);
- builder.append(1).unwrap();
- builder.append_null().unwrap();
- builder.append(3).unwrap();
- let array: ArrayRef = Arc::new(builder.finish());
-
- let expected = vec!["1", "null", "3"];
-
- // Test casting TO PrimitiveArray, different dictionary type
- let cast_array = cast(&array, &Utf8).expect("cast to UTF-8 failed");
- assert_eq!(array_to_strings(&cast_array), expected);
- assert_eq!(cast_array.data_type(), &Utf8);
-
- let cast_array = cast(&array, &Int64).expect("cast to int64 failed");
- assert_eq!(array_to_strings(&cast_array), expected);
- assert_eq!(cast_array.data_type(), &Int64);
- }
-
- #[test]
- fn test_cast_primitive_array_to_dict() {
- use DataType::*;
-
- let mut builder = PrimitiveBuilder::<Int32Type>::new(10);
- builder.append_value(1).unwrap();
- builder.append_null().unwrap();
- builder.append_value(3).unwrap();
- let array: ArrayRef = Arc::new(builder.finish());
-
- let expected = vec!["1", "null", "3"];
-
- // Cast to a dictionary (same value type, Int32)
- let cast_type = Dictionary(Box::new(UInt8), Box::new(Int32));
- let cast_array = cast(&array, &cast_type).expect("cast failed");
- assert_eq!(cast_array.data_type(), &cast_type);
- assert_eq!(array_to_strings(&cast_array), expected);
-
- // Cast to a dictionary (different value type, Int8)
- let cast_type = Dictionary(Box::new(UInt8), Box::new(Int8));
- let cast_array = cast(&array, &cast_type).expect("cast failed");
- assert_eq!(cast_array.data_type(), &cast_type);
- assert_eq!(array_to_strings(&cast_array), expected);
- }
-
- #[test]
- fn test_cast_string_array_to_dict() {
- use DataType::*;
-
- let array = Arc::new(StringArray::from(vec![Some("one"), None, Some("three")]))
- as ArrayRef;
-
- let expected = vec!["one", "null", "three"];
-
- // Cast to a dictionary (same value type, Utf8)
- let cast_type = Dictionary(Box::new(UInt8), Box::new(Utf8));
- let cast_array = cast(&array, &cast_type).expect("cast failed");
- assert_eq!(cast_array.data_type(), &cast_type);
- assert_eq!(array_to_strings(&cast_array), expected);
- }
-
- #[test]
- fn test_cast_null_array_to_int32() {
- let array = Arc::new(NullArray::new(6)) as ArrayRef;
-
- let expected = Int32Array::from(vec![None; 6]);
-
- // Cast to a dictionary (same value type, Utf8)
- let cast_type = DataType::Int32;
- let cast_array = cast(&array, &cast_type).expect("cast failed");
- let cast_array = as_primitive_array::<Int32Type>(&cast_array);
- assert_eq!(cast_array.data_type(), &cast_type);
- assert_eq!(cast_array, &expected);
- }
-
- /// Print the `DictionaryArray` `array` as a vector of strings
- fn array_to_strings(array: &ArrayRef) -> Vec<String> {
- (0..array.len())
- .map(|i| {
- if array.is_null(i) {
- "null".to_string()
- } else {
- array_value_to_string(array, i).expect("Convert array to String")
- }
- })
- .collect()
- }
-
- #[test]
- fn test_cast_utf8_to_date32() {
- use chrono::NaiveDate;
- let from_ymd = chrono::NaiveDate::from_ymd;
- let since = chrono::NaiveDate::signed_duration_since;
-
- let a = StringArray::from(vec![
- "2000-01-01", // valid date with leading 0s
- "2000-2-2", // valid date without leading 0s
- "2000-00-00", // invalid month and day
- "2000-01-01T12:00:00", // date + time is invalid
- "2000", // just a year is invalid
- ]);
- let array = Arc::new(a) as ArrayRef;
- let b = cast(&array, &DataType::Date32).unwrap();
- let c = b.as_any().downcast_ref::<Date32Array>().unwrap();
-
- // test valid inputs
- let date_value = since(NaiveDate::from_ymd(2000, 1, 1), from_ymd(1970, 1, 1))
- .num_days() as i32;
- assert_eq!(true, c.is_valid(0)); // "2000-01-01"
- assert_eq!(date_value, c.value(0));
-
- let date_value = since(NaiveDate::from_ymd(2000, 2, 2), from_ymd(1970, 1, 1))
- .num_days() as i32;
- assert_eq!(true, c.is_valid(1)); // "2000-2-2"
- assert_eq!(date_value, c.value(1));
-
- // test invalid inputs
- assert_eq!(false, c.is_valid(2)); // "2000-00-00"
- assert_eq!(false, c.is_valid(3)); // "2000-01-01T12:00:00"
- assert_eq!(false, c.is_valid(4)); // "2000"
- }
-
- #[test]
- fn test_cast_utf8_to_date64() {
- let a = StringArray::from(vec![
- "2000-01-01T12:00:00", // date + time valid
- "2020-12-15T12:34:56", // date + time valid
- "2020-2-2T12:34:56", // valid date time without leading 0s
- "2000-00-00T12:00:00", // invalid month and day
- "2000-01-01 12:00:00", // missing the 'T'
- "2000-01-01", // just a date is invalid
- ]);
- let array = Arc::new(a) as ArrayRef;
- let b = cast(&array, &DataType::Date64).unwrap();
- let c = b.as_any().downcast_ref::<Date64Array>().unwrap();
-
- // test valid inputs
- assert_eq!(true, c.is_valid(0)); // "2000-01-01T12:00:00"
- assert_eq!(946728000000, c.value(0));
- assert_eq!(true, c.is_valid(1)); // "2020-12-15T12:34:56"
- assert_eq!(1608035696000, c.value(1));
- assert_eq!(true, c.is_valid(2)); // "2020-2-2T12:34:56"
- assert_eq!(1580646896000, c.value(2));
-
- // test invalid inputs
- assert_eq!(false, c.is_valid(3)); // "2000-00-00T12:00:00"
- assert_eq!(false, c.is_valid(4)); // "2000-01-01 12:00:00"
- assert_eq!(false, c.is_valid(5)); // "2000-01-01"
- }
-
- #[test]
- #[cfg_attr(miri, ignore)] // running forever
- fn test_can_cast_types() {
- // this function attempts to ensure that can_cast_types stays
- // in sync with cast. It simply tries all combinations of
- // types and makes sure that if `can_cast_types` returns
- // true, so does `cast`
-
- let all_types = get_all_types();
-
- for array in get_arrays_of_all_types() {
- for to_type in &all_types {
- println!("Test casting {:?} --> {:?}", array.data_type(), to_type);
- let cast_result = cast(&array, &to_type);
- let reported_cast_ability = can_cast_types(array.data_type(), to_type);
-
- // check for mismatch
- match (cast_result, reported_cast_ability) {
- (Ok(_), false) => {
- panic!("Was able to cast array {:?} from {:?} to {:?} but can_cast_types reported false",
- array, array.data_type(), to_type)
- }
- (Err(e), true) => {
- panic!("Was not able to cast array {:?} from {:?} to {:?} but can_cast_types reported true. \
- Error was {:?}",
- array, array.data_type(), to_type, e)
- }
- // otherwise it was a match
- _ => {}
- };
- }
- }
- }
-
- #[test]
- fn test_cast_list_containers() {
- // large-list to list
- let array = Arc::new(make_large_list_array()) as ArrayRef;
- let list_array = cast(
- &array,
- &DataType::List(Box::new(Field::new("", DataType::Int32, false))),
- )
- .unwrap();
- let actual = list_array.as_any().downcast_ref::<ListArray>().unwrap();
- let expected = array.as_any().downcast_ref::<LargeListArray>().unwrap();
-
- assert_eq!(&expected.value(0), &actual.value(0));
- assert_eq!(&expected.value(1), &actual.value(1));
- assert_eq!(&expected.value(2), &actual.value(2));
-
- // list to large-list
- let array = Arc::new(make_list_array()) as ArrayRef;
- let large_list_array = cast(
- &array,
- &DataType::LargeList(Box::new(Field::new("", DataType::Int32, false))),
- )
- .unwrap();
- let actual = large_list_array
- .as_any()
- .downcast_ref::<LargeListArray>()
- .unwrap();
- let expected = array.as_any().downcast_ref::<ListArray>().unwrap();
-
- assert_eq!(&expected.value(0), &actual.value(0));
- assert_eq!(&expected.value(1), &actual.value(1));
- assert_eq!(&expected.value(2), &actual.value(2));
- }
-
- /// Create instances of arrays with varying types for cast tests
- fn get_arrays_of_all_types() -> Vec<ArrayRef> {
- let tz_name = String::from("America/New_York");
- let binary_data: Vec<&[u8]> = vec![b"foo", b"bar"];
- vec![
- Arc::new(BinaryArray::from(binary_data.clone())),
- Arc::new(LargeBinaryArray::from(binary_data.clone())),
- make_dictionary_primitive::<Int8Type>(),
- make_dictionary_primitive::<Int16Type>(),
- make_dictionary_primitive::<Int32Type>(),
- make_dictionary_primitive::<Int64Type>(),
- make_dictionary_primitive::<UInt8Type>(),
- make_dictionary_primitive::<UInt16Type>(),
- make_dictionary_primitive::<UInt32Type>(),
- make_dictionary_primitive::<UInt64Type>(),
- make_dictionary_utf8::<Int8Type>(),
- make_dictionary_utf8::<Int16Type>(),
- make_dictionary_utf8::<Int32Type>(),
- make_dictionary_utf8::<Int64Type>(),
- make_dictionary_utf8::<UInt8Type>(),
- make_dictionary_utf8::<UInt16Type>(),
- make_dictionary_utf8::<UInt32Type>(),
- make_dictionary_utf8::<UInt64Type>(),
- Arc::new(make_list_array()),
- Arc::new(make_large_list_array()),
- Arc::new(make_fixed_size_list_array()),
- Arc::new(make_fixed_size_binary_array()),
- Arc::new(StructArray::from(vec![
- (
- Field::new("a", DataType::Boolean, false),
- Arc::new(BooleanArray::from(vec![false, false, true, true]))
- as Arc<Array>,
- ),
- (
- Field::new("b", DataType::Int32, false),
- Arc::new(Int32Array::from(vec![42, 28, 19, 31])),
- ),
- ])),
- //Arc::new(make_union_array()),
- Arc::new(NullArray::new(10)),
- Arc::new(StringArray::from(vec!["foo", "bar"])),
- Arc::new(LargeStringArray::from(vec!["foo", "bar"])),
- Arc::new(BooleanArray::from(vec![true, false])),
- Arc::new(Int8Array::from(vec![1, 2])),
- Arc::new(Int16Array::from(vec![1, 2])),
- Arc::new(Int32Array::from(vec![1, 2])),
- Arc::new(Int64Array::from(vec![1, 2])),
- Arc::new(UInt8Array::from(vec![1, 2])),
- Arc::new(UInt16Array::from(vec![1, 2])),
- Arc::new(UInt32Array::from(vec![1, 2])),
- Arc::new(UInt64Array::from(vec![1, 2])),
- Arc::new(Float32Array::from(vec![1.0, 2.0])),
- Arc::new(Float64Array::from(vec![1.0, 2.0])),
- Arc::new(TimestampSecondArray::from_vec(vec![1000, 2000], None)),
- Arc::new(TimestampMillisecondArray::from_vec(vec![1000, 2000], None)),
- Arc::new(TimestampMicrosecondArray::from_vec(vec![1000, 2000], None)),
- Arc::new(TimestampNanosecondArray::from_vec(vec![1000, 2000], None)),
- Arc::new(TimestampSecondArray::from_vec(
- vec![1000, 2000],
- Some(tz_name.clone()),
- )),
- Arc::new(TimestampMillisecondArray::from_vec(
- vec![1000, 2000],
- Some(tz_name.clone()),
- )),
- Arc::new(TimestampMicrosecondArray::from_vec(
- vec![1000, 2000],
- Some(tz_name.clone()),
- )),
- Arc::new(TimestampNanosecondArray::from_vec(
- vec![1000, 2000],
- Some(tz_name),
- )),
- Arc::new(Date32Array::from(vec![1000, 2000])),
- Arc::new(Date64Array::from(vec![1000, 2000])),
- Arc::new(Time32SecondArray::from(vec![1000, 2000])),
- Arc::new(Time32MillisecondArray::from(vec![1000, 2000])),
- Arc::new(Time64MicrosecondArray::from(vec![1000, 2000])),
- Arc::new(Time64NanosecondArray::from(vec![1000, 2000])),
- Arc::new(IntervalYearMonthArray::from(vec![1000, 2000])),
- Arc::new(IntervalDayTimeArray::from(vec![1000, 2000])),
- Arc::new(DurationSecondArray::from(vec![1000, 2000])),
- Arc::new(DurationMillisecondArray::from(vec![1000, 2000])),
- Arc::new(DurationMicrosecondArray::from(vec![1000, 2000])),
- Arc::new(DurationNanosecondArray::from(vec![1000, 2000])),
- ]
- }
-
- fn make_list_array() -> ListArray {
- // Construct a value array
- let value_data = ArrayData::builder(DataType::Int32)
- .len(8)
- .add_buffer(Buffer::from_slice_ref(&[0, 1, 2, 3, 4, 5, 6, 7]))
- .build();
-
- // Construct a buffer for value offsets, for the nested array:
- // [[0, 1, 2], [3, 4, 5], [6, 7]]
- let value_offsets = Buffer::from_slice_ref(&[0, 3, 6, 8]);
-
- // Construct a list array from the above two
- let list_data_type =
- DataType::List(Box::new(Field::new("item", DataType::Int32, true)));
- let list_data = ArrayData::builder(list_data_type)
- .len(3)
- .add_buffer(value_offsets)
- .add_child_data(value_data)
- .build();
- ListArray::from(list_data)
- }
-
- fn make_large_list_array() -> LargeListArray {
- // Construct a value array
- let value_data = ArrayData::builder(DataType::Int32)
- .len(8)
- .add_buffer(Buffer::from_slice_ref(&[0, 1, 2, 3, 4, 5, 6, 7]))
- .build();
-
- // Construct a buffer for value offsets, for the nested array:
- // [[0, 1, 2], [3, 4, 5], [6, 7]]
- let value_offsets = Buffer::from_slice_ref(&[0i64, 3, 6, 8]);
-
- // Construct a list array from the above two
- let list_data_type =
- DataType::LargeList(Box::new(Field::new("item", DataType::Int32, true)));
- let list_data = ArrayData::builder(list_data_type)
- .len(3)
- .add_buffer(value_offsets)
- .add_child_data(value_data)
- .build();
- LargeListArray::from(list_data)
- }
-
- fn make_fixed_size_list_array() -> FixedSizeListArray {
- // Construct a value array
- let value_data = ArrayData::builder(DataType::Int32)
- .len(10)
- .add_buffer(Buffer::from_slice_ref(&[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]))
- .build();
-
- // Construct a fixed size list array from the above two
- let list_data_type = DataType::FixedSizeList(
- Box::new(Field::new("item", DataType::Int32, true)),
- 2,
- );
- let list_data = ArrayData::builder(list_data_type)
- .len(5)
- .add_child_data(value_data)
- .build();
- FixedSizeListArray::from(list_data)
- }
-
- fn make_fixed_size_binary_array() -> FixedSizeBinaryArray {
- let values: [u8; 15] = *b"hellotherearrow";
-
- let array_data = ArrayData::builder(DataType::FixedSizeBinary(5))
- .len(3)
- .add_buffer(Buffer::from(&values[..]))
- .build();
- FixedSizeBinaryArray::from(array_data)
- }
-
- fn make_union_array() -> UnionArray {
- let mut builder = UnionBuilder::new_dense(7);
- builder.append::<Int32Type>("a", 1).unwrap();
- builder.append::<Int64Type>("b", 2).unwrap();
- builder.build().unwrap()
- }
-
- /// Creates a dictionary with primitive dictionary values, and keys of type K
- fn make_dictionary_primitive<K: ArrowDictionaryKeyType>() -> ArrayRef {
- let keys_builder = PrimitiveBuilder::<K>::new(2);
- // Pick Int32 arbitrarily for dictionary values
- let values_builder = PrimitiveBuilder::<Int32Type>::new(2);
- let mut b = PrimitiveDictionaryBuilder::new(keys_builder, values_builder);
- b.append(1).unwrap();
- b.append(2).unwrap();
- Arc::new(b.finish())
- }
-
- /// Creates a dictionary with utf8 values, and keys of type K
- fn make_dictionary_utf8<K: ArrowDictionaryKeyType>() -> ArrayRef {
- let keys_builder = PrimitiveBuilder::<K>::new(2);
- // Pick Int32 arbitrarily for dictionary values
- let values_builder = StringBuilder::new(2);
- let mut b = StringDictionaryBuilder::new(keys_builder, values_builder);
- b.append("foo").unwrap();
- b.append("bar").unwrap();
- Arc::new(b.finish())
- }
-
- // Get a selection of datatypes to try and cast to
- fn get_all_types() -> Vec<DataType> {
- use DataType::*;
- let tz_name = String::from("America/New_York");
-
- vec![
- Null,
- Boolean,
- Int8,
- Int16,
- Int32,
- UInt64,
- UInt8,
- UInt16,
- UInt32,
- UInt64,
- Float16,
- Float32,
- Float64,
- Timestamp(TimeUnit::Second, None),
- Timestamp(TimeUnit::Millisecond, None),
- Timestamp(TimeUnit::Microsecond, None),
- Timestamp(TimeUnit::Nanosecond, None),
- Timestamp(TimeUnit::Second, Some(tz_name.clone())),
- Timestamp(TimeUnit::Millisecond, Some(tz_name.clone())),
- Timestamp(TimeUnit::Microsecond, Some(tz_name.clone())),
- Timestamp(TimeUnit::Nanosecond, Some(tz_name)),
- Date32,
- Date64,
- Time32(TimeUnit::Second),
- Time32(TimeUnit::Millisecond),
- Time64(TimeUnit::Microsecond),
- Time64(TimeUnit::Nanosecond),
- Duration(TimeUnit::Second),
- Duration(TimeUnit::Millisecond),
- Duration(TimeUnit::Microsecond),
- Duration(TimeUnit::Nanosecond),
- Interval(IntervalUnit::YearMonth),
- Interval(IntervalUnit::DayTime),
- Binary,
- FixedSizeBinary(10),
- LargeBinary,
- Utf8,
- LargeUtf8,
- List(Box::new(Field::new("item", DataType::Int8, true))),
- List(Box::new(Field::new("item", DataType::Utf8, true))),
- FixedSizeList(Box::new(Field::new("item", DataType::Int8, true)), 10),
- FixedSizeList(Box::new(Field::new("item", DataType::Utf8, false)), 10),
- LargeList(Box::new(Field::new("item", DataType::Int8, true))),
- LargeList(Box::new(Field::new("item", DataType::Utf8, false))),
- Struct(vec![
- Field::new("f1", DataType::Int32, false),
- Field::new("f2", DataType::Utf8, true),
- ]),
- Union(vec![
- Field::new("f1", DataType::Int32, false),
- Field::new("f2", DataType::Utf8, true),
- ]),
- Dictionary(Box::new(DataType::Int8), Box::new(DataType::Int32)),
- Dictionary(Box::new(DataType::Int16), Box::new(DataType::Utf8)),
- Dictionary(Box::new(DataType::UInt32), Box::new(DataType::Utf8)),
- ]
- }
-
- #[test]
- fn test_utf8_cast_offsets() {
- // test if offset of the array is taken into account during cast
- let str_array = StringArray::from(vec!["a", "b", "c"]);
- let str_array = str_array.slice(1, 2);
-
- let out = cast(&str_array, &DataType::LargeUtf8).unwrap();
-
- let large_str_array = out.as_any().downcast_ref::<LargeStringArray>().unwrap();
- let strs = large_str_array.into_iter().flatten().collect::<Vec<_>>();
- assert_eq!(strs, &["b", "c"])
- }
-
- #[test]
- fn test_list_cast_offsets() {
- // test if offset of the array is taken into account during cast
- let array1 = make_list_array().slice(1, 2);
- let array2 = Arc::new(make_list_array()) as ArrayRef;
-
- let dt = DataType::LargeList(Box::new(Field::new("item", DataType::Int32, true)));
- let out1 = cast(&array1, &dt).unwrap();
- let out2 = cast(&array2, &dt).unwrap();
-
- assert_eq!(&out1, &out2.slice(1, 2))
- }
-}
diff --git a/arrow/src/compute/kernels/cast_utils.rs b/arrow/src/compute/kernels/cast_utils.rs
deleted file mode 100644
index 201d209..0000000
--- a/arrow/src/compute/kernels/cast_utils.rs
+++ /dev/null
@@ -1,300 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-use crate::error::{ArrowError, Result};
-use chrono::{prelude::*, LocalResult};
-
-/// Accepts a string in RFC3339 / ISO8601 standard format and some
-/// variants and converts it to a nanosecond precision timestamp.
-///
-/// Implements the `to_timestamp` function to convert a string to a
-/// timestamp, following the model of spark SQL’s to_`timestamp`.
-///
-/// In addition to RFC3339 / ISO8601 standard timestamps, it also
-/// accepts strings that use a space ` ` to separate the date and time
-/// as well as strings that have no explicit timezone offset.
-///
-/// Examples of accepted inputs:
-/// * `1997-01-31T09:26:56.123Z` # RCF3339
-/// * `1997-01-31T09:26:56.123-05:00` # RCF3339
-/// * `1997-01-31 09:26:56.123-05:00` # close to RCF3339 but with a space rather than T
-/// * `1997-01-31T09:26:56.123` # close to RCF3339 but no timezone offset specified
-/// * `1997-01-31 09:26:56.123` # close to RCF3339 but uses a space and no timezone offset
-/// * `1997-01-31 09:26:56` # close to RCF3339, no fractional seconds
-//
-/// Internally, this function uses the `chrono` library for the
-/// datetime parsing
-///
-/// We hope to extend this function in the future with a second
-/// parameter to specifying the format string.
-///
-/// ## Timestamp Precision
-///
-/// Function uses the maximum precision timestamps supported by
-/// Arrow (nanoseconds stored as a 64-bit integer) timestamps. This
-/// means the range of dates that timestamps can represent is ~1677 AD
-/// to 2262 AM
-///
-///
-/// ## Timezone / Offset Handling
-///
-/// Numerical values of timestamps are stored compared to offset UTC.
-///
-/// This function intertprets strings without an explicit time zone as
-/// timestamps with offsets of the local time on the machine
-///
-/// For example, `1997-01-31 09:26:56.123Z` is interpreted as UTC, as
-/// it has an explicit timezone specifier (“Z” for Zulu/UTC)
-///
-/// `1997-01-31T09:26:56.123` is interpreted as a local timestamp in
-/// the timezone of the machine. For example, if
-/// the system timezone is set to Americas/New_York (UTC-5) the
-/// timestamp will be interpreted as though it were
-/// `1997-01-31T09:26:56.123-05:00`
-#[inline]
-pub fn string_to_timestamp_nanos(s: &str) -> Result<i64> {
- // Fast path: RFC3339 timestamp (with a T)
- // Example: 2020-09-08T13:42:29.190855Z
- if let Ok(ts) = DateTime::parse_from_rfc3339(s) {
- return Ok(ts.timestamp_nanos());
- }
-
- // Implement quasi-RFC3339 support by trying to parse the
- // timestamp with various other format specifiers to to support
- // separating the date and time with a space ' ' rather than 'T' to be
- // (more) compatible with Apache Spark SQL
-
- // timezone offset, using ' ' as a separator
- // Example: 2020-09-08 13:42:29.190855-05:00
- if let Ok(ts) = DateTime::parse_from_str(s, "%Y-%m-%d %H:%M:%S%.f%:z") {
- return Ok(ts.timestamp_nanos());
- }
-
- // with an explicit Z, using ' ' as a separator
- // Example: 2020-09-08 13:42:29Z
- if let Ok(ts) = Utc.datetime_from_str(s, "%Y-%m-%d %H:%M:%S%.fZ") {
- return Ok(ts.timestamp_nanos());
- }
-
- // Support timestamps without an explicit timezone offset, again
- // to be compatible with what Apache Spark SQL does.
-
- // without a timezone specifier as a local time, using T as a separator
- // Example: 2020-09-08T13:42:29.190855
- if let Ok(ts) = NaiveDateTime::parse_from_str(s, "%Y-%m-%dT%H:%M:%S.%f") {
- return naive_datetime_to_timestamp(s, ts);
- }
-
- // without a timezone specifier as a local time, using T as a
- // separator, no fractional seconds
- // Example: 2020-09-08T13:42:29
- if let Ok(ts) = NaiveDateTime::parse_from_str(s, "%Y-%m-%dT%H:%M:%S") {
- return naive_datetime_to_timestamp(s, ts);
- }
-
- // without a timezone specifier as a local time, using ' ' as a separator
- // Example: 2020-09-08 13:42:29.190855
- if let Ok(ts) = NaiveDateTime::parse_from_str(s, "%Y-%m-%d %H:%M:%S.%f") {
- return naive_datetime_to_timestamp(s, ts);
- }
-
- // without a timezone specifier as a local time, using ' ' as a
- // separator, no fractional seconds
- // Example: 2020-09-08 13:42:29
- if let Ok(ts) = NaiveDateTime::parse_from_str(s, "%Y-%m-%d %H:%M:%S") {
- return naive_datetime_to_timestamp(s, ts);
- }
-
- // Note we don't pass along the error message from the underlying
- // chrono parsing because we tried several different format
- // strings and we don't know which the user was trying to
- // match. Ths any of the specific error messages is likely to be
- // be more confusing than helpful
- Err(ArrowError::CastError(format!(
- "Error parsing '{}' as timestamp",
- s
- )))
-}
-
-/// Converts the naive datetime (which has no specific timezone) to a
-/// nanosecond epoch timestamp relative to UTC.
-fn naive_datetime_to_timestamp(s: &str, datetime: NaiveDateTime) -> Result<i64> {
- let l = Local {};
-
- match l.from_local_datetime(&datetime) {
- LocalResult::None => Err(ArrowError::CastError(format!(
- "Error parsing '{}' as timestamp: local time representation is invalid",
- s
- ))),
- LocalResult::Single(local_datetime) => {
- Ok(local_datetime.with_timezone(&Utc).timestamp_nanos())
- }
- // Ambiguous times can happen if the timestamp is exactly when
- // a daylight savings time transition occurs, for example, and
- // so the datetime could validly be said to be in two
- // potential offsets. However, since we are about to convert
- // to UTC anyways, we can pick one arbitrarily
- LocalResult::Ambiguous(local_datetime, _) => {
- Ok(local_datetime.with_timezone(&Utc).timestamp_nanos())
- }
- }
-}
-
-#[cfg(test)]
-mod tests {
- use super::*;
-
- #[test]
- fn string_to_timestamp_timezone() -> Result<()> {
- // Explicit timezone
- assert_eq!(
- 1599572549190855000,
- parse_timestamp("2020-09-08T13:42:29.190855+00:00")?
- );
- assert_eq!(
- 1599572549190855000,
- parse_timestamp("2020-09-08T13:42:29.190855Z")?
- );
- assert_eq!(
- 1599572549000000000,
- parse_timestamp("2020-09-08T13:42:29Z")?
- ); // no fractional part
- assert_eq!(
- 1599590549190855000,
- parse_timestamp("2020-09-08T13:42:29.190855-05:00")?
- );
- Ok(())
- }
-
- #[test]
- fn string_to_timestamp_timezone_space() -> Result<()> {
- // Ensure space rather than T between time and date is accepted
- assert_eq!(
- 1599572549190855000,
- parse_timestamp("2020-09-08 13:42:29.190855+00:00")?
- );
- assert_eq!(
- 1599572549190855000,
- parse_timestamp("2020-09-08 13:42:29.190855Z")?
- );
- assert_eq!(
- 1599572549000000000,
- parse_timestamp("2020-09-08 13:42:29Z")?
- ); // no fractional part
- assert_eq!(
- 1599590549190855000,
- parse_timestamp("2020-09-08 13:42:29.190855-05:00")?
- );
- Ok(())
- }
-
- /// Interprets a naive_datetime (with no explicit timzone offset)
- /// using the local timezone and returns the timestamp in UTC (0
- /// offset)
- fn naive_datetime_to_timestamp(naive_datetime: &NaiveDateTime) -> i64 {
- // Note: Use chrono APIs that are different than
- // naive_datetime_to_timestamp to compute the utc offset to
- // try and double check the logic
- let utc_offset_secs = match Local.offset_from_local_datetime(&naive_datetime) {
- LocalResult::Single(local_offset) => {
- local_offset.fix().local_minus_utc() as i64
- }
- _ => panic!("Unexpected failure converting to local datetime"),
- };
- let utc_offset_nanos = utc_offset_secs * 1_000_000_000;
- naive_datetime.timestamp_nanos() - utc_offset_nanos
- }
-
- #[test]
- #[cfg_attr(miri, ignore)] // unsupported operation: can't call foreign function: mktime
- fn string_to_timestamp_no_timezone() -> Result<()> {
- // This test is designed to succeed in regardless of the local
- // timezone the test machine is running. Thus it is still
- // somewhat suceptable to bugs in the use of chrono
- let naive_datetime = NaiveDateTime::new(
- NaiveDate::from_ymd(2020, 9, 8),
- NaiveTime::from_hms_nano(13, 42, 29, 190855),
- );
-
- // Ensure both T and ' ' variants work
- assert_eq!(
- naive_datetime_to_timestamp(&naive_datetime),
- parse_timestamp("2020-09-08T13:42:29.190855")?
- );
-
- assert_eq!(
- naive_datetime_to_timestamp(&naive_datetime),
- parse_timestamp("2020-09-08 13:42:29.190855")?
- );
-
- // Also ensure that parsing timestamps with no fractional
- // second part works as well
- let naive_datetime_whole_secs = NaiveDateTime::new(
- NaiveDate::from_ymd(2020, 9, 8),
- NaiveTime::from_hms(13, 42, 29),
- );
-
- // Ensure both T and ' ' variants work
- assert_eq!(
- naive_datetime_to_timestamp(&naive_datetime_whole_secs),
- parse_timestamp("2020-09-08T13:42:29")?
- );
-
- assert_eq!(
- naive_datetime_to_timestamp(&naive_datetime_whole_secs),
- parse_timestamp("2020-09-08 13:42:29")?
- );
-
- Ok(())
- }
-
- #[test]
- fn string_to_timestamp_invalid() {
- // Test parsing invalid formats
-
- // It would be nice to make these messages better
- expect_timestamp_parse_error("", "Error parsing '' as timestamp");
- expect_timestamp_parse_error("SS", "Error parsing 'SS' as timestamp");
- expect_timestamp_parse_error(
- "Wed, 18 Feb 2015 23:16:09 GMT",
- "Error parsing 'Wed, 18 Feb 2015 23:16:09 GMT' as timestamp",
- );
- }
-
- // Parse a timestamp to timestamp int with a useful human readable error message
- fn parse_timestamp(s: &str) -> Result<i64> {
- let result = string_to_timestamp_nanos(s);
- if let Err(e) = &result {
- eprintln!("Error parsing timestamp '{}': {:?}", s, e);
- }
- result
- }
-
- fn expect_timestamp_parse_error(s: &str, expected_err: &str) {
- match string_to_timestamp_nanos(s) {
- Ok(v) => panic!(
- "Expected error '{}' while parsing '{}', but parsed {} instead",
- expected_err, s, v
- ),
- Err(e) => {
- assert!(e.to_string().contains(expected_err),
- "Can not find expected error '{}' while parsing '{}'. Actual error '{}'",
- expected_err, s, e);
- }
- }
- }
-}
diff --git a/arrow/src/compute/kernels/comparison.rs b/arrow/src/compute/kernels/comparison.rs
deleted file mode 100644
index 8e5da22..0000000
--- a/arrow/src/compute/kernels/comparison.rs
+++ /dev/null
@@ -1,1619 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-//! Defines basic comparison kernels for [`PrimitiveArray`]s.
-//!
-//! These kernels can leverage SIMD if available on your system. Currently no runtime
-//! detection is provided, you should enable the specific SIMD intrinsics using
-//! `RUSTFLAGS="-C target-feature=+avx2"` for example. See the documentation
-//! [here](https://doc.rust-lang.org/stable/core/arch/) for more information.
-
-use regex::Regex;
-use std::collections::HashMap;
-
-use crate::array::*;
-use crate::buffer::{Buffer, MutableBuffer};
-use crate::compute::util::combine_option_bitmap;
-use crate::datatypes::{ArrowNumericType, DataType};
-use crate::error::{ArrowError, Result};
-use crate::util::bit_util;
-
-/// Helper function to perform boolean lambda function on values from two arrays, this
-/// version does not attempt to use SIMD.
-macro_rules! compare_op {
- ($left: expr, $right:expr, $op:expr) => {{
- if $left.len() != $right.len() {
- return Err(ArrowError::ComputeError(
- "Cannot perform comparison operation on arrays of different length"
- .to_string(),
- ));
- }
-
- let null_bit_buffer =
- combine_option_bitmap($left.data_ref(), $right.data_ref(), $left.len())?;
-
- let comparison = (0..$left.len()).map(|i| $op($left.value(i), $right.value(i)));
- // same size as $left.len() and $right.len()
- let buffer = unsafe { MutableBuffer::from_trusted_len_iter_bool(comparison) };
-
- let data = ArrayData::new(
- DataType::Boolean,
- $left.len(),
- None,
- null_bit_buffer,
- 0,
- vec![Buffer::from(buffer)],
- vec![],
- );
- Ok(BooleanArray::from(data))
- }};
-}
-
-macro_rules! compare_op_primitive {
- ($left: expr, $right:expr, $op:expr) => {{
- if $left.len() != $right.len() {
- return Err(ArrowError::ComputeError(
- "Cannot perform comparison operation on arrays of different length"
- .to_string(),
- ));
- }
-
- let null_bit_buffer =
- combine_option_bitmap($left.data_ref(), $right.data_ref(), $left.len())?;
-
- let mut values = MutableBuffer::from_len_zeroed(($left.len() + 7) / 8);
- let lhs_chunks_iter = $left.values().chunks_exact(8);
- let lhs_remainder = lhs_chunks_iter.remainder();
- let rhs_chunks_iter = $right.values().chunks_exact(8);
- let rhs_remainder = rhs_chunks_iter.remainder();
- let chunks = $left.len() / 8;
-
- values[..chunks]
- .iter_mut()
- .zip(lhs_chunks_iter)
- .zip(rhs_chunks_iter)
- .for_each(|((byte, lhs), rhs)| {
- lhs.iter()
- .zip(rhs.iter())
- .enumerate()
- .for_each(|(i, (&lhs, &rhs))| {
- *byte |= if $op(lhs, rhs) { 1 << i } else { 0 };
- });
- });
-
- if !lhs_remainder.is_empty() {
- let last = &mut values[chunks];
- lhs_remainder
- .iter()
- .zip(rhs_remainder.iter())
- .enumerate()
- .for_each(|(i, (&lhs, &rhs))| {
- *last |= if $op(lhs, rhs) { 1 << i } else { 0 };
- });
- };
- let data = ArrayData::new(
- DataType::Boolean,
- $left.len(),
- None,
- null_bit_buffer,
- 0,
- vec![Buffer::from(values)],
- vec![],
- );
- Ok(BooleanArray::from(data))
- }};
-}
-
-macro_rules! compare_op_scalar {
- ($left: expr, $right:expr, $op:expr) => {{
- let null_bit_buffer = $left.data().null_buffer().cloned();
-
- let comparison = (0..$left.len()).map(|i| $op($left.value(i), $right));
- // same as $left.len()
- let buffer = unsafe { MutableBuffer::from_trusted_len_iter_bool(comparison) };
-
- let data = ArrayData::new(
- DataType::Boolean,
- $left.len(),
- None,
- null_bit_buffer,
- 0,
- vec![Buffer::from(buffer)],
- vec![],
- );
- Ok(BooleanArray::from(data))
- }};
-}
-
-macro_rules! compare_op_scalar_primitive {
- ($left: expr, $right:expr, $op:expr) => {{
- let null_bit_buffer = $left.data().null_buffer().cloned();
-
- let mut values = MutableBuffer::from_len_zeroed(($left.len() + 7) / 8);
- let lhs_chunks_iter = $left.values().chunks_exact(8);
- let lhs_remainder = lhs_chunks_iter.remainder();
- let chunks = $left.len() / 8;
-
- values[..chunks]
- .iter_mut()
- .zip(lhs_chunks_iter)
- .for_each(|(byte, chunk)| {
- chunk.iter().enumerate().for_each(|(i, &c_i)| {
- *byte |= if $op(c_i, $right) { 1 << i } else { 0 };
- });
- });
- if !lhs_remainder.is_empty() {
- let last = &mut values[chunks];
- lhs_remainder.iter().enumerate().for_each(|(i, &lhs)| {
- *last |= if $op(lhs, $right) { 1 << i } else { 0 };
- });
- };
-
- let data = ArrayData::new(
- DataType::Boolean,
- $left.len(),
- None,
- null_bit_buffer,
- 0,
- vec![Buffer::from(values)],
- vec![],
- );
- Ok(BooleanArray::from(data))
- }};
-}
-
-/// Evaluate `op(left, right)` for [`PrimitiveArray`]s using a specified
-/// comparison function.
-pub fn no_simd_compare_op<T, F>(
- left: &PrimitiveArray<T>,
- right: &PrimitiveArray<T>,
- op: F,
-) -> Result<BooleanArray>
-where
- T: ArrowNumericType,
- F: Fn(T::Native, T::Native) -> bool,
-{
- compare_op_primitive!(left, right, op)
-}
-
-/// Evaluate `op(left, right)` for [`PrimitiveArray`] and scalar using
-/// a specified comparison function.
-pub fn no_simd_compare_op_scalar<T, F>(
- left: &PrimitiveArray<T>,
- right: T::Native,
- op: F,
-) -> Result<BooleanArray>
-where
- T: ArrowNumericType,
- F: Fn(T::Native, T::Native) -> bool,
-{
- compare_op_scalar_primitive!(left, right, op)
-}
-
-/// Perform SQL `left LIKE right` operation on [`StringArray`] / [`LargeStringArray`].
-///
-/// There are two wildcards supported with the LIKE operator:
-///
-/// 1. `%` - The percent sign represents zero, one, or multiple characters
-/// 2. `_` - The underscore represents a single character
-///
-/// For example:
-/// ```
-/// use arrow::array::{StringArray, BooleanArray};
-/// use arrow::compute::like_utf8;
-///
-/// let strings = StringArray::from(vec!["Arrow", "Arrow", "Arrow", "Ar"]);
-/// let patterns = StringArray::from(vec!["A%", "B%", "A.", "A."]);
-///
-/// let result = like_utf8(&strings, &patterns).unwrap();
-/// assert_eq!(result, BooleanArray::from(vec![true, false, false, true]));
-/// ```
-pub fn like_utf8<OffsetSize: StringOffsetSizeTrait>(
- left: &GenericStringArray<OffsetSize>,
- right: &GenericStringArray<OffsetSize>,
-) -> Result<BooleanArray> {
- let mut map = HashMap::new();
- if left.len() != right.len() {
- return Err(ArrowError::ComputeError(
- "Cannot perform comparison operation on arrays of different length"
- .to_string(),
- ));
- }
-
- let null_bit_buffer =
- combine_option_bitmap(left.data_ref(), right.data_ref(), left.len())?;
-
- let mut result = BooleanBufferBuilder::new(left.len());
- for i in 0..left.len() {
- let haystack = left.value(i);
- let pat = right.value(i);
- let re = if let Some(ref regex) = map.get(pat) {
- regex
- } else {
- let re_pattern = pat.replace("%", ".*").replace("_", ".");
- let re = Regex::new(&format!("^{}$", re_pattern)).map_err(|e| {
- ArrowError::ComputeError(format!(
- "Unable to build regex from LIKE pattern: {}",
- e
- ))
- })?;
- map.insert(pat, re);
- map.get(pat).unwrap()
- };
-
- result.append(re.is_match(haystack));
- }
-
- let data = ArrayData::new(
- DataType::Boolean,
- left.len(),
- None,
- null_bit_buffer,
- 0,
- vec![result.finish()],
- vec![],
- );
- Ok(BooleanArray::from(data))
-}
-
-fn is_like_pattern(c: char) -> bool {
- c == '%' || c == '_'
-}
-
-/// Perform SQL `left LIKE right` operation on [`StringArray`] /
-/// [`LargeStringArray`] and a scalar.
-///
-/// See the documentation on [`like_utf8`] for more details.
-pub fn like_utf8_scalar<OffsetSize: StringOffsetSizeTrait>(
- left: &GenericStringArray<OffsetSize>,
- right: &str,
-) -> Result<BooleanArray> {
- let null_bit_buffer = left.data().null_buffer().cloned();
- let bytes = bit_util::ceil(left.len(), 8);
- let mut bool_buf = MutableBuffer::from_len_zeroed(bytes);
- let bool_slice = bool_buf.as_slice_mut();
-
- if !right.contains(is_like_pattern) {
- // fast path, can use equals
- for i in 0..left.len() {
- if left.value(i) == right {
- bit_util::set_bit(bool_slice, i);
- }
- }
- } else if right.ends_with('%') && !right[..right.len() - 1].contains(is_like_pattern)
- {
- // fast path, can use starts_with
- let starts_with = &right[..right.len() - 1];
- for i in 0..left.len() {
- if left.value(i).starts_with(starts_with) {
- bit_util::set_bit(bool_slice, i);
- }
- }
- } else if right.starts_with('%') && !right[1..].contains(is_like_pattern) {
- // fast path, can use ends_with
- let ends_with = &right[1..];
- for i in 0..left.len() {
- if left.value(i).ends_with(ends_with) {
- bit_util::set_bit(bool_slice, i);
- }
- }
- } else {
- let re_pattern = right.replace("%", ".*").replace("_", ".");
- let re = Regex::new(&format!("^{}$", re_pattern)).map_err(|e| {
- ArrowError::ComputeError(format!(
- "Unable to build regex from LIKE pattern: {}",
- e
- ))
- })?;
-
- for i in 0..left.len() {
- let haystack = left.value(i);
- if re.is_match(haystack) {
- bit_util::set_bit(bool_slice, i);
- }
- }
- };
-
- let data = ArrayData::new(
- DataType::Boolean,
- left.len(),
- None,
- null_bit_buffer,
- 0,
- vec![bool_buf.into()],
- vec![],
- );
- Ok(BooleanArray::from(data))
-}
-
-/// Perform SQL `left NOT LIKE right` operation on [`StringArray`] /
-/// [`LargeStringArray`].
-///
-/// See the documentation on [`like_utf8`] for more details.
-pub fn nlike_utf8<OffsetSize: StringOffsetSizeTrait>(
- left: &GenericStringArray<OffsetSize>,
- right: &GenericStringArray<OffsetSize>,
-) -> Result<BooleanArray> {
- let mut map = HashMap::new();
- if left.len() != right.len() {
- return Err(ArrowError::ComputeError(
- "Cannot perform comparison operation on arrays of different length"
- .to_string(),
- ));
- }
-
- let null_bit_buffer =
- combine_option_bitmap(left.data_ref(), right.data_ref(), left.len())?;
-
- let mut result = BooleanBufferBuilder::new(left.len());
- for i in 0..left.len() {
- let haystack = left.value(i);
- let pat = right.value(i);
- let re = if let Some(ref regex) = map.get(pat) {
- regex
- } else {
- let re_pattern = pat.replace("%", ".*").replace("_", ".");
- let re = Regex::new(&format!("^{}$", re_pattern)).map_err(|e| {
- ArrowError::ComputeError(format!(
- "Unable to build regex from LIKE pattern: {}",
- e
- ))
- })?;
- map.insert(pat, re);
- map.get(pat).unwrap()
- };
-
- result.append(!re.is_match(haystack));
- }
-
- let data = ArrayData::new(
- DataType::Boolean,
- left.len(),
- None,
- null_bit_buffer,
- 0,
- vec![result.finish()],
- vec![],
- );
- Ok(BooleanArray::from(data))
-}
-
-/// Perform SQL `left NOT LIKE right` operation on [`StringArray`] /
-/// [`LargeStringArray`] and a scalar.
-///
-/// See the documentation on [`like_utf8`] for more details.
-pub fn nlike_utf8_scalar<OffsetSize: StringOffsetSizeTrait>(
- left: &GenericStringArray<OffsetSize>,
- right: &str,
-) -> Result<BooleanArray> {
- let null_bit_buffer = left.data().null_buffer().cloned();
- let mut result = BooleanBufferBuilder::new(left.len());
-
- if !right.contains(is_like_pattern) {
- // fast path, can use equals
- for i in 0..left.len() {
- result.append(left.value(i) != right);
- }
- } else if right.ends_with('%') && !right[..right.len() - 1].contains(is_like_pattern)
- {
- // fast path, can use ends_with
- for i in 0..left.len() {
- result.append(!left.value(i).starts_with(&right[..right.len() - 1]));
- }
- } else if right.starts_with('%') && !right[1..].contains(is_like_pattern) {
- // fast path, can use starts_with
- for i in 0..left.len() {
- result.append(!left.value(i).ends_with(&right[1..]));
- }
- } else {
- let re_pattern = right.replace("%", ".*").replace("_", ".");
- let re = Regex::new(&format!("^{}$", re_pattern)).map_err(|e| {
- ArrowError::ComputeError(format!(
- "Unable to build regex from LIKE pattern: {}",
- e
- ))
- })?;
- for i in 0..left.len() {
- let haystack = left.value(i);
- result.append(!re.is_match(haystack));
- }
- }
-
- let data = ArrayData::new(
- DataType::Boolean,
- left.len(),
- None,
- null_bit_buffer,
- 0,
- vec![result.finish()],
- vec![],
- );
- Ok(BooleanArray::from(data))
-}
-
-/// Perform `left == right` operation on [`StringArray`] / [`LargeStringArray`].
-pub fn eq_utf8<OffsetSize: StringOffsetSizeTrait>(
- left: &GenericStringArray<OffsetSize>,
- right: &GenericStringArray<OffsetSize>,
-) -> Result<BooleanArray> {
- compare_op!(left, right, |a, b| a == b)
-}
-
-/// Perform `left == right` operation on [`StringArray`] / [`LargeStringArray`] and a scalar.
-pub fn eq_utf8_scalar<OffsetSize: StringOffsetSizeTrait>(
- left: &GenericStringArray<OffsetSize>,
- right: &str,
-) -> Result<BooleanArray> {
- compare_op_scalar!(left, right, |a, b| a == b)
-}
-
-/// Perform `left != right` operation on [`StringArray`] / [`LargeStringArray`].
-pub fn neq_utf8<OffsetSize: StringOffsetSizeTrait>(
- left: &GenericStringArray<OffsetSize>,
- right: &GenericStringArray<OffsetSize>,
-) -> Result<BooleanArray> {
- compare_op!(left, right, |a, b| a != b)
-}
-
-/// Perform `left != right` operation on [`StringArray`] / [`LargeStringArray`] and a scalar.
-pub fn neq_utf8_scalar<OffsetSize: StringOffsetSizeTrait>(
- left: &GenericStringArray<OffsetSize>,
- right: &str,
-) -> Result<BooleanArray> {
- compare_op_scalar!(left, right, |a, b| a != b)
-}
-
-/// Perform `left < right` operation on [`StringArray`] / [`LargeStringArray`].
-pub fn lt_utf8<OffsetSize: StringOffsetSizeTrait>(
- left: &GenericStringArray<OffsetSize>,
- right: &GenericStringArray<OffsetSize>,
-) -> Result<BooleanArray> {
- compare_op!(left, right, |a, b| a < b)
-}
-
-/// Perform `left < right` operation on [`StringArray`] / [`LargeStringArray`] and a scalar.
-pub fn lt_utf8_scalar<OffsetSize: StringOffsetSizeTrait>(
- left: &GenericStringArray<OffsetSize>,
- right: &str,
-) -> Result<BooleanArray> {
- compare_op_scalar!(left, right, |a, b| a < b)
-}
-
-/// Perform `left <= right` operation on [`StringArray`] / [`LargeStringArray`].
-pub fn lt_eq_utf8<OffsetSize: StringOffsetSizeTrait>(
- left: &GenericStringArray<OffsetSize>,
- right: &GenericStringArray<OffsetSize>,
-) -> Result<BooleanArray> {
- compare_op!(left, right, |a, b| a <= b)
-}
-
-/// Perform `left <= right` operation on [`StringArray`] / [`LargeStringArray`] and a scalar.
-pub fn lt_eq_utf8_scalar<OffsetSize: StringOffsetSizeTrait>(
- left: &GenericStringArray<OffsetSize>,
- right: &str,
-) -> Result<BooleanArray> {
- compare_op_scalar!(left, right, |a, b| a <= b)
-}
-
-/// Perform `left > right` operation on [`StringArray`] / [`LargeStringArray`].
-pub fn gt_utf8<OffsetSize: StringOffsetSizeTrait>(
- left: &GenericStringArray<OffsetSize>,
- right: &GenericStringArray<OffsetSize>,
-) -> Result<BooleanArray> {
- compare_op!(left, right, |a, b| a > b)
-}
-
-/// Perform `left > right` operation on [`StringArray`] / [`LargeStringArray`] and a scalar.
-pub fn gt_utf8_scalar<OffsetSize: StringOffsetSizeTrait>(
- left: &GenericStringArray<OffsetSize>,
- right: &str,
-) -> Result<BooleanArray> {
- compare_op_scalar!(left, right, |a, b| a > b)
-}
-
-/// Perform `left >= right` operation on [`StringArray`] / [`LargeStringArray`].
-pub fn gt_eq_utf8<OffsetSize: StringOffsetSizeTrait>(
- left: &GenericStringArray<OffsetSize>,
- right: &GenericStringArray<OffsetSize>,
-) -> Result<BooleanArray> {
- compare_op!(left, right, |a, b| a >= b)
-}
-
-/// Perform `left >= right` operation on [`StringArray`] / [`LargeStringArray`] and a scalar.
-pub fn gt_eq_utf8_scalar<OffsetSize: StringOffsetSizeTrait>(
- left: &GenericStringArray<OffsetSize>,
- right: &str,
-) -> Result<BooleanArray> {
- compare_op_scalar!(left, right, |a, b| a >= b)
-}
-
-/// Helper function to perform boolean lambda function on values from two arrays using
-/// SIMD.
-#[cfg(feature = "simd")]
-fn simd_compare_op<T, SIMD_OP, SCALAR_OP>(
- left: &PrimitiveArray<T>,
- right: &PrimitiveArray<T>,
- simd_op: SIMD_OP,
- scalar_op: SCALAR_OP,
-) -> Result<BooleanArray>
-where
- T: ArrowNumericType,
- SIMD_OP: Fn(T::Simd, T::Simd) -> T::SimdMask,
- SCALAR_OP: Fn(T::Native, T::Native) -> bool,
-{
- use std::borrow::BorrowMut;
-
- let len = left.len();
- if len != right.len() {
- return Err(ArrowError::ComputeError(
- "Cannot perform comparison operation on arrays of different length"
- .to_string(),
- ));
- }
-
- let null_bit_buffer = combine_option_bitmap(left.data_ref(), right.data_ref(), len)?;
-
- let lanes = T::lanes();
- let buffer_size = bit_util::ceil(len, 8);
- let mut result = MutableBuffer::new(buffer_size).with_bitset(buffer_size, false);
-
- // this is currently the case for all our datatypes and allows us to always append full bytes
- assert!(
- lanes % 8 == 0,
- "Number of vector lanes must be multiple of 8"
- );
- let mut left_chunks = left.values().chunks_exact(lanes);
- let mut right_chunks = right.values().chunks_exact(lanes);
-
- let result_remainder = left_chunks
- .borrow_mut()
- .zip(right_chunks.borrow_mut())
- .fold(
- result.typed_data_mut(),
- |result_slice, (left_slice, right_slice)| {
- let simd_left = T::load(left_slice);
- let simd_right = T::load(right_slice);
- let simd_result = simd_op(simd_left, simd_right);
-
- let bitmask = T::mask_to_u64(&simd_result);
- let bytes = bitmask.to_le_bytes();
- &result_slice[0..lanes / 8].copy_from_slice(&bytes[0..lanes / 8]);
-
- &mut result_slice[lanes / 8..]
- },
- );
-
- let left_remainder = left_chunks.remainder();
- let right_remainder = right_chunks.remainder();
-
- assert_eq!(left_remainder.len(), right_remainder.len());
-
- let remainder_bitmask = left_remainder
- .iter()
- .zip(right_remainder.iter())
- .enumerate()
- .fold(0_u64, |mut mask, (i, (scalar_left, scalar_right))| {
- let bit = if scalar_op(*scalar_left, *scalar_right) {
- 1_u64
- } else {
- 0_u64
- };
- mask |= bit << i;
- mask
- });
- let remainder_mask_as_bytes =
- &remainder_bitmask.to_le_bytes()[0..bit_util::ceil(left_remainder.len(), 8)];
- result_remainder.copy_from_slice(remainder_mask_as_bytes);
-
- let data = ArrayData::new(
- DataType::Boolean,
- len,
- None,
- null_bit_buffer,
- 0,
- vec![result.into()],
- vec![],
- );
- Ok(BooleanArray::from(data))
-}
-
-/// Helper function to perform boolean lambda function on values from an array and a scalar value using
-/// SIMD.
-#[cfg(feature = "simd")]
-fn simd_compare_op_scalar<T, SIMD_OP, SCALAR_OP>(
- left: &PrimitiveArray<T>,
- right: T::Native,
- simd_op: SIMD_OP,
- scalar_op: SCALAR_OP,
-) -> Result<BooleanArray>
-where
- T: ArrowNumericType,
- SIMD_OP: Fn(T::Simd, T::Simd) -> T::SimdMask,
- SCALAR_OP: Fn(T::Native, T::Native) -> bool,
-{
- use std::borrow::BorrowMut;
-
- let len = left.len();
-
- let lanes = T::lanes();
- let buffer_size = bit_util::ceil(len, 8);
- let mut result = MutableBuffer::new(buffer_size).with_bitset(buffer_size, false);
-
- // this is currently the case for all our datatypes and allows us to always append full bytes
- assert!(
- lanes % 8 == 0,
- "Number of vector lanes must be multiple of 8"
- );
- let mut left_chunks = left.values().chunks_exact(lanes);
- let simd_right = T::init(right);
-
- let result_remainder = left_chunks.borrow_mut().fold(
- result.typed_data_mut(),
- |result_slice, left_slice| {
- let simd_left = T::load(left_slice);
- let simd_result = simd_op(simd_left, simd_right);
-
- let bitmask = T::mask_to_u64(&simd_result);
- let bytes = bitmask.to_le_bytes();
- &result_slice[0..lanes / 8].copy_from_slice(&bytes[0..lanes / 8]);
-
- &mut result_slice[lanes / 8..]
- },
- );
-
- let left_remainder = left_chunks.remainder();
-
- let remainder_bitmask =
- left_remainder
- .iter()
- .enumerate()
- .fold(0_u64, |mut mask, (i, scalar_left)| {
- let bit = if scalar_op(*scalar_left, right) {
- 1_u64
- } else {
- 0_u64
- };
- mask |= bit << i;
- mask
- });
- let remainder_mask_as_bytes =
- &remainder_bitmask.to_le_bytes()[0..bit_util::ceil(left_remainder.len(), 8)];
- result_remainder.copy_from_slice(remainder_mask_as_bytes);
-
- let null_bit_buffer = left
- .data_ref()
- .null_buffer()
- .map(|b| b.bit_slice(left.offset(), left.len()));
-
- // null count is the same as in the input since the right side of the scalar comparison cannot be null
- let null_count = left.null_count();
-
- let data = ArrayData::new(
- DataType::Boolean,
- len,
- Some(null_count),
- null_bit_buffer,
- 0,
- vec![result.into()],
- vec![],
- );
- Ok(BooleanArray::from(data))
-}
-
-/// Perform `left == right` operation on two arrays.
-pub fn eq<T>(left: &PrimitiveArray<T>, right: &PrimitiveArray<T>) -> Result<BooleanArray>
-where
- T: ArrowNumericType,
-{
- #[cfg(feature = "simd")]
- return simd_compare_op(left, right, T::eq, |a, b| a == b);
- #[cfg(not(feature = "simd"))]
- return compare_op!(left, right, |a, b| a == b);
-}
-
-/// Perform `left == right` operation on an array and a scalar value.
-pub fn eq_scalar<T>(left: &PrimitiveArray<T>, right: T::Native) -> Result<BooleanArray>
-where
- T: ArrowNumericType,
-{
- #[cfg(feature = "simd")]
- return simd_compare_op_scalar(left, right, T::eq, |a, b| a == b);
- #[cfg(not(feature = "simd"))]
- return compare_op_scalar!(left, right, |a, b| a == b);
-}
-
-/// Perform `left != right` operation on two arrays.
-pub fn neq<T>(left: &PrimitiveArray<T>, right: &PrimitiveArray<T>) -> Result<BooleanArray>
-where
- T: ArrowNumericType,
-{
- #[cfg(feature = "simd")]
- return simd_compare_op(left, right, T::ne, |a, b| a != b);
- #[cfg(not(feature = "simd"))]
- return compare_op!(left, right, |a, b| a != b);
-}
-
-/// Perform `left != right` operation on an array and a scalar value.
-pub fn neq_scalar<T>(left: &PrimitiveArray<T>, right: T::Native) -> Result<BooleanArray>
-where
- T: ArrowNumericType,
-{
- #[cfg(feature = "simd")]
- return simd_compare_op_scalar(left, right, T::ne, |a, b| a != b);
- #[cfg(not(feature = "simd"))]
- return compare_op_scalar!(left, right, |a, b| a != b);
-}
-
-/// Perform `left < right` operation on two arrays. Null values are less than non-null
-/// values.
-pub fn lt<T>(left: &PrimitiveArray<T>, right: &PrimitiveArray<T>) -> Result<BooleanArray>
-where
- T: ArrowNumericType,
-{
- #[cfg(feature = "simd")]
- return simd_compare_op(left, right, T::lt, |a, b| a < b);
- #[cfg(not(feature = "simd"))]
- return compare_op!(left, right, |a, b| a < b);
-}
-
-/// Perform `left < right` operation on an array and a scalar value.
-/// Null values are less than non-null values.
-pub fn lt_scalar<T>(left: &PrimitiveArray<T>, right: T::Native) -> Result<BooleanArray>
-where
- T: ArrowNumericType,
-{
- #[cfg(feature = "simd")]
- return simd_compare_op_scalar(left, right, T::lt, |a, b| a < b);
- #[cfg(not(feature = "simd"))]
- return compare_op_scalar!(left, right, |a, b| a < b);
-}
-
-/// Perform `left <= right` operation on two arrays. Null values are less than non-null
-/// values.
-pub fn lt_eq<T>(
- left: &PrimitiveArray<T>,
- right: &PrimitiveArray<T>,
-) -> Result<BooleanArray>
-where
- T: ArrowNumericType,
-{
- #[cfg(feature = "simd")]
- return simd_compare_op(left, right, T::le, |a, b| a <= b);
- #[cfg(not(feature = "simd"))]
- return compare_op!(left, right, |a, b| a <= b);
-}
-
-/// Perform `left <= right` operation on an array and a scalar value.
-/// Null values are less than non-null values.
-pub fn lt_eq_scalar<T>(left: &PrimitiveArray<T>, right: T::Native) -> Result<BooleanArray>
-where
- T: ArrowNumericType,
-{
- #[cfg(feature = "simd")]
- return simd_compare_op_scalar(left, right, T::le, |a, b| a <= b);
- #[cfg(not(feature = "simd"))]
- return compare_op_scalar!(left, right, |a, b| a <= b);
-}
-
-/// Perform `left > right` operation on two arrays. Non-null values are greater than null
-/// values.
-pub fn gt<T>(left: &PrimitiveArray<T>, right: &PrimitiveArray<T>) -> Result<BooleanArray>
-where
- T: ArrowNumericType,
-{
- #[cfg(feature = "simd")]
- return simd_compare_op(left, right, T::gt, |a, b| a > b);
- #[cfg(not(feature = "simd"))]
- return compare_op!(left, right, |a, b| a > b);
-}
-
-/// Perform `left > right` operation on an array and a scalar value.
-/// Non-null values are greater than null values.
-pub fn gt_scalar<T>(left: &PrimitiveArray<T>, right: T::Native) -> Result<BooleanArray>
-where
- T: ArrowNumericType,
-{
- #[cfg(feature = "simd")]
- return simd_compare_op_scalar(left, right, T::gt, |a, b| a > b);
- #[cfg(not(feature = "simd"))]
- return compare_op_scalar!(left, right, |a, b| a > b);
-}
-
-/// Perform `left >= right` operation on two arrays. Non-null values are greater than null
-/// values.
-pub fn gt_eq<T>(
- left: &PrimitiveArray<T>,
- right: &PrimitiveArray<T>,
-) -> Result<BooleanArray>
-where
- T: ArrowNumericType,
-{
- #[cfg(feature = "simd")]
- return simd_compare_op(left, right, T::ge, |a, b| a >= b);
- #[cfg(not(feature = "simd"))]
- return compare_op!(left, right, |a, b| a >= b);
-}
-
-/// Perform `left >= right` operation on an array and a scalar value.
-/// Non-null values are greater than null values.
-pub fn gt_eq_scalar<T>(left: &PrimitiveArray<T>, right: T::Native) -> Result<BooleanArray>
-where
- T: ArrowNumericType,
-{
- #[cfg(feature = "simd")]
- return simd_compare_op_scalar(left, right, T::ge, |a, b| a >= b);
- #[cfg(not(feature = "simd"))]
- return compare_op_scalar!(left, right, |a, b| a >= b);
-}
-
-/// Checks if a [`GenericListArray`] contains a value in the [`PrimitiveArray`]
-pub fn contains<T, OffsetSize>(
- left: &PrimitiveArray<T>,
- right: &GenericListArray<OffsetSize>,
-) -> Result<BooleanArray>
-where
- T: ArrowNumericType,
- OffsetSize: OffsetSizeTrait,
-{
- let left_len = left.len();
- if left_len != right.len() {
- return Err(ArrowError::ComputeError(
- "Cannot perform comparison operation on arrays of different length"
- .to_string(),
- ));
- }
-
- let num_bytes = bit_util::ceil(left_len, 8);
-
- let not_both_null_bit_buffer =
- match combine_option_bitmap(left.data_ref(), right.data_ref(), left_len)? {
- Some(buff) => buff,
- None => new_all_set_buffer(num_bytes),
- };
- let not_both_null_bitmap = not_both_null_bit_buffer.as_slice();
-
- let mut bool_buf = MutableBuffer::from_len_zeroed(num_bytes);
- let bool_slice = bool_buf.as_slice_mut();
-
- // if both array slots are valid, check if list contains primitive
- for i in 0..left_len {
- if bit_util::get_bit(not_both_null_bitmap, i) {
- let list = right.value(i);
- let list = list.as_any().downcast_ref::<PrimitiveArray<T>>().unwrap();
-
- for j in 0..list.len() {
- if list.is_valid(j) && (left.value(i) == list.value(j)) {
- bit_util::set_bit(bool_slice, i);
- continue;
- }
- }
- }
- }
-
- let data = ArrayData::new(
- DataType::Boolean,
- left.len(),
- None,
- None,
- 0,
- vec![bool_buf.into()],
- vec![],
- );
- Ok(BooleanArray::from(data))
-}
-
-/// Checks if a [`GenericListArray`] contains a value in the [`GenericStringArray`]
-pub fn contains_utf8<OffsetSize>(
- left: &GenericStringArray<OffsetSize>,
- right: &ListArray,
-) -> Result<BooleanArray>
-where
- OffsetSize: StringOffsetSizeTrait,
-{
- let left_len = left.len();
- if left_len != right.len() {
- return Err(ArrowError::ComputeError(
- "Cannot perform comparison operation on arrays of different length"
- .to_string(),
- ));
- }
-
- let num_bytes = bit_util::ceil(left_len, 8);
-
- let not_both_null_bit_buffer =
- match combine_option_bitmap(left.data_ref(), right.data_ref(), left_len)? {
- Some(buff) => buff,
- None => new_all_set_buffer(num_bytes),
- };
- let not_both_null_bitmap = not_both_null_bit_buffer.as_slice();
-
- let mut bool_buf = MutableBuffer::from_len_zeroed(num_bytes);
- let bool_slice = &mut bool_buf;
-
- for i in 0..left_len {
- // contains(null, null) = false
- if bit_util::get_bit(not_both_null_bitmap, i) {
- let list = right.value(i);
- let list = list
- .as_any()
- .downcast_ref::<GenericStringArray<OffsetSize>>()
- .unwrap();
-
- for j in 0..list.len() {
- if list.is_valid(j) && (left.value(i) == list.value(j)) {
- bit_util::set_bit(bool_slice, i);
- continue;
- }
- }
- }
- }
-
- let data = ArrayData::new(
- DataType::Boolean,
- left.len(),
- None,
- None,
- 0,
- vec![bool_buf.into()],
- vec![],
- );
- Ok(BooleanArray::from(data))
-}
-
-// create a buffer and fill it with valid bits
-#[inline]
-fn new_all_set_buffer(len: usize) -> Buffer {
- let buffer = MutableBuffer::new(len);
- let buffer = buffer.with_bitset(len, true);
-
- buffer.into()
-}
-
-// disable wrapping inside literal vectors used for test data and assertions
-#[rustfmt::skip::macros(vec)]
-#[cfg(test)]
-mod tests {
- use super::*;
- use crate::datatypes::Int8Type;
- use crate::{array::Int32Array, array::Int64Array, datatypes::Field};
-
- /// Evaluate `KERNEL` with two vectors as inputs and assert against the expected output.
- /// `A_VEC` and `B_VEC` can be of type `Vec<i64>` or `Vec<Option<i64>>`.
- /// `EXPECTED` can be either `Vec<bool>` or `Vec<Option<bool>>`.
- /// The main reason for this macro is that inputs and outputs align nicely after `cargo fmt`.
- macro_rules! cmp_i64 {
- ($KERNEL:ident, $A_VEC:expr, $B_VEC:expr, $EXPECTED:expr) => {
- let a = Int64Array::from($A_VEC);
- let b = Int64Array::from($B_VEC);
- let c = $KERNEL(&a, &b).unwrap();
- assert_eq!(BooleanArray::from($EXPECTED), c);
- };
- }
-
- /// Evaluate `KERNEL` with one vectors and one scalar as inputs and assert against the expected output.
- /// `A_VEC` can be of type `Vec<i64>` or `Vec<Option<i64>>`.
- /// `EXPECTED` can be either `Vec<bool>` or `Vec<Option<bool>>`.
- /// The main reason for this macro is that inputs and outputs align nicely after `cargo fmt`.
- macro_rules! cmp_i64_scalar {
- ($KERNEL:ident, $A_VEC:expr, $B:literal, $EXPECTED:expr) => {
- let a = Int64Array::from($A_VEC);
- let c = $KERNEL(&a, $B).unwrap();
- assert_eq!(BooleanArray::from($EXPECTED), c);
- };
- }
-
- #[test]
- fn test_primitive_array_eq() {
- cmp_i64!(
- eq,
- vec![8, 8, 8, 8, 8, 8, 8, 8, 8, 8],
- vec![6, 7, 8, 9, 10, 6, 7, 8, 9, 10],
- vec![false, false, true, false, false, false, false, true, false, false]
- );
- }
-
- #[test]
- fn test_primitive_array_eq_scalar() {
- cmp_i64_scalar!(
- eq_scalar,
- vec![6, 7, 8, 9, 10, 6, 7, 8, 9, 10],
- 8,
- vec![false, false, true, false, false, false, false, true, false, false]
- );
- }
-
- #[test]
- fn test_primitive_array_eq_with_slice() {
- let a = Int32Array::from(vec![6, 7, 8, 8, 10]);
- let b = Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10]);
- let b_slice = b.slice(5, 5);
- let c = b_slice.as_any().downcast_ref().unwrap();
- let d = eq(&c, &a).unwrap();
- assert_eq!(true, d.value(0));
- assert_eq!(true, d.value(1));
- assert_eq!(true, d.value(2));
- assert_eq!(false, d.value(3));
- assert_eq!(true, d.value(4));
- }
-
- #[test]
- fn test_primitive_array_neq() {
- cmp_i64!(
- neq,
- vec![8, 8, 8, 8, 8, 8, 8, 8, 8, 8],
- vec![6, 7, 8, 9, 10, 6, 7, 8, 9, 10],
- vec![true, true, false, true, true, true, true, false, true, true]
- );
- }
-
- #[test]
- fn test_primitive_array_neq_scalar() {
- cmp_i64_scalar!(
- neq_scalar,
- vec![6, 7, 8, 9, 10, 6, 7, 8, 9, 10],
- 8,
- vec![true, true, false, true, true, true, true, false, true, true]
- );
- }
-
- #[test]
- fn test_primitive_array_lt() {
- cmp_i64!(
- lt,
- vec![8, 8, 8, 8, 8, 8, 8, 8, 8, 8],
- vec![6, 7, 8, 9, 10, 6, 7, 8, 9, 10],
- vec![false, false, false, true, true, false, false, false, true, true]
- );
- }
-
- #[test]
- fn test_primitive_array_lt_scalar() {
- cmp_i64_scalar!(
- lt_scalar,
- vec![6, 7, 8, 9, 10, 6, 7, 8, 9, 10],
- 8,
- vec![true, true, false, false, false, true, true, false, false, false]
- );
- }
-
- #[test]
- fn test_primitive_array_lt_nulls() {
- cmp_i64!(
- lt,
- vec![None, None, Some(1), Some(1), None, None, Some(2), Some(2),],
- vec![None, Some(1), None, Some(1), None, Some(3), None, Some(3),],
- vec![None, None, None, Some(false), None, None, None, Some(true)]
- );
- }
-
- #[test]
- fn test_primitive_array_lt_scalar_nulls() {
- cmp_i64_scalar!(
- lt_scalar,
- vec![None, Some(1), Some(2), Some(3), None, Some(1), Some(2), Some(3), Some(2), None],
- 2,
- vec![None, Some(true), Some(false), Some(false), None, Some(true), Some(false), Some(false), Some(false), None]
- );
- }
-
- #[test]
- fn test_primitive_array_lt_eq() {
- cmp_i64!(
- lt_eq,
- vec![8, 8, 8, 8, 8, 8, 8, 8, 8, 8],
- vec![6, 7, 8, 9, 10, 6, 7, 8, 9, 10],
- vec![false, false, true, true, true, false, false, true, true, true]
- );
- }
-
- #[test]
- fn test_primitive_array_lt_eq_scalar() {
- cmp_i64_scalar!(
- lt_eq_scalar,
- vec![6, 7, 8, 9, 10, 6, 7, 8, 9, 10],
- 8,
- vec![true, true, true, false, false, true, true, true, false, false]
- );
- }
-
- #[test]
- fn test_primitive_array_lt_eq_nulls() {
- cmp_i64!(
- lt_eq,
- vec![None, None, Some(1), None, None, Some(1), None, None, Some(1)],
- vec![None, Some(1), Some(0), None, Some(1), Some(2), None, None, Some(3)],
- vec![None, None, Some(false), None, None, Some(true), None, None, Some(true)]
- );
- }
-
- #[test]
- fn test_primitive_array_lt_eq_scalar_nulls() {
- cmp_i64_scalar!(
- lt_eq_scalar,
- vec![None, Some(1), Some(2), None, Some(1), Some(2), None, Some(1), Some(2)],
- 1,
- vec![None, Some(true), Some(false), None, Some(true), Some(false), None, Some(true), Some(false)]
- );
- }
-
- #[test]
- fn test_primitive_array_gt() {
- cmp_i64!(
- gt,
- vec![8, 8, 8, 8, 8, 8, 8, 8, 8, 8],
- vec![6, 7, 8, 9, 10, 6, 7, 8, 9, 10],
- vec![true, true, false, false, false, true, true, false, false, false]
- );
- }
-
- #[test]
- fn test_primitive_array_gt_scalar() {
- cmp_i64_scalar!(
- gt_scalar,
- vec![6, 7, 8, 9, 10, 6, 7, 8, 9, 10],
- 8,
- vec![false, false, false, true, true, false, false, false, true, true]
- );
- }
-
- #[test]
- fn test_primitive_array_gt_nulls() {
- cmp_i64!(
- gt,
- vec![None, None, Some(1), None, None, Some(2), None, None, Some(3)],
- vec![None, Some(1), Some(1), None, Some(1), Some(1), None, Some(1), Some(1)],
- vec![None, None, Some(false), None, None, Some(true), None, None, Some(true)]
- );
- }
-
- #[test]
- fn test_primitive_array_gt_scalar_nulls() {
- cmp_i64_scalar!(
- gt_scalar,
- vec![None, Some(1), Some(2), None, Some(1), Some(2), None, Some(1), Some(2)],
- 1,
- vec![None, Some(false), Some(true), None, Some(false), Some(true), None, Some(false), Some(true)]
- );
- }
-
- #[test]
- fn test_primitive_array_gt_eq() {
- cmp_i64!(
- gt_eq,
- vec![8, 8, 8, 8, 8, 8, 8, 8, 8, 8],
- vec![6, 7, 8, 9, 10, 6, 7, 8, 9, 10],
- vec![true, true, true, false, false, true, true, true, false, false]
- );
- }
-
- #[test]
- fn test_primitive_array_gt_eq_scalar() {
- cmp_i64_scalar!(
- gt_eq_scalar,
- vec![6, 7, 8, 9, 10, 6, 7, 8, 9, 10],
- 8,
- vec![false, false, true, true, true, false, false, true, true, true]
- );
- }
-
- #[test]
- fn test_primitive_array_gt_eq_nulls() {
- cmp_i64!(
- gt_eq,
- vec![None, None, Some(1), None, Some(1), Some(2), None, None, Some(1)],
- vec![None, Some(1), None, None, Some(1), Some(1), None, Some(2), Some(2)],
- vec![None, None, None, None, Some(true), Some(true), None, None, Some(false)]
- );
- }
-
- #[test]
- fn test_primitive_array_gt_eq_scalar_nulls() {
- cmp_i64_scalar!(
- gt_eq_scalar,
- vec![None, Some(1), Some(2), None, Some(2), Some(3), None, Some(3), Some(4)],
- 2,
- vec![None, Some(false), Some(true), None, Some(true), Some(true), None, Some(true), Some(true)]
- );
- }
-
- #[test]
- fn test_primitive_array_compare_slice() {
- let a: Int32Array = (0..100).map(Some).collect();
- let a = a.slice(50, 50);
- let a = a.as_any().downcast_ref::<Int32Array>().unwrap();
- let b: Int32Array = (100..200).map(Some).collect();
- let b = b.slice(50, 50);
- let b = b.as_any().downcast_ref::<Int32Array>().unwrap();
- let actual = lt(&a, &b).unwrap();
- let expected: BooleanArray = (0..50).map(|_| Some(true)).collect();
- assert_eq!(expected, actual);
- }
-
- #[test]
- fn test_primitive_array_compare_scalar_slice() {
- let a: Int32Array = (0..100).map(Some).collect();
- let a = a.slice(50, 50);
- let a = a.as_any().downcast_ref::<Int32Array>().unwrap();
- let actual = lt_scalar(&a, 200).unwrap();
- let expected: BooleanArray = (0..50).map(|_| Some(true)).collect();
- assert_eq!(expected, actual);
- }
-
- #[test]
- fn test_length_of_result_buffer() {
- // `item_count` is chosen to not be a multiple of the number of SIMD lanes for this
- // type (`Int8Type`), 64.
- let item_count = 130;
-
- let select_mask: BooleanArray = vec![true; item_count].into();
-
- let array_a: PrimitiveArray<Int8Type> = vec![1; item_count].into();
- let array_b: PrimitiveArray<Int8Type> = vec![2; item_count].into();
- let result_mask = gt_eq(&array_a, &array_b).unwrap();
-
- assert_eq!(
- result_mask.data().buffers()[0].len(),
- select_mask.data().buffers()[0].len()
- );
- }
-
- // Expected behaviour:
- // contains(1, [1, 2, null]) = true
- // contains(3, [1, 2, null]) = false
- // contains(null, [1, 2, null]) = false
- // contains(null, null) = false
- #[test]
- fn test_contains() {
- let value_data = Int32Array::from(vec![
- Some(0),
- Some(1),
- Some(2),
- Some(3),
- Some(4),
- Some(5),
- Some(6),
- None,
- Some(7),
- ])
- .data()
- .clone();
- let value_offsets = Buffer::from_slice_ref(&[0i64, 3, 6, 6, 9]);
- let list_data_type =
- DataType::LargeList(Box::new(Field::new("item", DataType::Int32, true)));
- let list_data = ArrayData::builder(list_data_type)
- .len(4)
- .add_buffer(value_offsets)
- .add_child_data(value_data)
- .null_bit_buffer(Buffer::from([0b00001011]))
- .build();
-
- // [[0, 1, 2], [3, 4, 5], null, [6, null, 7]]
- let list_array = LargeListArray::from(list_data);
-
- let nulls = Int32Array::from(vec![None, None, None, None]);
- let nulls_result = contains(&nulls, &list_array).unwrap();
- assert_eq!(
- nulls_result
- .as_any()
- .downcast_ref::<BooleanArray>()
- .unwrap(),
- &BooleanArray::from(vec![false, false, false, false]),
- );
-
- let values = Int32Array::from(vec![Some(0), Some(0), Some(0), Some(0)]);
- let values_result = contains(&values, &list_array).unwrap();
- assert_eq!(
- values_result
- .as_any()
- .downcast_ref::<BooleanArray>()
- .unwrap(),
- &BooleanArray::from(vec![true, false, false, false]),
- );
- }
-
- // Expected behaviour:
- // contains("ab", ["ab", "cd", null]) = true
- // contains("ef", ["ab", "cd", null]) = false
- // contains(null, ["ab", "cd", null]) = false
- // contains(null, null) = false
- #[test]
- fn test_contains_utf8() {
- let values_builder = StringBuilder::new(10);
- let mut builder = ListBuilder::new(values_builder);
-
- builder.values().append_value("Lorem").unwrap();
- builder.values().append_value("ipsum").unwrap();
- builder.values().append_null().unwrap();
- builder.append(true).unwrap();
- builder.values().append_value("sit").unwrap();
- builder.values().append_value("amet").unwrap();
- builder.values().append_value("Lorem").unwrap();
- builder.append(true).unwrap();
- builder.append(false).unwrap();
- builder.values().append_value("ipsum").unwrap();
- builder.append(true).unwrap();
-
- // [["Lorem", "ipsum", null], ["sit", "amet", "Lorem"], null, ["ipsum"]]
- // value_offsets = [0, 3, 6, 6]
- let list_array = builder.finish();
-
- let nulls = StringArray::from(vec![None, None, None, None]);
- let nulls_result = contains_utf8(&nulls, &list_array).unwrap();
- assert_eq!(
- nulls_result
- .as_any()
- .downcast_ref::<BooleanArray>()
- .unwrap(),
- &BooleanArray::from(vec![false, false, false, false]),
- );
-
- let values = StringArray::from(vec![
- Some("Lorem"),
- Some("Lorem"),
- Some("Lorem"),
- Some("Lorem"),
- ]);
- let values_result = contains_utf8(&values, &list_array).unwrap();
- assert_eq!(
- values_result
- .as_any()
- .downcast_ref::<BooleanArray>()
- .unwrap(),
- &BooleanArray::from(vec![true, true, false, false]),
- );
- }
-
- macro_rules! test_utf8 {
- ($test_name:ident, $left:expr, $right:expr, $op:expr, $expected:expr) => {
- #[test]
- fn $test_name() {
- let left = StringArray::from($left);
- let right = StringArray::from($right);
- let res = $op(&left, &right).unwrap();
- let expected = $expected;
- assert_eq!(expected.len(), res.len());
- for i in 0..res.len() {
- let v = res.value(i);
- assert_eq!(v, expected[i]);
- }
- }
- };
- }
-
- macro_rules! test_utf8_scalar {
- ($test_name:ident, $left:expr, $right:expr, $op:expr, $expected:expr) => {
- #[test]
- fn $test_name() {
- let left = StringArray::from($left);
- let res = $op(&left, $right).unwrap();
- let expected = $expected;
- assert_eq!(expected.len(), res.len());
- for i in 0..res.len() {
- let v = res.value(i);
- assert_eq!(
- v,
- expected[i],
- "unexpected result when comparing {} at position {} to {} ",
- left.value(i),
- i,
- $right
- );
- }
-
- let left = LargeStringArray::from($left);
- let res = $op(&left, $right).unwrap();
- let expected = $expected;
- assert_eq!(expected.len(), res.len());
- for i in 0..res.len() {
- let v = res.value(i);
- assert_eq!(
- v,
- expected[i],
- "unexpected result when comparing {} at position {} to {} ",
- left.value(i),
- i,
- $right
- );
- }
- }
- };
- }
-
- test_utf8!(
- test_utf8_array_like,
- vec!["arrow", "arrow", "arrow", "arrow", "arrow", "arrows", "arrow"],
- vec!["arrow", "ar%", "%ro%", "foo", "arr", "arrow_", "arrow_"],
- like_utf8,
- vec![true, true, true, false, false, true, false]
- );
-
- test_utf8_scalar!(
- test_utf8_array_like_scalar,
- vec!["arrow", "parquet", "datafusion", "flight"],
- "%ar%",
- like_utf8_scalar,
- vec![true, true, false, false]
- );
- test_utf8_scalar!(
- test_utf8_array_like_scalar_start,
- vec!["arrow", "parrow", "arrows", "arr"],
- "arrow%",
- like_utf8_scalar,
- vec![true, false, true, false]
- );
-
- test_utf8_scalar!(
- test_utf8_array_like_scalar_end,
- vec!["arrow", "parrow", "arrows", "arr"],
- "%arrow",
- like_utf8_scalar,
- vec![true, true, false, false]
- );
-
- test_utf8_scalar!(
- test_utf8_array_like_scalar_equals,
- vec!["arrow", "parrow", "arrows", "arr"],
- "arrow",
- like_utf8_scalar,
- vec![true, false, false, false]
- );
-
- test_utf8_scalar!(
- test_utf8_array_like_scalar_one,
- vec!["arrow", "arrows", "parrow", "arr"],
- "arrow_",
- like_utf8_scalar,
- vec![false, true, false, false]
- );
-
- test_utf8!(
- test_utf8_array_nlike,
- vec!["arrow", "arrow", "arrow", "arrow", "arrow", "arrows", "arrow"],
- vec!["arrow", "ar%", "%ro%", "foo", "arr", "arrow_", "arrow_"],
- nlike_utf8,
- vec![false, false, false, true, true, false, true]
- );
- test_utf8_scalar!(
- test_utf8_array_nlike_scalar,
- vec!["arrow", "parquet", "datafusion", "flight"],
- "%ar%",
- nlike_utf8_scalar,
- vec![false, false, true, true]
- );
-
- test_utf8!(
- test_utf8_array_eq,
- vec!["arrow", "arrow", "arrow", "arrow"],
- vec!["arrow", "parquet", "datafusion", "flight"],
- eq_utf8,
- vec![true, false, false, false]
- );
- test_utf8_scalar!(
- test_utf8_array_eq_scalar,
- vec!["arrow", "parquet", "datafusion", "flight"],
- "arrow",
- eq_utf8_scalar,
- vec![true, false, false, false]
- );
-
- test_utf8_scalar!(
- test_utf8_array_nlike_scalar_start,
- vec!["arrow", "parrow", "arrows", "arr"],
- "arrow%",
- nlike_utf8_scalar,
- vec![false, true, false, true]
- );
-
- test_utf8_scalar!(
- test_utf8_array_nlike_scalar_end,
- vec!["arrow", "parrow", "arrows", "arr"],
- "%arrow",
- nlike_utf8_scalar,
- vec![false, false, true, true]
- );
-
- test_utf8_scalar!(
- test_utf8_array_nlike_scalar_equals,
- vec!["arrow", "parrow", "arrows", "arr"],
- "arrow",
- nlike_utf8_scalar,
- vec![false, true, true, true]
- );
-
- test_utf8_scalar!(
- test_utf8_array_nlike_scalar_one,
- vec!["arrow", "arrows", "parrow", "arr"],
- "arrow_",
- nlike_utf8_scalar,
- vec![true, false, true, true]
- );
-
- test_utf8!(
- test_utf8_array_neq,
- vec!["arrow", "arrow", "arrow", "arrow"],
- vec!["arrow", "parquet", "datafusion", "flight"],
- neq_utf8,
- vec![false, true, true, true]
- );
- test_utf8_scalar!(
- test_utf8_array_neq_scalar,
- vec!["arrow", "parquet", "datafusion", "flight"],
- "arrow",
- neq_utf8_scalar,
- vec![false, true, true, true]
- );
-
- test_utf8!(
- test_utf8_array_lt,
- vec!["arrow", "datafusion", "flight", "parquet"],
- vec!["flight", "flight", "flight", "flight"],
- lt_utf8,
- vec![true, true, false, false]
- );
- test_utf8_scalar!(
- test_utf8_array_lt_scalar,
- vec!["arrow", "datafusion", "flight", "parquet"],
- "flight",
- lt_utf8_scalar,
- vec![true, true, false, false]
- );
-
- test_utf8!(
- test_utf8_array_lt_eq,
- vec!["arrow", "datafusion", "flight", "parquet"],
- vec!["flight", "flight", "flight", "flight"],
- lt_eq_utf8,
- vec![true, true, true, false]
- );
- test_utf8_scalar!(
- test_utf8_array_lt_eq_scalar,
- vec!["arrow", "datafusion", "flight", "parquet"],
- "flight",
- lt_eq_utf8_scalar,
- vec![true, true, true, false]
- );
-
- test_utf8!(
- test_utf8_array_gt,
- vec!["arrow", "datafusion", "flight", "parquet"],
- vec!["flight", "flight", "flight", "flight"],
- gt_utf8,
- vec![false, false, false, true]
- );
- test_utf8_scalar!(
- test_utf8_array_gt_scalar,
- vec!["arrow", "datafusion", "flight", "parquet"],
- "flight",
- gt_utf8_scalar,
- vec![false, false, false, true]
- );
-
- test_utf8!(
- test_utf8_array_gt_eq,
- vec!["arrow", "datafusion", "flight", "parquet"],
- vec!["flight", "flight", "flight", "flight"],
- gt_eq_utf8,
- vec![false, false, true, true]
- );
- test_utf8_scalar!(
- test_utf8_array_gt_eq_scalar,
- vec!["arrow", "datafusion", "flight", "parquet"],
- "flight",
- gt_eq_utf8_scalar,
- vec![false, false, true, true]
- );
-}
diff --git a/arrow/src/compute/kernels/concat.rs b/arrow/src/compute/kernels/concat.rs
deleted file mode 100644
index cc976a4..0000000
--- a/arrow/src/compute/kernels/concat.rs
+++ /dev/null
@@ -1,510 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-//! Defines concat kernel for `ArrayRef`
-//!
-//! Example:
-//!
-//! ```
-//! use arrow::array::{ArrayRef, StringArray};
-//! use arrow::compute::concat;
-//!
-//! let arr = concat(&[
-//! &StringArray::from(vec!["hello", "world"]),
-//! &StringArray::from(vec!["!"]),
-//! ]).unwrap();
-//! assert_eq!(arr.len(), 3);
-//! ```
-
-use crate::array::*;
-use crate::datatypes::DataType;
-use crate::error::{ArrowError, Result};
-
-fn compute_str_values_length<Offset: StringOffsetSizeTrait>(
- arrays: &[&ArrayData],
-) -> usize {
- arrays
- .iter()
- .map(|&data| {
- // get the length of the value buffer
- let buf_len = data.buffers()[1].len();
- // find the offset of the buffer
- // this returns a slice of offsets, starting from the offset of the array
- // so we can take the first value
- let offset = data.buffer::<Offset>(0)[0];
- buf_len - offset.to_usize().unwrap()
- })
- .sum()
-}
-
-/// Concatenate multiple [Array] of the same type into a single [ArrayRef].
-pub fn concat(arrays: &[&Array]) -> Result<ArrayRef> {
- if arrays.is_empty() {
- return Err(ArrowError::ComputeError(
- "concat requires input of at least one array".to_string(),
- ));
- }
-
- if arrays
- .iter()
- .any(|array| array.data_type() != arrays[0].data_type())
- {
- return Err(ArrowError::InvalidArgumentError(
- "It is not possible to concatenate arrays of different data types."
- .to_string(),
- ));
- }
-
- let lengths = arrays.iter().map(|array| array.len()).collect::<Vec<_>>();
- let capacity = lengths.iter().sum();
-
- let arrays = arrays.iter().map(|a| a.data()).collect::<Vec<_>>();
-
- let mut mutable = match arrays[0].data_type() {
- DataType::Utf8 => {
- let str_values_size = compute_str_values_length::<i32>(&arrays);
- MutableArrayData::with_capacities(
- arrays,
- false,
- Capacities::Binary(capacity, Some(str_values_size)),
- )
- }
- DataType::LargeUtf8 => {
- let str_values_size = compute_str_values_length::<i64>(&arrays);
- MutableArrayData::with_capacities(
- arrays,
- false,
- Capacities::Binary(capacity, Some(str_values_size)),
- )
- }
- _ => MutableArrayData::new(arrays, false, capacity),
- };
-
- for (i, len) in lengths.iter().enumerate() {
- mutable.extend(i, 0, *len)
- }
-
- Ok(make_array(mutable.freeze()))
-}
-
-#[cfg(test)]
-mod tests {
- use super::*;
- use crate::datatypes::*;
- use std::sync::Arc;
-
- #[test]
- fn test_concat_empty_vec() {
- let re = concat(&[]);
- assert!(re.is_err());
- }
-
- #[test]
- fn test_concat_incompatible_datatypes() {
- let re = concat(&[
- &PrimitiveArray::<Int64Type>::from(vec![Some(-1), Some(2), None]),
- &StringArray::from(vec![Some("hello"), Some("bar"), Some("world")]),
- ]);
- assert!(re.is_err());
- }
-
- #[test]
- fn test_concat_string_arrays() -> Result<()> {
- let arr = concat(&[
- &StringArray::from(vec!["hello", "world"]),
- &StringArray::from(vec!["2", "3", "4"]),
- &StringArray::from(vec![Some("foo"), Some("bar"), None, Some("baz")]),
- ])?;
-
- let expected_output = Arc::new(StringArray::from(vec![
- Some("hello"),
- Some("world"),
- Some("2"),
- Some("3"),
- Some("4"),
- Some("foo"),
- Some("bar"),
- None,
- Some("baz"),
- ])) as ArrayRef;
-
- assert_eq!(&arr, &expected_output);
-
- Ok(())
- }
-
- #[test]
- fn test_concat_primitive_arrays() -> Result<()> {
- let arr = concat(&[
- &PrimitiveArray::<Int64Type>::from(vec![
- Some(-1),
- Some(-1),
- Some(2),
- None,
- None,
- ]),
- &PrimitiveArray::<Int64Type>::from(vec![
- Some(101),
- Some(102),
- Some(103),
- None,
- ]),
- &PrimitiveArray::<Int64Type>::from(vec![Some(256), Some(512), Some(1024)]),
- ])?;
-
- let expected_output = Arc::new(PrimitiveArray::<Int64Type>::from(vec![
- Some(-1),
- Some(-1),
- Some(2),
- None,
- None,
- Some(101),
- Some(102),
- Some(103),
- None,
- Some(256),
- Some(512),
- Some(1024),
- ])) as ArrayRef;
-
- assert_eq!(&arr, &expected_output);
-
- Ok(())
- }
-
- #[test]
- fn test_concat_primitive_array_slices() -> Result<()> {
- let input_1 = PrimitiveArray::<Int64Type>::from(vec![
- Some(-1),
- Some(-1),
- Some(2),
- None,
- None,
- ])
- .slice(1, 3);
-
- let input_2 = PrimitiveArray::<Int64Type>::from(vec![
- Some(101),
- Some(102),
- Some(103),
- None,
- ])
- .slice(1, 3);
- let arr = concat(&[input_1.as_ref(), input_2.as_ref()])?;
-
- let expected_output = Arc::new(PrimitiveArray::<Int64Type>::from(vec![
- Some(-1),
- Some(2),
- None,
- Some(102),
- Some(103),
- None,
- ])) as ArrayRef;
-
- assert_eq!(&arr, &expected_output);
-
- Ok(())
- }
-
- #[test]
- fn test_concat_boolean_primitive_arrays() -> Result<()> {
- let arr = concat(&[
- &BooleanArray::from(vec![
- Some(true),
- Some(true),
- Some(false),
- None,
- None,
- Some(false),
- ]),
- &BooleanArray::from(vec![None, Some(false), Some(true), Some(false)]),
- ])?;
-
- let expected_output = Arc::new(BooleanArray::from(vec![
- Some(true),
- Some(true),
- Some(false),
- None,
- None,
- Some(false),
- None,
- Some(false),
- Some(true),
- Some(false),
- ])) as ArrayRef;
-
- assert_eq!(&arr, &expected_output);
-
- Ok(())
- }
-
- #[test]
- fn test_concat_primitive_list_arrays() -> Result<()> {
- let list1 = vec![
- Some(vec![Some(-1), Some(-1), Some(2), None, None]),
- Some(vec![]),
- None,
- Some(vec![Some(10)]),
- ];
- let list1_array =
- ListArray::from_iter_primitive::<Int64Type, _, _>(list1.clone());
-
- let list2 = vec![
- None,
- Some(vec![Some(100), None, Some(101)]),
- Some(vec![Some(102)]),
- ];
- let list2_array =
- ListArray::from_iter_primitive::<Int64Type, _, _>(list2.clone());
-
- let list3 = vec![Some(vec![Some(1000), Some(1001)])];
- let list3_array =
- ListArray::from_iter_primitive::<Int64Type, _, _>(list3.clone());
-
- let array_result = concat(&[&list1_array, &list2_array, &list3_array])?;
-
- let expected = list1
- .into_iter()
- .chain(list2.into_iter())
- .chain(list3.into_iter());
- let array_expected = ListArray::from_iter_primitive::<Int64Type, _, _>(expected);
-
- assert_eq!(array_result.as_ref(), &array_expected as &dyn Array);
-
- Ok(())
- }
-
- #[test]
- fn test_concat_struct_arrays() -> Result<()> {
- let field = Field::new("field", DataType::Int64, true);
- let input_primitive_1: ArrayRef =
- Arc::new(PrimitiveArray::<Int64Type>::from(vec![
- Some(-1),
- Some(-1),
- Some(2),
- None,
- None,
- ]));
- let input_struct_1 = StructArray::from(vec![(field.clone(), input_primitive_1)]);
-
- let input_primitive_2: ArrayRef =
- Arc::new(PrimitiveArray::<Int64Type>::from(vec![
- Some(101),
- Some(102),
- Some(103),
- None,
- ]));
- let input_struct_2 = StructArray::from(vec![(field.clone(), input_primitive_2)]);
-
- let input_primitive_3: ArrayRef =
- Arc::new(PrimitiveArray::<Int64Type>::from(vec![
- Some(256),
- Some(512),
- Some(1024),
- ]));
- let input_struct_3 = StructArray::from(vec![(field, input_primitive_3)]);
-
- let arr = concat(&[&input_struct_1, &input_struct_2, &input_struct_3])?;
-
- let expected_primitive_output = Arc::new(PrimitiveArray::<Int64Type>::from(vec![
- Some(-1),
- Some(-1),
- Some(2),
- None,
- None,
- Some(101),
- Some(102),
- Some(103),
- None,
- Some(256),
- Some(512),
- Some(1024),
- ])) as ArrayRef;
-
- let actual_primitive = arr
- .as_any()
- .downcast_ref::<StructArray>()
- .unwrap()
- .column(0);
- assert_eq!(actual_primitive, &expected_primitive_output);
-
- Ok(())
- }
-
- #[test]
- fn test_concat_struct_array_slices() -> Result<()> {
- let field = Field::new("field", DataType::Int64, true);
- let input_primitive_1: ArrayRef =
- Arc::new(PrimitiveArray::<Int64Type>::from(vec![
- Some(-1),
- Some(-1),
- Some(2),
- None,
- None,
- ]));
- let input_struct_1 = StructArray::from(vec![(field.clone(), input_primitive_1)]);
-
- let input_primitive_2: ArrayRef =
- Arc::new(PrimitiveArray::<Int64Type>::from(vec![
- Some(101),
- Some(102),
- Some(103),
- None,
- ]));
- let input_struct_2 = StructArray::from(vec![(field, input_primitive_2)]);
-
- let arr = concat(&[
- input_struct_1.slice(1, 3).as_ref(),
- input_struct_2.slice(1, 2).as_ref(),
- ])?;
-
- let expected_primitive_output = Arc::new(PrimitiveArray::<Int64Type>::from(vec![
- Some(-1),
- Some(2),
- None,
- Some(102),
- Some(103),
- ])) as ArrayRef;
-
- let actual_primitive = arr
- .as_any()
- .downcast_ref::<StructArray>()
- .unwrap()
- .column(0);
- assert_eq!(actual_primitive, &expected_primitive_output);
-
- Ok(())
- }
-
- #[test]
- fn test_string_array_slices() -> Result<()> {
- let input_1 = StringArray::from(vec!["hello", "A", "B", "C"]);
- let input_2 = StringArray::from(vec!["world", "D", "E", "Z"]);
-
- let arr = concat(&[input_1.slice(1, 3).as_ref(), input_2.slice(1, 2).as_ref()])?;
-
- let expected_output = StringArray::from(vec!["A", "B", "C", "D", "E"]);
-
- let actual_output = arr.as_any().downcast_ref::<StringArray>().unwrap();
- assert_eq!(actual_output, &expected_output);
-
- Ok(())
- }
-
- #[test]
- fn test_string_array_with_null_slices() -> Result<()> {
- let input_1 = StringArray::from(vec![Some("hello"), None, Some("A"), Some("C")]);
- let input_2 = StringArray::from(vec![None, Some("world"), Some("D"), None]);
-
- let arr = concat(&[input_1.slice(1, 3).as_ref(), input_2.slice(1, 2).as_ref()])?;
-
- let expected_output =
- StringArray::from(vec![None, Some("A"), Some("C"), Some("world"), Some("D")]);
-
- let actual_output = arr.as_any().downcast_ref::<StringArray>().unwrap();
- assert_eq!(actual_output, &expected_output);
-
- Ok(())
- }
-
- fn collect_string_dictionary(
- dictionary: &DictionaryArray<Int32Type>,
- ) -> Vec<Option<String>> {
- let values = dictionary.values();
- let values = values.as_any().downcast_ref::<StringArray>().unwrap();
-
- dictionary
- .keys()
- .iter()
- .map(|key| key.map(|key| values.value(key as _).to_string()))
- .collect()
- }
-
- fn concat_dictionary(
- input_1: DictionaryArray<Int32Type>,
- input_2: DictionaryArray<Int32Type>,
- ) -> Vec<Option<String>> {
- let concat = concat(&[&input_1 as _, &input_2 as _]).unwrap();
- let concat = concat
- .as_any()
- .downcast_ref::<DictionaryArray<Int32Type>>()
- .unwrap();
-
- collect_string_dictionary(concat)
- }
-
- #[test]
- fn test_string_dictionary_array() {
- let input_1: DictionaryArray<Int32Type> =
- vec!["hello", "A", "B", "hello", "hello", "C"]
- .into_iter()
- .collect();
- let input_2: DictionaryArray<Int32Type> =
- vec!["hello", "E", "E", "hello", "F", "E"]
- .into_iter()
- .collect();
-
- let expected: Vec<_> = vec![
- "hello", "A", "B", "hello", "hello", "C", "hello", "E", "E", "hello", "F",
- "E",
- ]
- .into_iter()
- .map(|x| Some(x.to_string()))
- .collect();
-
- let concat = concat_dictionary(input_1, input_2);
- assert_eq!(concat, expected);
- }
-
- #[test]
- fn test_string_dictionary_array_nulls() {
- let input_1: DictionaryArray<Int32Type> =
- vec![Some("foo"), Some("bar"), None, Some("fiz")]
- .into_iter()
- .collect();
- let input_2: DictionaryArray<Int32Type> = vec![None].into_iter().collect();
- let expected = vec![
- Some("foo".to_string()),
- Some("bar".to_string()),
- None,
- Some("fiz".to_string()),
- None,
- ];
-
- let concat = concat_dictionary(input_1, input_2);
- assert_eq!(concat, expected);
- }
-
- #[test]
- fn test_concat_string_sizes() -> Result<()> {
- let a: LargeStringArray = ((0..150).map(|_| Some("foo"))).collect();
- let b: LargeStringArray = ((0..150).map(|_| Some("foo"))).collect();
- let c = LargeStringArray::from(vec![Some("foo"), Some("bar"), None, Some("baz")]);
- // 150 * 3 = 450
- // 150 * 3 = 450
- // 3 * 3 = 9
- // ------------+
- // 909
- // closest 64 byte aligned cap = 960
-
- let arr = concat(&[&a, &b, &c])?;
- // this would have been 1280 if we did not precompute the value lengths.
- assert_eq!(arr.data().buffers()[1].capacity(), 960);
-
- Ok(())
- }
-}
diff --git a/arrow/src/compute/kernels/filter.rs b/arrow/src/compute/kernels/filter.rs
deleted file mode 100644
index c5efadd..0000000
--- a/arrow/src/compute/kernels/filter.rs
+++ /dev/null
@@ -1,681 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-//! Defines miscellaneous array kernels.
-
-use crate::buffer::buffer_bin_and;
-use crate::datatypes::DataType;
-use crate::error::Result;
-use crate::record_batch::RecordBatch;
-use crate::{array::*, util::bit_chunk_iterator::BitChunkIterator};
-use std::iter::Enumerate;
-
-/// Function that can filter arbitrary arrays
-pub type Filter<'a> = Box<Fn(&ArrayData) -> ArrayData + 'a>;
-
-/// Internal state of [SlicesIterator]
-#[derive(Debug, PartialEq)]
-enum State {
- // it is iterating over bits of a mask (`u64`, steps of size of 1 slot)
- Bits(u64),
- // it is iterating over chunks (steps of size of 64 slots)
- Chunks,
- // it is iterating over the remainding bits (steps of size of 1 slot)
- Remainder,
- // nothing more to iterate.
- Finish,
-}
-
-/// An iterator of `(usize, usize)` each representing an interval `[start,end[` whose
-/// slots of a [BooleanArray] are true. Each interval corresponds to a contiguous region of memory to be
-/// "taken" from an array to be filtered.
-#[derive(Debug)]
-pub struct SlicesIterator<'a> {
- iter: Enumerate<BitChunkIterator<'a>>,
- state: State,
- filter: &'a BooleanArray,
- remainder_mask: u64,
- remainder_len: usize,
- chunk_len: usize,
- len: usize,
- start: usize,
- on_region: bool,
- current_chunk: usize,
- current_bit: usize,
-}
-
-impl<'a> SlicesIterator<'a> {
- pub fn new(filter: &'a BooleanArray) -> Self {
- let values = &filter.data_ref().buffers()[0];
- let chunks = values.bit_chunks(filter.offset(), filter.len());
-
- Self {
- iter: chunks.iter().enumerate(),
- state: State::Chunks,
- filter,
- remainder_len: chunks.remainder_len(),
- chunk_len: chunks.chunk_len(),
- remainder_mask: chunks.remainder_bits(),
- len: 0,
- start: 0,
- on_region: false,
- current_chunk: 0,
- current_bit: 0,
- }
- }
-
- /// Counts the number of set bits in the filter array.
- fn filter_count(&self) -> usize {
- let values = self.filter.values();
- values.count_set_bits_offset(self.filter.offset(), self.filter.len())
- }
-
- #[inline]
- fn current_start(&self) -> usize {
- self.current_chunk * 64 + self.current_bit
- }
-
- #[inline]
- fn iterate_bits(&mut self, mask: u64, max: usize) -> Option<(usize, usize)> {
- while self.current_bit < max {
- if (mask & (1 << self.current_bit)) != 0 {
- if !self.on_region {
- self.start = self.current_start();
- self.on_region = true;
- }
- self.len += 1;
- } else if self.on_region {
- let result = (self.start, self.start + self.len);
- self.len = 0;
- self.on_region = false;
- self.current_bit += 1;
- return Some(result);
- }
- self.current_bit += 1;
- }
- self.current_bit = 0;
- None
- }
-
- /// iterates over chunks.
- #[inline]
- fn iterate_chunks(&mut self) -> Option<(usize, usize)> {
- while let Some((i, mask)) = self.iter.next() {
- self.current_chunk = i;
- if mask == 0 {
- if self.on_region {
- let result = (self.start, self.start + self.len);
- self.len = 0;
- self.on_region = false;
- return Some(result);
- }
- } else if mask == 18446744073709551615u64 {
- // = !0u64
- if !self.on_region {
- self.start = self.current_start();
- self.on_region = true;
- }
- self.len += 64;
- } else {
- // there is a chunk that has a non-trivial mask => iterate over bits.
- self.state = State::Bits(mask);
- return None;
- }
- }
- // no more chunks => start iterating over the remainder
- self.current_chunk = self.chunk_len;
- self.state = State::Remainder;
- None
- }
-}
-
-impl<'a> Iterator for SlicesIterator<'a> {
- type Item = (usize, usize);
-
- fn next(&mut self) -> Option<Self::Item> {
- match self.state {
- State::Chunks => {
- match self.iterate_chunks() {
- None => {
- // iterating over chunks does not yield any new slice => continue to the next
- self.current_bit = 0;
- self.next()
- }
- other => other,
- }
- }
- State::Bits(mask) => {
- match self.iterate_bits(mask, 64) {
- None => {
- // iterating over bits does not yield any new slice => change back
- // to chunks and continue to the next
- self.state = State::Chunks;
- self.next()
- }
- other => other,
- }
- }
- State::Remainder => {
- match self.iterate_bits(self.remainder_mask, self.remainder_len) {
- None => {
- self.state = State::Finish;
- if self.on_region {
- Some((self.start, self.start + self.len))
- } else {
- None
- }
- }
- other => other,
- }
- }
- State::Finish => None,
- }
- }
-}
-
-/// Returns a prepared function optimized to filter multiple arrays.
-/// Creating this function requires time, but using it is faster than [filter] when the
-/// same filter needs to be applied to multiple arrays (e.g. a multi-column `RecordBatch`).
-/// WARNING: the nulls of `filter` are ignored and the value on its slot is considered.
-/// Therefore, it is considered undefined behavior to pass `filter` with null values.
-pub fn build_filter(filter: &BooleanArray) -> Result<Filter> {
- let iter = SlicesIterator::new(filter);
- let filter_count = iter.filter_count();
- let chunks = iter.collect::<Vec<_>>();
-
- Ok(Box::new(move |array: &ArrayData| {
- match filter_count {
- // return all
- len if len == array.len() => array.clone(),
- 0 => ArrayData::new_empty(array.data_type()),
- _ => {
- let mut mutable = MutableArrayData::new(vec![array], false, filter_count);
- chunks
- .iter()
- .for_each(|(start, end)| mutable.extend(0, *start, *end));
- mutable.freeze()
- }
- }
- }))
-}
-
-/// Remove null values by do a bitmask AND operation with null bits and the boolean bits.
-fn prep_null_mask_filter(filter: &BooleanArray) -> BooleanArray {
- let array_data = filter.data_ref();
- let null_bitmap = array_data.null_buffer().unwrap();
- let mask = filter.values();
- let offset = filter.offset();
-
- let new_mask = buffer_bin_and(mask, offset, null_bitmap, offset, filter.len());
-
- let array_data = ArrayData::builder(DataType::Boolean)
- .len(filter.len())
- .add_buffer(new_mask)
- .build();
- BooleanArray::from(array_data)
-}
-
-/// Filters an [Array], returning elements matching the filter (i.e. where the values are true).
-///
-/// # Example
-/// ```rust
-/// # use arrow::array::{Int32Array, BooleanArray};
-/// # use arrow::error::Result;
-/// # use arrow::compute::kernels::filter::filter;
-/// # fn main() -> Result<()> {
-/// let array = Int32Array::from(vec![5, 6, 7, 8, 9]);
-/// let filter_array = BooleanArray::from(vec![true, false, false, true, false]);
-/// let c = filter(&array, &filter_array)?;
-/// let c = c.as_any().downcast_ref::<Int32Array>().unwrap();
-/// assert_eq!(c, &Int32Array::from(vec![5, 8]));
-/// # Ok(())
-/// # }
-/// ```
-pub fn filter(array: &Array, predicate: &BooleanArray) -> Result<ArrayRef> {
- if predicate.null_count() > 0 {
- // this greatly simplifies subsequent filtering code
- // now we only have a boolean mask to deal with
- let predicate = prep_null_mask_filter(predicate);
- return filter(array, &predicate);
- }
-
- let iter = SlicesIterator::new(predicate);
- let filter_count = iter.filter_count();
- match filter_count {
- 0 => {
- // return empty
- Ok(new_empty_array(array.data_type()))
- }
- len if len == array.len() => {
- // return all
- let data = array.data().clone();
- Ok(make_array(data))
- }
- _ => {
- // actually filter
- let mut mutable =
- MutableArrayData::new(vec![array.data_ref()], false, filter_count);
- iter.for_each(|(start, end)| mutable.extend(0, start, end));
- let data = mutable.freeze();
- Ok(make_array(data))
- }
- }
-}
-
-/// Returns a new [RecordBatch] with arrays containing only values matching the filter.
-pub fn filter_record_batch(
- record_batch: &RecordBatch,
- predicate: &BooleanArray,
-) -> Result<RecordBatch> {
- if predicate.null_count() > 0 {
- // this greatly simplifies subsequent filtering code
- // now we only have a boolean mask to deal with
- let predicate = prep_null_mask_filter(predicate);
- return filter_record_batch(record_batch, &predicate);
- }
-
- let filter = build_filter(predicate)?;
- let filtered_arrays = record_batch
- .columns()
- .iter()
- .map(|a| make_array(filter(&a.data())))
- .collect();
- RecordBatch::try_new(record_batch.schema(), filtered_arrays)
-}
-
-#[cfg(test)]
-mod tests {
- use super::*;
- use crate::datatypes::Int64Type;
- use crate::{
- buffer::Buffer,
- datatypes::{DataType, Field},
- };
-
- macro_rules! def_temporal_test {
- ($test:ident, $array_type: ident, $data: expr) => {
- #[test]
- fn $test() {
- let a = $data;
- let b = BooleanArray::from(vec![true, false, true, false]);
- let c = filter(&a, &b).unwrap();
- let d = c.as_ref().as_any().downcast_ref::<$array_type>().unwrap();
- assert_eq!(2, d.len());
- assert_eq!(1, d.value(0));
- assert_eq!(3, d.value(1));
- }
- };
- }
-
- def_temporal_test!(
- test_filter_date32,
- Date32Array,
- Date32Array::from(vec![1, 2, 3, 4])
- );
- def_temporal_test!(
- test_filter_date64,
- Date64Array,
- Date64Array::from(vec![1, 2, 3, 4])
- );
- def_temporal_test!(
- test_filter_time32_second,
- Time32SecondArray,
- Time32SecondArray::from(vec![1, 2, 3, 4])
- );
- def_temporal_test!(
- test_filter_time32_millisecond,
- Time32MillisecondArray,
- Time32MillisecondArray::from(vec![1, 2, 3, 4])
- );
- def_temporal_test!(
- test_filter_time64_microsecond,
- Time64MicrosecondArray,
- Time64MicrosecondArray::from(vec![1, 2, 3, 4])
- );
- def_temporal_test!(
- test_filter_time64_nanosecond,
- Time64NanosecondArray,
- Time64NanosecondArray::from(vec![1, 2, 3, 4])
- );
- def_temporal_test!(
- test_filter_duration_second,
- DurationSecondArray,
- DurationSecondArray::from(vec![1, 2, 3, 4])
- );
- def_temporal_test!(
- test_filter_duration_millisecond,
- DurationMillisecondArray,
- DurationMillisecondArray::from(vec![1, 2, 3, 4])
- );
- def_temporal_test!(
- test_filter_duration_microsecond,
- DurationMicrosecondArray,
- DurationMicrosecondArray::from(vec![1, 2, 3, 4])
- );
- def_temporal_test!(
- test_filter_duration_nanosecond,
- DurationNanosecondArray,
- DurationNanosecondArray::from(vec![1, 2, 3, 4])
- );
- def_temporal_test!(
- test_filter_timestamp_second,
- TimestampSecondArray,
- TimestampSecondArray::from_vec(vec![1, 2, 3, 4], None)
- );
- def_temporal_test!(
- test_filter_timestamp_millisecond,
- TimestampMillisecondArray,
- TimestampMillisecondArray::from_vec(vec![1, 2, 3, 4], None)
- );
- def_temporal_test!(
- test_filter_timestamp_microsecond,
- TimestampMicrosecondArray,
- TimestampMicrosecondArray::from_vec(vec![1, 2, 3, 4], None)
- );
- def_temporal_test!(
- test_filter_timestamp_nanosecond,
- TimestampNanosecondArray,
- TimestampNanosecondArray::from_vec(vec![1, 2, 3, 4], None)
- );
-
- #[test]
- fn test_filter_array_slice() {
- let a_slice = Int32Array::from(vec![5, 6, 7, 8, 9]).slice(1, 4);
- let a = a_slice.as_ref();
- let b = BooleanArray::from(vec![true, false, false, true]);
- // filtering with sliced filter array is not currently supported
- // let b_slice = BooleanArray::from(vec![true, false, false, true, false]).slice(1, 4);
- // let b = b_slice.as_any().downcast_ref().unwrap();
- let c = filter(a, &b).unwrap();
- let d = c.as_ref().as_any().downcast_ref::<Int32Array>().unwrap();
- assert_eq!(2, d.len());
- assert_eq!(6, d.value(0));
- assert_eq!(9, d.value(1));
- }
-
- #[test]
- fn test_filter_array_low_density() {
- // this test exercises the all 0's branch of the filter algorithm
- let mut data_values = (1..=65).collect::<Vec<i32>>();
- let mut filter_values =
- (1..=65).map(|i| matches!(i % 65, 0)).collect::<Vec<bool>>();
- // set up two more values after the batch
- data_values.extend_from_slice(&[66, 67]);
- filter_values.extend_from_slice(&[false, true]);
- let a = Int32Array::from(data_values);
- let b = BooleanArray::from(filter_values);
- let c = filter(&a, &b).unwrap();
- let d = c.as_ref().as_any().downcast_ref::<Int32Array>().unwrap();
- assert_eq!(2, d.len());
- assert_eq!(65, d.value(0));
- assert_eq!(67, d.value(1));
- }
-
- #[test]
- fn test_filter_array_high_density() {
- // this test exercises the all 1's branch of the filter algorithm
- let mut data_values = (1..=65).map(Some).collect::<Vec<_>>();
- let mut filter_values = (1..=65)
- .map(|i| !matches!(i % 65, 0))
- .collect::<Vec<bool>>();
- // set second data value to null
- data_values[1] = None;
- // set up two more values after the batch
- data_values.extend_from_slice(&[Some(66), None, Some(67), None]);
- filter_values.extend_from_slice(&[false, true, true, true]);
- let a = Int32Array::from(data_values);
- let b = BooleanArray::from(filter_values);
- let c = filter(&a, &b).unwrap();
- let d = c.as_ref().as_any().downcast_ref::<Int32Array>().unwrap();
- assert_eq!(67, d.len());
- assert_eq!(3, d.null_count());
- assert_eq!(1, d.value(0));
- assert_eq!(true, d.is_null(1));
- assert_eq!(64, d.value(63));
- assert_eq!(true, d.is_null(64));
- assert_eq!(67, d.value(65));
- }
-
- #[test]
- fn test_filter_string_array_simple() {
- let a = StringArray::from(vec!["hello", " ", "world", "!"]);
- let b = BooleanArray::from(vec![true, false, true, false]);
- let c = filter(&a, &b).unwrap();
- let d = c.as_ref().as_any().downcast_ref::<StringArray>().unwrap();
- assert_eq!(2, d.len());
- assert_eq!("hello", d.value(0));
- assert_eq!("world", d.value(1));
- }
-
- #[test]
- fn test_filter_primative_array_with_null() {
- let a = Int32Array::from(vec![Some(5), None]);
- let b = BooleanArray::from(vec![false, true]);
- let c = filter(&a, &b).unwrap();
- let d = c.as_ref().as_any().downcast_ref::<Int32Array>().unwrap();
- assert_eq!(1, d.len());
- assert_eq!(true, d.is_null(0));
- }
-
- #[test]
- fn test_filter_string_array_with_null() {
- let a = StringArray::from(vec![Some("hello"), None, Some("world"), None]);
- let b = BooleanArray::from(vec![true, false, false, true]);
- let c = filter(&a, &b).unwrap();
- let d = c.as_ref().as_any().downcast_ref::<StringArray>().unwrap();
- assert_eq!(2, d.len());
- assert_eq!("hello", d.value(0));
- assert_eq!(false, d.is_null(0));
- assert_eq!(true, d.is_null(1));
- }
-
- #[test]
- fn test_filter_binary_array_with_null() {
- let data: Vec<Option<&[u8]>> = vec![Some(b"hello"), None, Some(b"world"), None];
- let a = BinaryArray::from(data);
- let b = BooleanArray::from(vec![true, false, false, true]);
- let c = filter(&a, &b).unwrap();
- let d = c.as_ref().as_any().downcast_ref::<BinaryArray>().unwrap();
- assert_eq!(2, d.len());
- assert_eq!(b"hello", d.value(0));
- assert_eq!(false, d.is_null(0));
- assert_eq!(true, d.is_null(1));
- }
-
- #[test]
- fn test_filter_array_slice_with_null() {
- let a_slice =
- Int32Array::from(vec![Some(5), None, Some(7), Some(8), Some(9)]).slice(1, 4);
- let a = a_slice.as_ref();
- let b = BooleanArray::from(vec![true, false, false, true]);
- // filtering with sliced filter array is not currently supported
- // let b_slice = BooleanArray::from(vec![true, false, false, true, false]).slice(1, 4);
- // let b = b_slice.as_any().downcast_ref().unwrap();
- let c = filter(a, &b).unwrap();
- let d = c.as_ref().as_any().downcast_ref::<Int32Array>().unwrap();
- assert_eq!(2, d.len());
- assert_eq!(true, d.is_null(0));
- assert_eq!(false, d.is_null(1));
- assert_eq!(9, d.value(1));
- }
-
- #[test]
- fn test_filter_dictionary_array() {
- let values = vec![Some("hello"), None, Some("world"), Some("!")];
- let a: Int8DictionaryArray = values.iter().copied().collect();
- let b = BooleanArray::from(vec![false, true, true, false]);
- let c = filter(&a, &b).unwrap();
- let d = c
- .as_ref()
- .as_any()
- .downcast_ref::<Int8DictionaryArray>()
- .unwrap();
- let value_array = d.values();
- let values = value_array.as_any().downcast_ref::<StringArray>().unwrap();
- // values are cloned in the filtered dictionary array
- assert_eq!(3, values.len());
- // but keys are filtered
- assert_eq!(2, d.len());
- assert_eq!(true, d.is_null(0));
- assert_eq!("world", values.value(d.keys().value(1) as usize));
- }
-
- #[test]
- fn test_filter_string_array_with_negated_boolean_array() {
- let a = StringArray::from(vec!["hello", " ", "world", "!"]);
- let mut bb = BooleanBuilder::new(2);
- bb.append_value(false).unwrap();
- bb.append_value(true).unwrap();
- bb.append_value(false).unwrap();
- bb.append_value(true).unwrap();
- let b = bb.finish();
- let b = crate::compute::not(&b).unwrap();
-
- let c = filter(&a, &b).unwrap();
- let d = c.as_ref().as_any().downcast_ref::<StringArray>().unwrap();
- assert_eq!(2, d.len());
- assert_eq!("hello", d.value(0));
- assert_eq!("world", d.value(1));
- }
-
- #[test]
- fn test_filter_list_array() {
- let value_data = ArrayData::builder(DataType::Int32)
- .len(8)
- .add_buffer(Buffer::from_slice_ref(&[0, 1, 2, 3, 4, 5, 6, 7]))
- .build();
-
- let value_offsets = Buffer::from_slice_ref(&[0i64, 3, 6, 8, 8]);
-
- let list_data_type =
- DataType::LargeList(Box::new(Field::new("item", DataType::Int32, false)));
- let list_data = ArrayData::builder(list_data_type)
- .len(4)
- .add_buffer(value_offsets)
- .add_child_data(value_data)
- .null_bit_buffer(Buffer::from([0b00000111]))
- .build();
-
- // a = [[0, 1, 2], [3, 4, 5], [6, 7], null]
- let a = LargeListArray::from(list_data);
- let b = BooleanArray::from(vec![false, true, false, true]);
- let result = filter(&a, &b).unwrap();
-
- // expected: [[3, 4, 5], null]
- let value_data = ArrayData::builder(DataType::Int32)
- .len(3)
- .add_buffer(Buffer::from_slice_ref(&[3, 4, 5]))
- .build();
-
- let value_offsets = Buffer::from_slice_ref(&[0i64, 3, 3]);
-
- let list_data_type =
- DataType::LargeList(Box::new(Field::new("item", DataType::Int32, false)));
- let expected = ArrayData::builder(list_data_type)
- .len(2)
- .add_buffer(value_offsets)
- .add_child_data(value_data)
- .null_bit_buffer(Buffer::from([0b00000001]))
- .build();
-
- assert_eq!(&make_array(expected), &result);
- }
-
- #[test]
- fn test_slice_iterator_bits() {
- let filter_values = (0..64).map(|i| i == 1).collect::<Vec<bool>>();
- let filter = BooleanArray::from(filter_values);
-
- let iter = SlicesIterator::new(&filter);
- let filter_count = iter.filter_count();
- let chunks = iter.collect::<Vec<_>>();
-
- assert_eq!(chunks, vec![(1, 2)]);
- assert_eq!(filter_count, 1);
- }
-
- #[test]
- fn test_slice_iterator_bits1() {
- let filter_values = (0..64).map(|i| i != 1).collect::<Vec<bool>>();
- let filter = BooleanArray::from(filter_values);
-
- let iter = SlicesIterator::new(&filter);
- let filter_count = iter.filter_count();
- let chunks = iter.collect::<Vec<_>>();
-
- assert_eq!(chunks, vec![(0, 1), (2, 64)]);
- assert_eq!(filter_count, 64 - 1);
- }
-
- #[test]
- fn test_slice_iterator_chunk_and_bits() {
- let filter_values = (0..130).map(|i| i % 62 != 0).collect::<Vec<bool>>();
- let filter = BooleanArray::from(filter_values);
-
- let iter = SlicesIterator::new(&filter);
- let filter_count = iter.filter_count();
- let chunks = iter.collect::<Vec<_>>();
-
- assert_eq!(chunks, vec![(1, 62), (63, 124), (125, 130)]);
- assert_eq!(filter_count, 61 + 61 + 5);
- }
-
- #[test]
- fn test_null_mask() -> Result<()> {
- use crate::compute::kernels::comparison;
- let a: PrimitiveArray<Int64Type> =
- PrimitiveArray::from(vec![Some(1), Some(2), None]);
- let mask0 = comparison::eq(&a, &a)?;
- let out0 = filter(&a, &mask0)?;
- let out_arr0 = out0
- .as_any()
- .downcast_ref::<PrimitiveArray<Int64Type>>()
- .unwrap();
-
- let mask1 = BooleanArray::from(vec![Some(true), Some(true), None]);
- let out1 = filter(&a, &mask1)?;
- let out_arr1 = out1
- .as_any()
- .downcast_ref::<PrimitiveArray<Int64Type>>()
- .unwrap();
- assert_eq!(mask0, mask1);
- assert_eq!(out_arr0, out_arr1);
- Ok(())
- }
-
- #[test]
- fn test_fast_path() -> Result<()> {
- let a: PrimitiveArray<Int64Type> =
- PrimitiveArray::from(vec![Some(1), Some(2), None]);
-
- // all true
- let mask = BooleanArray::from(vec![true, true, true]);
- let out = filter(&a, &mask)?;
- let b = out
- .as_any()
- .downcast_ref::<PrimitiveArray<Int64Type>>()
- .unwrap();
- assert_eq!(&a, b);
-
- // all false
- let mask = BooleanArray::from(vec![false, false, false]);
- let out = filter(&a, &mask)?;
- assert_eq!(out.len(), 0);
- assert_eq!(out.data_type(), &DataType::Int64);
- Ok(())
- }
-}
diff --git a/arrow/src/compute/kernels/length.rs b/arrow/src/compute/kernels/length.rs
deleted file mode 100644
index 1f9b6e5..0000000
--- a/arrow/src/compute/kernels/length.rs
+++ /dev/null
@@ -1,389 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-//! Defines kernel for length of a string array
-
-use crate::{
- array::*,
- buffer::Buffer,
- datatypes::{ArrowNativeType, ArrowPrimitiveType},
-};
-use crate::{
- datatypes::{DataType, Int32Type, Int64Type},
- error::{ArrowError, Result},
-};
-
-fn unary_offsets_string<O, F>(
- array: &GenericStringArray<O>,
- data_type: DataType,
- op: F,
-) -> ArrayRef
-where
- O: StringOffsetSizeTrait + ArrowNativeType,
- F: Fn(O) -> O,
-{
- // note: offsets are stored as u8, but they can be interpreted as OffsetSize
- let offsets = &array.data_ref().buffers()[0];
- // this is a 30% improvement over iterating over u8s and building OffsetSize, which
- // justifies the usage of `unsafe`.
- let slice: &[O] = &unsafe { offsets.typed_data::<O>() }[array.offset()..];
-
- let lengths = slice.windows(2).map(|offset| op(offset[1] - offset[0]));
-
- // JUSTIFICATION
- // Benefit
- // ~60% speedup
- // Soundness
- // `values` is an iterator with a known size.
- let buffer = unsafe { Buffer::from_trusted_len_iter(lengths) };
-
- let null_bit_buffer = array
- .data_ref()
- .null_bitmap()
- .as_ref()
- .map(|b| b.bits.clone());
-
- let data = ArrayData::new(
- data_type,
- array.len(),
- None,
- null_bit_buffer,
- 0,
- vec![buffer],
- vec![],
- );
- make_array(data)
-}
-
-fn octet_length<O: StringOffsetSizeTrait, T: ArrowPrimitiveType>(
- array: &dyn Array,
-) -> ArrayRef
-where
- T::Native: StringOffsetSizeTrait,
-{
- let array = array
- .as_any()
- .downcast_ref::<GenericStringArray<O>>()
- .unwrap();
- unary_offsets_string::<O, _>(array, T::DATA_TYPE, |x| x)
-}
-
-fn bit_length_impl<O: StringOffsetSizeTrait, T: ArrowPrimitiveType>(
- array: &dyn Array,
-) -> ArrayRef
-where
- T::Native: StringOffsetSizeTrait,
-{
- let array = array
- .as_any()
- .downcast_ref::<GenericStringArray<O>>()
- .unwrap();
- let bits_in_bytes = O::from_usize(8).unwrap();
- unary_offsets_string::<O, _>(array, T::DATA_TYPE, |x| x * bits_in_bytes)
-}
-
-/// Returns an array of Int32/Int64 denoting the number of bytes in each string in the array.
-///
-/// * this only accepts StringArray/Utf8 and LargeString/LargeUtf8
-/// * length of null is null.
-/// * length is in number of bytes
-pub fn length(array: &Array) -> Result<ArrayRef> {
- match array.data_type() {
- DataType::Utf8 => Ok(octet_length::<i32, Int32Type>(array)),
- DataType::LargeUtf8 => Ok(octet_length::<i64, Int64Type>(array)),
- _ => Err(ArrowError::ComputeError(format!(
- "length not supported for {:?}",
- array.data_type()
- ))),
- }
-}
-
-/// Returns an array of Int32/Int64 denoting the number of bits in each string in the array.
-///
-/// * this only accepts StringArray/Utf8 and LargeString/LargeUtf8
-/// * bit_length of null is null.
-/// * bit_length is in number of bits
-pub fn bit_length(array: &Array) -> Result<ArrayRef> {
- match array.data_type() {
- DataType::Utf8 => Ok(bit_length_impl::<i32, Int32Type>(array)),
- DataType::LargeUtf8 => Ok(bit_length_impl::<i64, Int64Type>(array)),
- _ => Err(ArrowError::ComputeError(format!(
- "bit_length not supported for {:?}",
- array.data_type()
- ))),
- }
-}
-
-#[cfg(test)]
-mod tests {
- use super::*;
-
- fn length_cases() -> Vec<(Vec<&'static str>, usize, Vec<i32>)> {
- fn double_vec<T: Clone>(v: Vec<T>) -> Vec<T> {
- [&v[..], &v[..]].concat()
- }
-
- // a large array
- let mut values = vec!["one", "on", "o", ""];
- let mut expected = vec![3, 2, 1, 0];
- for _ in 0..10 {
- values = double_vec(values);
- expected = double_vec(expected);
- }
-
- vec![
- (vec!["hello", " ", "world"], 3, vec![5, 1, 5]),
- (vec!["hello", " ", "world", "!"], 4, vec![5, 1, 5, 1]),
- (vec!["💖"], 1, vec![4]),
- (values, 4096, expected),
- ]
- }
-
- #[test]
- #[cfg_attr(miri, ignore)] // running forever
- fn length_test_string() -> Result<()> {
- length_cases()
- .into_iter()
- .try_for_each(|(input, len, expected)| {
- let array = StringArray::from(input);
- let result = length(&array)?;
- assert_eq!(len, result.len());
- let result = result.as_any().downcast_ref::<Int32Array>().unwrap();
- expected.iter().enumerate().for_each(|(i, value)| {
- assert_eq!(*value, result.value(i));
- });
- Ok(())
- })
- }
-
- #[test]
- #[cfg_attr(miri, ignore)] // running forever
- fn length_test_large_string() -> Result<()> {
- length_cases()
- .into_iter()
- .try_for_each(|(input, len, expected)| {
- let array = LargeStringArray::from(input);
- let result = length(&array)?;
- assert_eq!(len, result.len());
- let result = result.as_any().downcast_ref::<Int64Array>().unwrap();
- expected.iter().enumerate().for_each(|(i, value)| {
- assert_eq!(*value as i64, result.value(i));
- });
- Ok(())
- })
- }
-
- fn length_null_cases() -> Vec<(Vec<Option<&'static str>>, usize, Vec<Option<i32>>)> {
- vec![(
- vec![Some("one"), None, Some("three"), Some("four")],
- 4,
- vec![Some(3), None, Some(5), Some(4)],
- )]
- }
-
- #[test]
- fn length_null_string() -> Result<()> {
- length_null_cases()
- .into_iter()
- .try_for_each(|(input, len, expected)| {
- let array = StringArray::from(input);
- let result = length(&array)?;
- assert_eq!(len, result.len());
- let result = result.as_any().downcast_ref::<Int32Array>().unwrap();
-
- let expected: Int32Array = expected.into();
- assert_eq!(expected.data(), result.data());
- Ok(())
- })
- }
-
- #[test]
- fn length_null_large_string() -> Result<()> {
- length_null_cases()
- .into_iter()
- .try_for_each(|(input, len, expected)| {
- let array = LargeStringArray::from(input);
- let result = length(&array)?;
- assert_eq!(len, result.len());
- let result = result.as_any().downcast_ref::<Int64Array>().unwrap();
-
- // convert to i64
- let expected: Int64Array = expected
- .iter()
- .map(|e| e.map(|e| e as i64))
- .collect::<Vec<_>>()
- .into();
- assert_eq!(expected.data(), result.data());
- Ok(())
- })
- }
-
- /// Tests that length is not valid for u64.
- #[test]
- fn length_wrong_type() {
- let array: UInt64Array = vec![1u64].into();
-
- assert!(length(&array).is_err());
- }
-
- /// Tests with an offset
- #[test]
- fn length_offsets() -> Result<()> {
- let a = StringArray::from(vec!["hello", " ", "world"]);
- let b = make_array(
- ArrayData::builder(DataType::Utf8)
- .len(2)
- .offset(1)
- .buffers(a.data_ref().buffers().to_vec())
- .build(),
- );
- let result = length(b.as_ref())?;
-
- let expected = Int32Array::from(vec![1, 5]);
- assert_eq!(expected.data(), result.data());
-
- Ok(())
- }
-
- fn bit_length_cases() -> Vec<(Vec<&'static str>, usize, Vec<i32>)> {
- fn double_vec<T: Clone>(v: Vec<T>) -> Vec<T> {
- [&v[..], &v[..]].concat()
- }
-
- // a large array
- let mut values = vec!["one", "on", "o", ""];
- let mut expected = vec![24, 16, 8, 0];
- for _ in 0..10 {
- values = double_vec(values);
- expected = double_vec(expected);
- }
-
- vec![
- (vec!["hello", " ", "world", "!"], 4, vec![40, 8, 40, 8]),
- (vec!["💖"], 1, vec![32]),
- (vec!["josé"], 1, vec![40]),
- (values, 4096, expected),
- ]
- }
-
- #[test]
- #[cfg_attr(miri, ignore)] // error: this test uses too much memory to run on CI
- fn bit_length_test_string() -> Result<()> {
- bit_length_cases()
- .into_iter()
- .try_for_each(|(input, len, expected)| {
- let array = StringArray::from(input);
- let result = bit_length(&array)?;
- assert_eq!(len, result.len());
- let result = result.as_any().downcast_ref::<Int32Array>().unwrap();
- expected.iter().enumerate().for_each(|(i, value)| {
- assert_eq!(*value, result.value(i));
- });
- Ok(())
- })
- }
-
- #[test]
- #[cfg_attr(miri, ignore)] // error: this test uses too much memory to run on CI
- fn bit_length_test_large_string() -> Result<()> {
- bit_length_cases()
- .into_iter()
- .try_for_each(|(input, len, expected)| {
- let array = LargeStringArray::from(input);
- let result = bit_length(&array)?;
- assert_eq!(len, result.len());
- let result = result.as_any().downcast_ref::<Int64Array>().unwrap();
- expected.iter().enumerate().for_each(|(i, value)| {
- assert_eq!(*value as i64, result.value(i));
- });
- Ok(())
- })
- }
-
- fn bit_length_null_cases() -> Vec<(Vec<Option<&'static str>>, usize, Vec<Option<i32>>)>
- {
- vec![(
- vec![Some("one"), None, Some("three"), Some("four")],
- 4,
- vec![Some(24), None, Some(40), Some(32)],
- )]
- }
-
- #[test]
- fn bit_length_null_string() -> Result<()> {
- bit_length_null_cases()
- .into_iter()
- .try_for_each(|(input, len, expected)| {
- let array = StringArray::from(input);
- let result = bit_length(&array)?;
- assert_eq!(len, result.len());
- let result = result.as_any().downcast_ref::<Int32Array>().unwrap();
-
- let expected: Int32Array = expected.into();
- assert_eq!(expected.data(), result.data());
- Ok(())
- })
- }
-
- #[test]
- fn bit_length_null_large_string() -> Result<()> {
- bit_length_null_cases()
- .into_iter()
- .try_for_each(|(input, len, expected)| {
- let array = LargeStringArray::from(input);
- let result = bit_length(&array)?;
- assert_eq!(len, result.len());
- let result = result.as_any().downcast_ref::<Int64Array>().unwrap();
-
- // convert to i64
- let expected: Int64Array = expected
- .iter()
- .map(|e| e.map(|e| e as i64))
- .collect::<Vec<_>>()
- .into();
- assert_eq!(expected.data(), result.data());
- Ok(())
- })
- }
-
- /// Tests that bit_length is not valid for u64.
- #[test]
- fn bit_length_wrong_type() {
- let array: UInt64Array = vec![1u64].into();
-
- assert!(bit_length(&array).is_err());
- }
-
- /// Tests with an offset
- #[test]
- fn bit_length_offsets() -> Result<()> {
- let a = StringArray::from(vec!["hello", " ", "world"]);
- let b = make_array(
- ArrayData::builder(DataType::Utf8)
- .len(2)
- .offset(1)
- .buffers(a.data_ref().buffers().to_vec())
- .build(),
- );
- let result = bit_length(b.as_ref())?;
-
- let expected = Int32Array::from(vec![8, 40]);
- assert_eq!(expected.data(), result.data());
-
- Ok(())
- }
-}
diff --git a/arrow/src/compute/kernels/limit.rs b/arrow/src/compute/kernels/limit.rs
deleted file mode 100644
index 4b4b085..0000000
--- a/arrow/src/compute/kernels/limit.rs
+++ /dev/null
@@ -1,200 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-//! Defines miscellaneous array kernels.
-
-use crate::array::ArrayRef;
-
-/// Returns the array, taking only the number of elements specified
-///
-/// Limit performs a zero-copy slice of the array, and is a convenience method on slice
-/// where:
-/// * it performs a bounds-check on the array
-/// * it slices from offset 0
-pub fn limit(array: &ArrayRef, num_elements: usize) -> ArrayRef {
- let lim = num_elements.min(array.len());
- array.slice(0, lim)
-}
-
-#[cfg(test)]
-mod tests {
- use super::*;
- use crate::array::*;
- use crate::buffer::Buffer;
- use crate::datatypes::{DataType, Field};
- use crate::util::bit_util;
-
- use std::sync::Arc;
-
- #[test]
- fn test_limit_array() {
- let a: ArrayRef = Arc::new(Int32Array::from(vec![5, 6, 7, 8, 9]));
- let b = limit(&a, 3);
- let c = b.as_ref().as_any().downcast_ref::<Int32Array>().unwrap();
- assert_eq!(3, c.len());
- assert_eq!(5, c.value(0));
- assert_eq!(6, c.value(1));
- assert_eq!(7, c.value(2));
- }
-
- #[test]
- fn test_limit_string_array() {
- let a: ArrayRef = Arc::new(StringArray::from(vec!["hello", " ", "world", "!"]));
- let b = limit(&a, 2);
- let c = b.as_ref().as_any().downcast_ref::<StringArray>().unwrap();
- assert_eq!(2, c.len());
- assert_eq!("hello", c.value(0));
- assert_eq!(" ", c.value(1));
- }
-
- #[test]
- fn test_limit_array_with_null() {
- let a: ArrayRef = Arc::new(Int32Array::from(vec![None, Some(5)]));
- let b = limit(&a, 1);
- let c = b.as_ref().as_any().downcast_ref::<Int32Array>().unwrap();
- assert_eq!(1, c.len());
- assert_eq!(true, c.is_null(0));
- }
-
- #[test]
- fn test_limit_array_with_limit_too_large() {
- let a = Int32Array::from(vec![5, 6, 7, 8, 9]);
- let a_ref: ArrayRef = Arc::new(a);
- let b = limit(&a_ref, 6);
- let c = b.as_ref().as_any().downcast_ref::<Int32Array>().unwrap();
-
- assert_eq!(5, c.len());
- assert_eq!(5, c.value(0));
- assert_eq!(6, c.value(1));
- assert_eq!(7, c.value(2));
- assert_eq!(8, c.value(3));
- assert_eq!(9, c.value(4));
- }
-
- #[test]
- fn test_list_array_limit() {
- // adapted from crate::array::test::test_list_array_slice
- // Construct a value array
- let value_data = ArrayData::builder(DataType::Int32)
- .len(10)
- .add_buffer(Buffer::from_slice_ref(&[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]))
- .build();
-
- // Construct a buffer for value offsets, for the nested array:
- // [[0, 1], null, [2, 3], null, [4, 5], null, [6, 7, 8], null, [9]]
- let value_offsets = Buffer::from_slice_ref(&[0, 2, 2, 4, 4, 6, 6, 9, 9, 10]);
- // 01010101 00000001
- let mut null_bits: [u8; 2] = [0; 2];
- bit_util::set_bit(&mut null_bits, 0);
- bit_util::set_bit(&mut null_bits, 2);
- bit_util::set_bit(&mut null_bits, 4);
- bit_util::set_bit(&mut null_bits, 6);
- bit_util::set_bit(&mut null_bits, 8);
-
- // Construct a list array from the above two
- let list_data_type =
- DataType::List(Box::new(Field::new("item", DataType::Int32, false)));
- let list_data = ArrayData::builder(list_data_type)
- .len(9)
- .add_buffer(value_offsets)
- .add_child_data(value_data)
- .null_bit_buffer(Buffer::from(null_bits))
- .build();
- let list_array: ArrayRef = Arc::new(ListArray::from(list_data));
-
- let limit_array = limit(&list_array, 6);
- assert_eq!(6, limit_array.len());
- assert_eq!(0, limit_array.offset());
- assert_eq!(3, limit_array.null_count());
-
- // Check offset and length for each non-null value.
- let limit_array: &ListArray =
- limit_array.as_any().downcast_ref::<ListArray>().unwrap();
-
- for i in 0..limit_array.len() {
- let offset = limit_array.value_offsets()[i];
- let length = limit_array.value_length(i);
- if i % 2 == 0 {
- assert_eq!(2, length);
- assert_eq!(i as i32, offset);
- } else {
- assert_eq!(0, length);
- }
- }
- }
-
- #[test]
- fn test_struct_array_limit() {
- // adapted from crate::array::test::test_struct_array_slice
- let boolean_data = ArrayData::builder(DataType::Boolean)
- .len(5)
- .add_buffer(Buffer::from([0b00010000]))
- .null_bit_buffer(Buffer::from([0b00010001]))
- .build();
- let int_data = ArrayData::builder(DataType::Int32)
- .len(5)
- .add_buffer(Buffer::from_slice_ref(&[0, 28, 42, 0, 0]))
- .null_bit_buffer(Buffer::from([0b00000110]))
- .build();
-
- let mut field_types = vec![];
- field_types.push(Field::new("a", DataType::Boolean, false));
- field_types.push(Field::new("b", DataType::Int32, false));
- let struct_array_data = ArrayData::builder(DataType::Struct(field_types))
- .len(5)
- .add_child_data(boolean_data.clone())
- .add_child_data(int_data.clone())
- .null_bit_buffer(Buffer::from([0b00010111]))
- .build();
- let struct_array = StructArray::from(struct_array_data);
-
- assert_eq!(5, struct_array.len());
- assert_eq!(1, struct_array.null_count());
- assert_eq!(&boolean_data, struct_array.column(0).data());
- assert_eq!(&int_data, struct_array.column(1).data());
-
- let array: ArrayRef = Arc::new(struct_array);
-
- let sliced_array = limit(&array, 3);
- let sliced_array = sliced_array.as_any().downcast_ref::<StructArray>().unwrap();
- assert_eq!(3, sliced_array.len());
- assert_eq!(0, sliced_array.offset());
- assert_eq!(0, sliced_array.null_count());
- assert!(sliced_array.is_valid(0));
- assert!(sliced_array.is_valid(1));
- assert!(sliced_array.is_valid(2));
-
- let sliced_c0 = sliced_array.column(0);
- let sliced_c0 = sliced_c0.as_any().downcast_ref::<BooleanArray>().unwrap();
- assert_eq!(3, sliced_c0.len());
- assert_eq!(0, sliced_c0.offset());
- assert_eq!(2, sliced_c0.null_count());
- assert!(sliced_c0.is_valid(0));
- assert!(sliced_c0.is_null(1));
- assert!(sliced_c0.is_null(2));
- assert_eq!(false, sliced_c0.value(0));
-
- let sliced_c1 = sliced_array.column(1);
- let sliced_c1 = sliced_c1.as_any().downcast_ref::<Int32Array>().unwrap();
- assert_eq!(3, sliced_c1.len());
- assert_eq!(0, sliced_c1.offset());
- assert_eq!(1, sliced_c1.null_count());
- assert!(sliced_c1.is_null(0));
- assert_eq!(28, sliced_c1.value(1));
- assert_eq!(42, sliced_c1.value(2));
- }
-}
diff --git a/arrow/src/compute/kernels/mod.rs b/arrow/src/compute/kernels/mod.rs
deleted file mode 100644
index a0ef50a..0000000
--- a/arrow/src/compute/kernels/mod.rs
+++ /dev/null
@@ -1,38 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-//! Computation kernels on Arrow Arrays
-
-pub mod aggregate;
-pub mod arithmetic;
-pub mod arity;
-pub mod boolean;
-pub mod cast;
-pub mod cast_utils;
-pub mod comparison;
-pub mod concat;
-pub mod filter;
-pub mod length;
-pub mod limit;
-pub mod partition;
-pub mod regexp;
-pub mod sort;
-pub mod substring;
-pub mod take;
-pub mod temporal;
-pub mod window;
-pub mod zip;
diff --git a/arrow/src/compute/kernels/partition.rs b/arrow/src/compute/kernels/partition.rs
deleted file mode 100644
index ad35e92..0000000
--- a/arrow/src/compute/kernels/partition.rs
+++ /dev/null
@@ -1,301 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-//! Defines partition kernel for `ArrayRef`
-
-use crate::compute::kernels::sort::LexicographicalComparator;
-use crate::compute::SortColumn;
-use crate::error::{ArrowError, Result};
-use std::cmp::Ordering;
-use std::iter::Iterator;
-use std::ops::Range;
-
-/// Given a list of already sorted columns, find partition ranges that would partition
-/// lexicographically equal values across columns.
-///
-/// Here LexicographicalComparator is used in conjunction with binary
-/// search so the columns *MUST* be pre-sorted already.
-///
-/// The returned vec would be of size k where k is cardinality of the sorted values; Consecutive
-/// values will be connected: (a, b) and (b, c), where start = 0 and end = n for the first and last
-/// range.
-pub fn lexicographical_partition_ranges(
- columns: &[SortColumn],
-) -> Result<impl Iterator<Item = Range<usize>> + '_> {
- LexicographicalPartitionIterator::try_new(columns)
-}
-
-struct LexicographicalPartitionIterator<'a> {
- comparator: LexicographicalComparator<'a>,
- num_rows: usize,
- previous_partition_point: usize,
- partition_point: usize,
- value_indices: Vec<usize>,
-}
-
-impl<'a> LexicographicalPartitionIterator<'a> {
- fn try_new(columns: &'a [SortColumn]) -> Result<LexicographicalPartitionIterator> {
- if columns.is_empty() {
- return Err(ArrowError::InvalidArgumentError(
- "Sort requires at least one column".to_string(),
- ));
- }
- let num_rows = columns[0].values.len();
- if columns.iter().any(|item| item.values.len() != num_rows) {
- return Err(ArrowError::ComputeError(
- "Lexical sort columns have different row counts".to_string(),
- ));
- };
-
- let comparator = LexicographicalComparator::try_new(columns)?;
- let value_indices = (0..num_rows).collect::<Vec<usize>>();
- Ok(LexicographicalPartitionIterator {
- comparator,
- num_rows,
- previous_partition_point: 0,
- partition_point: 0,
- value_indices,
- })
- }
-}
-
-impl<'a> Iterator for LexicographicalPartitionIterator<'a> {
- type Item = Range<usize>;
-
- fn next(&mut self) -> Option<Self::Item> {
- if self.partition_point < self.num_rows {
- // invariant:
- // value_indices[0..previous_partition_point] all are values <= value_indices[previous_partition_point]
- // so in order to save time we can do binary search on the value_indices[previous_partition_point..]
- // and find when any value is greater than value_indices[previous_partition_point]; because we are using
- // new indices, the new offset is _added_ to the previous_partition_point.
- //
- // be careful that idx is of type &usize which points to the actual value within value_indices, which itself
- // contains usize (0..row_count), providing access to lexicographical_comparator as pointers into the
- // original columnar data.
- self.partition_point += self.value_indices[self.partition_point..]
- .partition_point(|idx| {
- self.comparator.compare(idx, &self.partition_point)
- != Ordering::Greater
- });
- let start = self.previous_partition_point;
- let end = self.partition_point;
- self.previous_partition_point = self.partition_point;
- Some(Range { start, end })
- } else {
- None
- }
- }
-}
-
-#[cfg(test)]
-mod tests {
- use super::*;
- use crate::array::*;
- use crate::compute::SortOptions;
- use crate::datatypes::DataType;
- use std::sync::Arc;
-
- #[test]
- fn test_lexicographical_partition_ranges_empty() {
- let input = vec![];
- assert!(
- lexicographical_partition_ranges(&input).is_err(),
- "lexicographical_partition_ranges should reject columns with empty rows"
- );
- }
-
- #[test]
- fn test_lexicographical_partition_ranges_unaligned_rows() {
- let input = vec![
- SortColumn {
- values: Arc::new(Int64Array::from(vec![None, Some(-1)])) as ArrayRef,
- options: None,
- },
- SortColumn {
- values: Arc::new(StringArray::from(vec![Some("foo")])) as ArrayRef,
- options: None,
- },
- ];
- assert!(
- lexicographical_partition_ranges(&input).is_err(),
- "lexicographical_partition_ranges should reject columns with different row counts"
- );
- }
-
- #[test]
- fn test_lexicographical_partition_single_column() -> Result<()> {
- let input = vec![SortColumn {
- values: Arc::new(Int64Array::from(vec![1, 2, 2, 2, 2, 2, 2, 2, 9]))
- as ArrayRef,
- options: Some(SortOptions {
- descending: false,
- nulls_first: true,
- }),
- }];
- {
- let results = lexicographical_partition_ranges(&input)?;
- assert_eq!(
- vec![(0_usize..1_usize), (1_usize..8_usize), (8_usize..9_usize)],
- results.collect::<Vec<_>>()
- );
- }
- Ok(())
- }
-
- #[test]
- fn test_lexicographical_partition_all_equal_values() -> Result<()> {
- let input = vec![SortColumn {
- values: Arc::new(Int64Array::from_value(1, 1000)) as ArrayRef,
- options: Some(SortOptions {
- descending: false,
- nulls_first: true,
- }),
- }];
-
- {
- let results = lexicographical_partition_ranges(&input)?;
- assert_eq!(vec![(0_usize..1000_usize)], results.collect::<Vec<_>>());
- }
- Ok(())
- }
-
- #[test]
- fn test_lexicographical_partition_all_null_values() -> Result<()> {
- let input = vec![
- SortColumn {
- values: new_null_array(&DataType::Int8, 1000),
- options: Some(SortOptions {
- descending: false,
- nulls_first: true,
- }),
- },
- SortColumn {
- values: new_null_array(&DataType::UInt16, 1000),
- options: Some(SortOptions {
- descending: false,
- nulls_first: false,
- }),
- },
- ];
- {
- let results = lexicographical_partition_ranges(&input)?;
- assert_eq!(vec![(0_usize..1000_usize)], results.collect::<Vec<_>>());
- }
- Ok(())
- }
-
- #[test]
- fn test_lexicographical_partition_unique_column_1() -> Result<()> {
- let input = vec![
- SortColumn {
- values: Arc::new(Int64Array::from(vec![None, Some(-1)])) as ArrayRef,
- options: Some(SortOptions {
- descending: false,
- nulls_first: true,
- }),
- },
- SortColumn {
- values: Arc::new(StringArray::from(vec![Some("foo"), Some("bar")]))
- as ArrayRef,
- options: Some(SortOptions {
- descending: true,
- nulls_first: true,
- }),
- },
- ];
- {
- let results = lexicographical_partition_ranges(&input)?;
- assert_eq!(
- vec![(0_usize..1_usize), (1_usize..2_usize)],
- results.collect::<Vec<_>>()
- );
- }
- Ok(())
- }
-
- #[test]
- fn test_lexicographical_partition_unique_column_2() -> Result<()> {
- let input = vec![
- SortColumn {
- values: Arc::new(Int64Array::from(vec![None, Some(-1), Some(-1)]))
- as ArrayRef,
- options: Some(SortOptions {
- descending: false,
- nulls_first: true,
- }),
- },
- SortColumn {
- values: Arc::new(StringArray::from(vec![
- Some("foo"),
- Some("bar"),
- Some("apple"),
- ])) as ArrayRef,
- options: Some(SortOptions {
- descending: true,
- nulls_first: true,
- }),
- },
- ];
- {
- let results = lexicographical_partition_ranges(&input)?;
- assert_eq!(
- vec![(0_usize..1_usize), (1_usize..2_usize), (2_usize..3_usize),],
- results.collect::<Vec<_>>()
- );
- }
- Ok(())
- }
-
- #[test]
- fn test_lexicographical_partition_non_unique_column_1() -> Result<()> {
- let input = vec![
- SortColumn {
- values: Arc::new(Int64Array::from(vec![
- None,
- Some(-1),
- Some(-1),
- Some(1),
- ])) as ArrayRef,
- options: Some(SortOptions {
- descending: false,
- nulls_first: true,
- }),
- },
- SortColumn {
- values: Arc::new(StringArray::from(vec![
- Some("foo"),
- Some("bar"),
- Some("bar"),
- Some("bar"),
- ])) as ArrayRef,
- options: Some(SortOptions {
- descending: true,
- nulls_first: true,
- }),
- },
- ];
- {
- let results = lexicographical_partition_ranges(&input)?;
- assert_eq!(
- vec![(0_usize..1_usize), (1_usize..3_usize), (3_usize..4_usize),],
- results.collect::<Vec<_>>()
- );
- }
- Ok(())
- }
-}
diff --git a/arrow/src/compute/kernels/regexp.rs b/arrow/src/compute/kernels/regexp.rs
deleted file mode 100644
index 5093dce..0000000
--- a/arrow/src/compute/kernels/regexp.rs
+++ /dev/null
@@ -1,158 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-//! Defines kernel to extract substrings based on a regular
-//! expression of a \[Large\]StringArray
-
-use crate::array::{
- ArrayRef, GenericStringArray, GenericStringBuilder, ListBuilder,
- StringOffsetSizeTrait,
-};
-use crate::error::{ArrowError, Result};
-use std::collections::HashMap;
-
-use std::sync::Arc;
-
-use regex::Regex;
-
-/// Extract all groups matched by a regular expression for a given String array.
-pub fn regexp_match<OffsetSize: StringOffsetSizeTrait>(
- array: &GenericStringArray<OffsetSize>,
- regex_array: &GenericStringArray<OffsetSize>,
- flags_array: Option<&GenericStringArray<OffsetSize>>,
-) -> Result<ArrayRef> {
- let mut patterns: HashMap<String, Regex> = HashMap::new();
- let builder: GenericStringBuilder<OffsetSize> = GenericStringBuilder::new(0);
- let mut list_builder = ListBuilder::new(builder);
-
- let complete_pattern = match flags_array {
- Some(flags) => Box::new(regex_array.iter().zip(flags.iter()).map(
- |(pattern, flags)| {
- pattern.map(|pattern| match flags {
- Some(value) => format!("(?{}){}", value, pattern),
- None => pattern.to_string(),
- })
- },
- )) as Box<dyn Iterator<Item = Option<String>>>,
- None => Box::new(
- regex_array
- .iter()
- .map(|pattern| pattern.map(|pattern| pattern.to_string())),
- ),
- };
- array
- .iter()
- .zip(complete_pattern)
- .map(|(value, pattern)| {
- match (value, pattern) {
- // Required for Postgres compatibility:
- // SELECT regexp_match('foobarbequebaz', ''); = {""}
- (Some(_), Some(pattern)) if pattern == *"" => {
- list_builder.values().append_value("")?;
- list_builder.append(true)?;
- }
- (Some(value), Some(pattern)) => {
- let existing_pattern = patterns.get(&pattern);
- let re = match existing_pattern {
- Some(re) => re.clone(),
- None => {
- let re = Regex::new(pattern.as_str()).map_err(|e| {
- ArrowError::ComputeError(format!(
- "Regular expression did not compile: {:?}",
- e
- ))
- })?;
- patterns.insert(pattern, re.clone());
- re
- }
- };
- match re.captures(value) {
- Some(caps) => {
- for m in caps.iter().skip(1).flatten() {
- list_builder.values().append_value(m.as_str())?;
- }
- list_builder.append(true)?
- }
- None => list_builder.append(false)?,
- }
- }
- _ => list_builder.append(false)?,
- }
- Ok(())
- })
- .collect::<Result<Vec<()>>>()?;
- Ok(Arc::new(list_builder.finish()))
-}
-
-#[cfg(test)]
-mod tests {
- use super::*;
- use crate::array::{ListArray, StringArray};
-
- #[test]
- fn match_single_group() -> Result<()> {
- let values = vec![
- Some("abc-005-def"),
- Some("X-7-5"),
- Some("X545"),
- None,
- Some("foobarbequebaz"),
- Some("foobarbequebaz"),
- ];
- let array = StringArray::from(values);
- let mut pattern_values = vec![r".*-(\d*)-.*"; 4];
- pattern_values.push(r"(bar)(bequ1e)");
- pattern_values.push("");
- let pattern = StringArray::from(pattern_values);
- let actual = regexp_match(&array, &pattern, None)?;
- let elem_builder: GenericStringBuilder<i32> = GenericStringBuilder::new(0);
- let mut expected_builder = ListBuilder::new(elem_builder);
- expected_builder.values().append_value("005")?;
- expected_builder.append(true)?;
- expected_builder.values().append_value("7")?;
- expected_builder.append(true)?;
- expected_builder.append(false)?;
- expected_builder.append(false)?;
- expected_builder.append(false)?;
- expected_builder.values().append_value("")?;
- expected_builder.append(true)?;
- let expected = expected_builder.finish();
- let result = actual.as_any().downcast_ref::<ListArray>().unwrap();
- assert_eq!(&expected, result);
- Ok(())
- }
-
- #[test]
- fn match_single_group_with_flags() -> Result<()> {
- let values = vec![Some("abc-005-def"), Some("X-7-5"), Some("X545"), None];
- let array = StringArray::from(values);
- let pattern = StringArray::from(vec![r"x.*-(\d*)-.*"; 4]);
- let flags = StringArray::from(vec!["i"; 4]);
- let actual = regexp_match(&array, &pattern, Some(&flags))?;
- let elem_builder: GenericStringBuilder<i32> = GenericStringBuilder::new(0);
- let mut expected_builder = ListBuilder::new(elem_builder);
- expected_builder.append(false)?;
- expected_builder.values().append_value("7")?;
- expected_builder.append(true)?;
- expected_builder.append(false)?;
- expected_builder.append(false)?;
- let expected = expected_builder.finish();
- let result = actual.as_any().downcast_ref::<ListArray>().unwrap();
- assert_eq!(&expected, result);
- Ok(())
- }
-}
diff --git a/arrow/src/compute/kernels/sort.rs b/arrow/src/compute/kernels/sort.rs
deleted file mode 100644
index b0eecb9..0000000
--- a/arrow/src/compute/kernels/sort.rs
+++ /dev/null
@@ -1,2664 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-//! Defines sort kernel for `ArrayRef`
-
-use crate::array::*;
-use crate::buffer::MutableBuffer;
-use crate::compute::take;
-use crate::datatypes::*;
-use crate::error::{ArrowError, Result};
-use std::cmp::Ordering;
-use TimeUnit::*;
-
-/// Sort the `ArrayRef` using `SortOptions`.
-///
-/// Performs a stable sort on values and indices. Nulls are ordered according to the `nulls_first` flag in `options`.
-/// Floats are sorted using IEEE 754 totalOrder
-///
-/// Returns an `ArrowError::ComputeError(String)` if the array type is either unsupported by `sort_to_indices` or `take`.
-///
-/// # Example
-/// ```rust
-/// # use std::sync::Arc;
-/// # use arrow::array::{Int32Array, ArrayRef};
-/// # use arrow::error::Result;
-/// # use arrow::compute::kernels::sort::sort;
-/// # fn main() -> Result<()> {
-/// let array: ArrayRef = Arc::new(Int32Array::from(vec![5, 4, 3, 2, 1]));
-/// let sorted_array = sort(&array, None).unwrap();
-/// let sorted_array = sorted_array.as_any().downcast_ref::<Int32Array>().unwrap();
-/// assert_eq!(sorted_array, &Int32Array::from(vec![1, 2, 3, 4, 5]));
-/// # Ok(())
-/// # }
-/// ```
-pub fn sort(values: &ArrayRef, options: Option<SortOptions>) -> Result<ArrayRef> {
- let indices = sort_to_indices(values, options, None)?;
- take(values.as_ref(), &indices, None)
-}
-
-/// Sort the `ArrayRef` partially.
-///
-/// If `limit` is specified, the resulting array will contain only
-/// first `limit` in the sort order. Any data data after the limit
-/// will be discarded.
-///
-/// Note: this is an unstable_sort, meaning it may not preserve the
-/// order of equal elements.
-///
-/// # Example
-/// ```rust
-/// # use std::sync::Arc;
-/// # use arrow::array::{Int32Array, ArrayRef};
-/// # use arrow::error::Result;
-/// # use arrow::compute::kernels::sort::{sort_limit, SortOptions};
-/// # fn main() -> Result<()> {
-/// let array: ArrayRef = Arc::new(Int32Array::from(vec![5, 4, 3, 2, 1]));
-///
-/// // Find the the top 2 items
-/// let sorted_array = sort_limit(&array, None, Some(2)).unwrap();
-/// let sorted_array = sorted_array.as_any().downcast_ref::<Int32Array>().unwrap();
-/// assert_eq!(sorted_array, &Int32Array::from(vec![1, 2]));
-///
-/// // Find the bottom top 2 items
-/// let options = Some(SortOptions {
-/// descending: true,
-/// ..Default::default()
-/// });
-/// let sorted_array = sort_limit(&array, options, Some(2)).unwrap();
-/// let sorted_array = sorted_array.as_any().downcast_ref::<Int32Array>().unwrap();
-/// assert_eq!(sorted_array, &Int32Array::from(vec![5, 4]));
-/// # Ok(())
-/// # }
-/// ```
-pub fn sort_limit(
- values: &ArrayRef,
- options: Option<SortOptions>,
- limit: Option<usize>,
-) -> Result<ArrayRef> {
- let indices = sort_to_indices(values, options, limit)?;
- take(values.as_ref(), &indices, None)
-}
-
-#[inline]
-fn sort_by<T, F>(array: &mut [T], limit: usize, cmp: F)
-where
- F: FnMut(&T, &T) -> Ordering,
-{
- if array.len() == limit {
- array.sort_by(cmp);
- } else {
- partial_sort(array, limit, cmp);
- }
-}
-
-// implements comparison using IEEE 754 total ordering for f32
-// Original implementation from https://doc.rust-lang.org/std/primitive.f64.html#method.total_cmp
-// TODO to change to use std when it becomes stable
-fn total_cmp_32(l: f32, r: f32) -> std::cmp::Ordering {
- let mut left = l.to_bits() as i32;
- let mut right = r.to_bits() as i32;
-
- left ^= (((left >> 31) as u32) >> 1) as i32;
- right ^= (((right >> 31) as u32) >> 1) as i32;
-
- left.cmp(&right)
-}
-
-// implements comparison using IEEE 754 total ordering for f64
-// Original implementation from https://doc.rust-lang.org/std/primitive.f64.html#method.total_cmp
-// TODO to change to use std when it becomes stable
-fn total_cmp_64(l: f64, r: f64) -> std::cmp::Ordering {
- let mut left = l.to_bits() as i64;
- let mut right = r.to_bits() as i64;
-
- left ^= (((left >> 63) as u64) >> 1) as i64;
- right ^= (((right >> 63) as u64) >> 1) as i64;
-
- left.cmp(&right)
-}
-
-fn cmp<T>(l: T, r: T) -> std::cmp::Ordering
-where
- T: Ord,
-{
- l.cmp(&r)
-}
-
-// partition indices into valid and null indices
-fn partition_validity(array: &ArrayRef) -> (Vec<u32>, Vec<u32>) {
- match array.null_count() {
- // faster path
- 0 => ((0..(array.len() as u32)).collect(), vec![]),
- _ => {
- let indices = 0..(array.len() as u32);
- indices.partition(|index| array.is_valid(*index as usize))
- }
- }
-}
-
-/// Sort elements from `ArrayRef` into an unsigned integer (`UInt32Array`) of indices.
-/// For floating point arrays any NaN values are considered to be greater than any other non-null value
-/// limit is an option for partial_sort
-pub fn sort_to_indices(
- values: &ArrayRef,
- options: Option<SortOptions>,
- limit: Option<usize>,
-) -> Result<UInt32Array> {
- let options = options.unwrap_or_default();
-
- let (v, n) = partition_validity(values);
-
- Ok(match values.data_type() {
- DataType::Boolean => sort_boolean(values, v, n, &options, limit),
- DataType::Int8 => {
- sort_primitive::<Int8Type, _>(values, v, n, cmp, &options, limit)
- }
- DataType::Int16 => {
- sort_primitive::<Int16Type, _>(values, v, n, cmp, &options, limit)
- }
- DataType::Int32 => {
- sort_primitive::<Int32Type, _>(values, v, n, cmp, &options, limit)
- }
- DataType::Int64 => {
- sort_primitive::<Int64Type, _>(values, v, n, cmp, &options, limit)
- }
- DataType::UInt8 => {
- sort_primitive::<UInt8Type, _>(values, v, n, cmp, &options, limit)
- }
- DataType::UInt16 => {
- sort_primitive::<UInt16Type, _>(values, v, n, cmp, &options, limit)
- }
- DataType::UInt32 => {
- sort_primitive::<UInt32Type, _>(values, v, n, cmp, &options, limit)
- }
- DataType::UInt64 => {
- sort_primitive::<UInt64Type, _>(values, v, n, cmp, &options, limit)
- }
- DataType::Float32 => {
- sort_primitive::<Float32Type, _>(values, v, n, total_cmp_32, &options, limit)
- }
- DataType::Float64 => {
- sort_primitive::<Float64Type, _>(values, v, n, total_cmp_64, &options, limit)
- }
- DataType::Date32 => {
- sort_primitive::<Date32Type, _>(values, v, n, cmp, &options, limit)
- }
- DataType::Date64 => {
- sort_primitive::<Date64Type, _>(values, v, n, cmp, &options, limit)
- }
- DataType::Time32(Second) => {
- sort_primitive::<Time32SecondType, _>(values, v, n, cmp, &options, limit)
- }
- DataType::Time32(Millisecond) => {
- sort_primitive::<Time32MillisecondType, _>(values, v, n, cmp, &options, limit)
- }
- DataType::Time64(Microsecond) => {
- sort_primitive::<Time64MicrosecondType, _>(values, v, n, cmp, &options, limit)
- }
- DataType::Time64(Nanosecond) => {
- sort_primitive::<Time64NanosecondType, _>(values, v, n, cmp, &options, limit)
- }
- DataType::Timestamp(Second, _) => {
- sort_primitive::<TimestampSecondType, _>(values, v, n, cmp, &options, limit)
- }
- DataType::Timestamp(Millisecond, _) => {
- sort_primitive::<TimestampMillisecondType, _>(
- values, v, n, cmp, &options, limit,
- )
- }
- DataType::Timestamp(Microsecond, _) => {
- sort_primitive::<TimestampMicrosecondType, _>(
- values, v, n, cmp, &options, limit,
- )
- }
- DataType::Timestamp(Nanosecond, _) => {
- sort_primitive::<TimestampNanosecondType, _>(
- values, v, n, cmp, &options, limit,
- )
- }
- DataType::Interval(IntervalUnit::YearMonth) => {
- sort_primitive::<IntervalYearMonthType, _>(values, v, n, cmp, &options, limit)
- }
- DataType::Interval(IntervalUnit::DayTime) => {
- sort_primitive::<IntervalDayTimeType, _>(values, v, n, cmp, &options, limit)
- }
- DataType::Duration(TimeUnit::Second) => {
- sort_primitive::<DurationSecondType, _>(values, v, n, cmp, &options, limit)
- }
- DataType::Duration(TimeUnit::Millisecond) => {
- sort_primitive::<DurationMillisecondType, _>(
- values, v, n, cmp, &options, limit,
- )
- }
- DataType::Duration(TimeUnit::Microsecond) => {
- sort_primitive::<DurationMicrosecondType, _>(
- values, v, n, cmp, &options, limit,
- )
- }
- DataType::Duration(TimeUnit::Nanosecond) => {
- sort_primitive::<DurationNanosecondType, _>(
- values, v, n, cmp, &options, limit,
- )
- }
- DataType::Utf8 => sort_string::<i32>(values, v, n, &options, limit),
- DataType::LargeUtf8 => sort_string::<i64>(values, v, n, &options, limit),
- DataType::List(field) => match field.data_type() {
- DataType::Int8 => sort_list::<i32, Int8Type>(values, v, n, &options, limit),
- DataType::Int16 => sort_list::<i32, Int16Type>(values, v, n, &options, limit),
- DataType::Int32 => sort_list::<i32, Int32Type>(values, v, n, &options, limit),
- DataType::Int64 => sort_list::<i32, Int64Type>(values, v, n, &options, limit),
- DataType::UInt8 => sort_list::<i32, UInt8Type>(values, v, n, &options, limit),
- DataType::UInt16 => {
- sort_list::<i32, UInt16Type>(values, v, n, &options, limit)
- }
- DataType::UInt32 => {
- sort_list::<i32, UInt32Type>(values, v, n, &options, limit)
- }
- DataType::UInt64 => {
- sort_list::<i32, UInt64Type>(values, v, n, &options, limit)
- }
- DataType::Float32 => {
- sort_list::<i32, Float32Type>(values, v, n, &options, limit)
- }
- DataType::Float64 => {
- sort_list::<i32, Float64Type>(values, v, n, &options, limit)
- }
- t => {
- return Err(ArrowError::ComputeError(format!(
- "Sort not supported for list type {:?}",
- t
- )))
- }
- },
- DataType::LargeList(field) => match field.data_type() {
- DataType::Int8 => sort_list::<i64, Int8Type>(values, v, n, &options, limit),
- DataType::Int16 => sort_list::<i64, Int16Type>(values, v, n, &options, limit),
- DataType::Int32 => sort_list::<i64, Int32Type>(values, v, n, &options, limit),
- DataType::Int64 => sort_list::<i64, Int64Type>(values, v, n, &options, limit),
- DataType::UInt8 => sort_list::<i64, UInt8Type>(values, v, n, &options, limit),
- DataType::UInt16 => {
- sort_list::<i64, UInt16Type>(values, v, n, &options, limit)
- }
- DataType::UInt32 => {
- sort_list::<i64, UInt32Type>(values, v, n, &options, limit)
- }
- DataType::UInt64 => {
- sort_list::<i64, UInt64Type>(values, v, n, &options, limit)
- }
- DataType::Float32 => {
- sort_list::<i64, Float32Type>(values, v, n, &options, limit)
- }
- DataType::Float64 => {
- sort_list::<i64, Float64Type>(values, v, n, &options, limit)
- }
- t => {
- return Err(ArrowError::ComputeError(format!(
- "Sort not supported for list type {:?}",
- t
- )))
- }
- },
- DataType::FixedSizeList(field, _) => match field.data_type() {
- DataType::Int8 => sort_list::<i32, Int8Type>(values, v, n, &options, limit),
- DataType::Int16 => sort_list::<i32, Int16Type>(values, v, n, &options, limit),
- DataType::Int32 => sort_list::<i32, Int32Type>(values, v, n, &options, limit),
- DataType::Int64 => sort_list::<i32, Int64Type>(values, v, n, &options, limit),
- DataType::UInt8 => sort_list::<i32, UInt8Type>(values, v, n, &options, limit),
- DataType::UInt16 => {
- sort_list::<i32, UInt16Type>(values, v, n, &options, limit)
- }
- DataType::UInt32 => {
- sort_list::<i32, UInt32Type>(values, v, n, &options, limit)
- }
- DataType::UInt64 => {
- sort_list::<i32, UInt64Type>(values, v, n, &options, limit)
- }
- DataType::Float32 => {
- sort_list::<i32, Float32Type>(values, v, n, &options, limit)
- }
- DataType::Float64 => {
- sort_list::<i32, Float64Type>(values, v, n, &options, limit)
- }
- t => {
- return Err(ArrowError::ComputeError(format!(
- "Sort not supported for list type {:?}",
- t
- )))
- }
- },
- DataType::Dictionary(key_type, value_type)
- if *value_type.as_ref() == DataType::Utf8 =>
- {
- match key_type.as_ref() {
- DataType::Int8 => {
- sort_string_dictionary::<Int8Type>(values, v, n, &options, limit)
- }
- DataType::Int16 => {
- sort_string_dictionary::<Int16Type>(values, v, n, &options, limit)
- }
- DataType::Int32 => {
- sort_string_dictionary::<Int32Type>(values, v, n, &options, limit)
- }
- DataType::Int64 => {
- sort_string_dictionary::<Int64Type>(values, v, n, &options, limit)
- }
- DataType::UInt8 => {
- sort_string_dictionary::<UInt8Type>(values, v, n, &options, limit)
- }
- DataType::UInt16 => {
- sort_string_dictionary::<UInt16Type>(values, v, n, &options, limit)
- }
- DataType::UInt32 => {
- sort_string_dictionary::<UInt32Type>(values, v, n, &options, limit)
- }
- DataType::UInt64 => {
- sort_string_dictionary::<UInt64Type>(values, v, n, &options, limit)
- }
- t => {
- return Err(ArrowError::ComputeError(format!(
- "Sort not supported for dictionary key type {:?}",
- t
- )))
- }
- }
- }
- t => {
- return Err(ArrowError::ComputeError(format!(
- "Sort not supported for data type {:?}",
- t
- )))
- }
- })
-}
-
-/// Options that define how sort kernels should behave
-#[derive(Clone, Copy, Debug, Eq, PartialEq)]
-pub struct SortOptions {
- /// Whether to sort in descending order
- pub descending: bool,
- /// Whether to sort nulls first
- pub nulls_first: bool,
-}
-
-impl Default for SortOptions {
- fn default() -> Self {
- Self {
- descending: false,
- // default to nulls first to match spark's behavior
- nulls_first: true,
- }
- }
-}
-
-/// Sort primitive values
-fn sort_boolean(
- values: &ArrayRef,
- value_indices: Vec<u32>,
- null_indices: Vec<u32>,
- options: &SortOptions,
- limit: Option<usize>,
-) -> UInt32Array {
- let values = values
- .as_any()
- .downcast_ref::<BooleanArray>()
- .expect("Unable to downcast to boolean array");
- let descending = options.descending;
-
- // create tuples that are used for sorting
- let mut valids = value_indices
- .into_iter()
- .map(|index| (index, values.value(index as usize)))
- .collect::<Vec<(u32, bool)>>();
-
- let mut nulls = null_indices;
-
- let valids_len = valids.len();
- let nulls_len = nulls.len();
-
- let mut len = values.len();
- if let Some(limit) = limit {
- len = limit.min(len);
- }
- if !descending {
- sort_by(&mut valids, len.saturating_sub(nulls_len), |a, b| {
- cmp(a.1, b.1)
- });
- } else {
- sort_by(&mut valids, len.saturating_sub(nulls_len), |a, b| {
- cmp(a.1, b.1).reverse()
- });
- // reverse to keep a stable ordering
- nulls.reverse();
- }
-
- // collect results directly into a buffer instead of a vec to avoid another aligned allocation
- let result_capacity = len * std::mem::size_of::<u32>();
- let mut result = MutableBuffer::new(result_capacity);
- // sets len to capacity so we can access the whole buffer as a typed slice
- result.resize(result_capacity, 0);
- let result_slice: &mut [u32] = result.typed_data_mut();
-
- if options.nulls_first {
- let size = nulls_len.min(len);
- result_slice[0..size].copy_from_slice(&nulls[0..size]);
- if nulls_len < len {
- insert_valid_values(result_slice, nulls_len, &valids[0..len - size]);
- }
- } else {
- // nulls last
- let size = valids.len().min(len);
- insert_valid_values(result_slice, 0, &valids[0..size]);
- if len > size {
- result_slice[valids_len..].copy_from_slice(&nulls[0..(len - valids_len)]);
- }
- }
-
- let result_data = ArrayData::new(
- DataType::UInt32,
- len,
- Some(0),
- None,
- 0,
- vec![result.into()],
- vec![],
- );
-
- UInt32Array::from(result_data)
-}
-
-/// Sort primitive values
-fn sort_primitive<T, F>(
- values: &ArrayRef,
- value_indices: Vec<u32>,
- null_indices: Vec<u32>,
- cmp: F,
- options: &SortOptions,
- limit: Option<usize>,
-) -> UInt32Array
-where
- T: ArrowPrimitiveType,
- T::Native: std::cmp::PartialOrd,
- F: Fn(T::Native, T::Native) -> std::cmp::Ordering,
-{
- let values = as_primitive_array::<T>(values);
- let descending = options.descending;
-
- // create tuples that are used for sorting
- let mut valids = value_indices
- .into_iter()
- .map(|index| (index, values.value(index as usize)))
- .collect::<Vec<(u32, T::Native)>>();
-
- let mut nulls = null_indices;
-
- let valids_len = valids.len();
- let nulls_len = nulls.len();
- let mut len = values.len();
-
- if let Some(limit) = limit {
- len = limit.min(len);
- }
- if !descending {
- sort_by(&mut valids, len.saturating_sub(nulls_len), |a, b| {
- cmp(a.1, b.1)
- });
- } else {
- sort_by(&mut valids, len.saturating_sub(nulls_len), |a, b| {
- cmp(a.1, b.1).reverse()
- });
- // reverse to keep a stable ordering
- nulls.reverse();
- }
-
- // collect results directly into a buffer instead of a vec to avoid another aligned allocation
- let result_capacity = len * std::mem::size_of::<u32>();
- let mut result = MutableBuffer::new(result_capacity);
- // sets len to capacity so we can access the whole buffer as a typed slice
- result.resize(result_capacity, 0);
- let result_slice: &mut [u32] = result.typed_data_mut();
-
- if options.nulls_first {
- let size = nulls_len.min(len);
- result_slice[0..size].copy_from_slice(&nulls[0..size]);
- if nulls_len < len {
- insert_valid_values(result_slice, nulls_len, &valids[0..len - size]);
- }
- } else {
- // nulls last
- let size = valids.len().min(len);
- insert_valid_values(result_slice, 0, &valids[0..size]);
- if len > size {
- result_slice[valids_len..].copy_from_slice(&nulls[0..(len - valids_len)]);
- }
- }
-
- let result_data = ArrayData::new(
- DataType::UInt32,
- len,
- Some(0),
- None,
- 0,
- vec![result.into()],
- vec![],
- );
-
- UInt32Array::from(result_data)
-}
-
-// insert valid and nan values in the correct order depending on the descending flag
-fn insert_valid_values<T>(result_slice: &mut [u32], offset: usize, valids: &[(u32, T)]) {
- let valids_len = valids.len();
- // helper to append the index part of the valid tuples
- let append_valids = move |dst_slice: &mut [u32]| {
- debug_assert_eq!(dst_slice.len(), valids_len);
- dst_slice
- .iter_mut()
- .zip(valids.iter())
- .for_each(|(dst, src)| *dst = src.0)
- };
-
- append_valids(&mut result_slice[offset..offset + valids.len()]);
-}
-
-/// Sort strings
-fn sort_string<Offset: StringOffsetSizeTrait>(
- values: &ArrayRef,
- value_indices: Vec<u32>,
- null_indices: Vec<u32>,
- options: &SortOptions,
- limit: Option<usize>,
-) -> UInt32Array {
- let values = values
- .as_any()
- .downcast_ref::<GenericStringArray<Offset>>()
- .unwrap();
-
- sort_string_helper(
- values,
- value_indices,
- null_indices,
- options,
- limit,
- |array, idx| array.value(idx as usize),
- )
-}
-
-/// Sort dictionary encoded strings
-fn sort_string_dictionary<T: ArrowDictionaryKeyType>(
- values: &ArrayRef,
- value_indices: Vec<u32>,
- null_indices: Vec<u32>,
- options: &SortOptions,
- limit: Option<usize>,
-) -> UInt32Array {
- let values: &DictionaryArray<T> = as_dictionary_array::<T>(values);
-
- let keys: &PrimitiveArray<T> = values.keys();
-
- let dict = values.values();
- let dict: &StringArray = as_string_array(&dict);
-
- sort_string_helper(
- keys,
- value_indices,
- null_indices,
- options,
- limit,
- |array: &PrimitiveArray<T>, idx| -> &str {
- let key: T::Native = array.value(idx as usize);
- dict.value(key.to_usize().unwrap())
- },
- )
-}
-
-/// shared implementation between dictionary encoded and plain string arrays
-#[inline]
-fn sort_string_helper<'a, A: Array, F>(
- values: &'a A,
- value_indices: Vec<u32>,
- null_indices: Vec<u32>,
- options: &SortOptions,
- limit: Option<usize>,
- value_fn: F,
-) -> UInt32Array
-where
- F: Fn(&'a A, u32) -> &str,
-{
- let mut valids = value_indices
- .into_iter()
- .map(|index| (index, value_fn(&values, index)))
- .collect::<Vec<(u32, &str)>>();
- let mut nulls = null_indices;
- let descending = options.descending;
- let mut len = values.len();
- let nulls_len = nulls.len();
-
- if let Some(limit) = limit {
- len = limit.min(len);
- }
- if !descending {
- sort_by(&mut valids, len.saturating_sub(nulls_len), |a, b| {
- cmp(a.1, b.1)
- });
- } else {
- sort_by(&mut valids, len.saturating_sub(nulls_len), |a, b| {
- cmp(a.1, b.1).reverse()
- });
- // reverse to keep a stable ordering
- nulls.reverse();
- }
- // collect the order of valid tuplies
- let mut valid_indices: Vec<u32> = valids.iter().map(|tuple| tuple.0).collect();
-
- if options.nulls_first {
- nulls.append(&mut valid_indices);
- nulls.truncate(len);
- UInt32Array::from(nulls)
- } else {
- // no need to sort nulls as they are in the correct order already
- valid_indices.append(&mut nulls);
- valid_indices.truncate(len);
- UInt32Array::from(valid_indices)
- }
-}
-
-fn sort_list<S, T>(
- values: &ArrayRef,
- value_indices: Vec<u32>,
- mut null_indices: Vec<u32>,
- options: &SortOptions,
- limit: Option<usize>,
-) -> UInt32Array
-where
- S: OffsetSizeTrait,
- T: ArrowPrimitiveType,
- T::Native: std::cmp::PartialOrd,
-{
- let mut valids: Vec<(u32, ArrayRef)> = values
- .as_any()
- .downcast_ref::<FixedSizeListArray>()
- .map_or_else(
- || {
- let values = as_generic_list_array::<S>(values);
- value_indices
- .iter()
- .copied()
- .map(|index| (index, values.value(index as usize)))
- .collect()
- },
- |values| {
- value_indices
- .iter()
- .copied()
- .map(|index| (index, values.value(index as usize)))
- .collect()
- },
- );
-
- let mut len = values.len();
- let nulls_len = null_indices.len();
- let descending = options.descending;
-
- if let Some(limit) = limit {
- len = limit.min(len);
- }
- if !descending {
- sort_by(&mut valids, len.saturating_sub(nulls_len), |a, b| {
- cmp_array(a.1.as_ref(), b.1.as_ref())
- });
- } else {
- sort_by(&mut valids, len.saturating_sub(nulls_len), |a, b| {
- cmp_array(a.1.as_ref(), b.1.as_ref()).reverse()
- });
- // reverse to keep a stable ordering
- null_indices.reverse();
- }
-
- let mut valid_indices: Vec<u32> = valids.iter().map(|tuple| tuple.0).collect();
- if options.nulls_first {
- null_indices.append(&mut valid_indices);
- null_indices.truncate(len);
- UInt32Array::from(null_indices)
- } else {
- valid_indices.append(&mut null_indices);
- valid_indices.truncate(len);
- UInt32Array::from(valid_indices)
- }
-}
-
-/// Compare two `Array`s based on the ordering defined in [ord](crate::array::ord).
-fn cmp_array(a: &Array, b: &Array) -> Ordering {
- let cmp_op = build_compare(a, b).unwrap();
- let length = a.len().max(b.len());
-
- for i in 0..length {
- let result = cmp_op(i, i);
- if result != Ordering::Equal {
- return result;
- }
- }
- Ordering::Equal
-}
-
-/// One column to be used in lexicographical sort
-#[derive(Clone, Debug)]
-pub struct SortColumn {
- pub values: ArrayRef,
- pub options: Option<SortOptions>,
-}
-
-/// Sort a list of `ArrayRef` using `SortOptions` provided for each array.
-///
-/// Performs a stable lexicographical sort on values and indices.
-///
-/// Returns an `ArrowError::ComputeError(String)` if any of the array type is either unsupported by
-/// `lexsort_to_indices` or `take`.
-///
-/// Example:
-///
-/// ```
-/// use std::convert::From;
-/// use std::sync::Arc;
-/// use arrow::array::{ArrayRef, StringArray, PrimitiveArray, as_primitive_array};
-/// use arrow::compute::kernels::sort::{SortColumn, SortOptions, lexsort};
-/// use arrow::datatypes::Int64Type;
-///
-/// let sorted_columns = lexsort(&vec![
-/// SortColumn {
-/// values: Arc::new(PrimitiveArray::<Int64Type>::from(vec![
-/// None,
-/// Some(-2),
-/// Some(89),
-/// Some(-64),
-/// Some(101),
-/// ])) as ArrayRef,
-/// options: None,
-/// },
-/// SortColumn {
-/// values: Arc::new(StringArray::from(vec![
-/// Some("hello"),
-/// Some("world"),
-/// Some(","),
-/// Some("foobar"),
-/// Some("!"),
-/// ])) as ArrayRef,
-/// options: Some(SortOptions {
-/// descending: true,
-/// nulls_first: false,
-/// }),
-/// },
-/// ], None).unwrap();
-///
-/// assert_eq!(as_primitive_array::<Int64Type>(&sorted_columns[0]).value(1), -64);
-/// assert!(sorted_columns[0].is_null(0));
-/// ```
-pub fn lexsort(columns: &[SortColumn], limit: Option<usize>) -> Result<Vec<ArrayRef>> {
- let indices = lexsort_to_indices(columns, limit)?;
- columns
- .iter()
- .map(|c| take(c.values.as_ref(), &indices, None))
- .collect()
-}
-
-/// Sort elements lexicographically from a list of `ArrayRef` into an unsigned integer
-/// (`UInt32Array`) of indices.
-pub fn lexsort_to_indices(
- columns: &[SortColumn],
- limit: Option<usize>,
-) -> Result<UInt32Array> {
- if columns.is_empty() {
- return Err(ArrowError::InvalidArgumentError(
- "Sort requires at least one column".to_string(),
- ));
- }
- if columns.len() == 1 {
- // fallback to non-lexical sort
- let column = &columns[0];
- return sort_to_indices(&column.values, column.options, limit);
- }
-
- let row_count = columns[0].values.len();
- if columns.iter().any(|item| item.values.len() != row_count) {
- return Err(ArrowError::ComputeError(
- "lexical sort columns have different row counts".to_string(),
- ));
- };
-
- let mut value_indices = (0..row_count).collect::<Vec<usize>>();
- let mut len = value_indices.len();
-
- if let Some(limit) = limit {
- len = limit.min(len);
- }
-
- let lexicographical_comparator = LexicographicalComparator::try_new(columns)?;
- sort_by(&mut value_indices, len, |a, b| {
- lexicographical_comparator.compare(a, b)
- });
-
- Ok(UInt32Array::from(
- (&value_indices)[0..len]
- .iter()
- .map(|i| *i as u32)
- .collect::<Vec<u32>>(),
- ))
-}
-
-/// It's unstable_sort, may not preserve the order of equal elements
-pub fn partial_sort<T, F>(v: &mut [T], limit: usize, mut is_less: F)
-where
- F: FnMut(&T, &T) -> Ordering,
-{
- let (before, _mid, _after) = v.select_nth_unstable_by(limit, &mut is_less);
- before.sort_unstable_by(is_less);
-}
-
-type LexicographicalCompareItem<'a> = (
- &'a ArrayData, // data
- Box<dyn Fn(usize, usize) -> Ordering + 'a>, // comparator
- SortOptions, // sort_option
-);
-
-/// A lexicographical comparator that wraps given array data (columns) and can lexicographically compare data
-/// at given two indices. The lifetime is the same at the data wrapped.
-pub(super) struct LexicographicalComparator<'a> {
- compare_items: Vec<LexicographicalCompareItem<'a>>,
-}
-
-impl LexicographicalComparator<'_> {
- /// lexicographically compare values at the wrapped columns with given indices.
- pub(super) fn compare<'a, 'b>(
- &'a self,
- a_idx: &'b usize,
- b_idx: &'b usize,
- ) -> Ordering {
- for (data, comparator, sort_option) in &self.compare_items {
- match (data.is_valid(*a_idx), data.is_valid(*b_idx)) {
- (true, true) => {
- match (comparator)(*a_idx, *b_idx) {
- // equal, move on to next column
- Ordering::Equal => continue,
- order => {
- if sort_option.descending {
- return order.reverse();
- } else {
- return order;
- }
- }
- }
- }
- (false, true) => {
- return if sort_option.nulls_first {
- Ordering::Less
- } else {
- Ordering::Greater
- };
- }
- (true, false) => {
- return if sort_option.nulls_first {
- Ordering::Greater
- } else {
- Ordering::Less
- };
- }
- // equal, move on to next column
- (false, false) => continue,
- }
- }
-
- Ordering::Equal
- }
-
- /// Create a new lex comparator that will wrap the given sort columns and give comparison
- /// results with two indices.
- pub(super) fn try_new(
- columns: &[SortColumn],
- ) -> Result<LexicographicalComparator<'_>> {
- let compare_items = columns
- .iter()
- .map(|column| {
- // flatten and convert build comparators
- // use ArrayData for is_valid checks later to avoid dynamic call
- let values = column.values.as_ref();
- let data = values.data_ref();
- Ok((
- data,
- build_compare(values, values)?,
- column.options.unwrap_or_default(),
- ))
- })
- .collect::<Result<Vec<_>>>()?;
- Ok(LexicographicalComparator { compare_items })
- }
-}
-
-#[cfg(test)]
-mod tests {
- use super::*;
- use crate::compute::util::tests::{
- build_fixed_size_list_nullable, build_generic_list_nullable,
- };
- use rand::rngs::StdRng;
- use rand::{Rng, RngCore, SeedableRng};
- use std::convert::TryFrom;
- use std::iter::FromIterator;
- use std::sync::Arc;
-
- fn test_sort_to_indices_boolean_arrays(
- data: Vec<Option<bool>>,
- options: Option<SortOptions>,
- limit: Option<usize>,
- expected_data: Vec<u32>,
- ) {
- let output = BooleanArray::from(data);
- let expected = UInt32Array::from(expected_data);
- let output =
- sort_to_indices(&(Arc::new(output) as ArrayRef), options, limit).unwrap();
- assert_eq!(output, expected)
- }
-
- fn test_sort_to_indices_primitive_arrays<T>(
- data: Vec<Option<T::Native>>,
- options: Option<SortOptions>,
- limit: Option<usize>,
- expected_data: Vec<u32>,
- ) where
- T: ArrowPrimitiveType,
- PrimitiveArray<T>: From<Vec<Option<T::Native>>>,
- {
- let output = PrimitiveArray::<T>::from(data);
- let expected = UInt32Array::from(expected_data);
- let output =
- sort_to_indices(&(Arc::new(output) as ArrayRef), options, limit).unwrap();
- assert_eq!(output, expected)
- }
-
- fn test_sort_primitive_arrays<T>(
- data: Vec<Option<T::Native>>,
- options: Option<SortOptions>,
- limit: Option<usize>,
- expected_data: Vec<Option<T::Native>>,
- ) where
- T: ArrowPrimitiveType,
- PrimitiveArray<T>: From<Vec<Option<T::Native>>>,
- {
- let output = PrimitiveArray::<T>::from(data);
- let expected = Arc::new(PrimitiveArray::<T>::from(expected_data)) as ArrayRef;
- let output = match limit {
- Some(_) => {
- sort_limit(&(Arc::new(output) as ArrayRef), options, limit).unwrap()
- }
- _ => sort(&(Arc::new(output) as ArrayRef), options).unwrap(),
- };
- assert_eq!(&output, &expected)
- }
-
- fn test_sort_to_indices_string_arrays(
- data: Vec<Option<&str>>,
- options: Option<SortOptions>,
- limit: Option<usize>,
- expected_data: Vec<u32>,
- ) {
- let output = StringArray::from(data);
- let expected = UInt32Array::from(expected_data);
- let output =
- sort_to_indices(&(Arc::new(output) as ArrayRef), options, limit).unwrap();
- assert_eq!(output, expected)
- }
-
- /// Tests both Utf8 and LargeUtf8
- fn test_sort_string_arrays(
- data: Vec<Option<&str>>,
- options: Option<SortOptions>,
- limit: Option<usize>,
- expected_data: Vec<Option<&str>>,
- ) {
- let output = StringArray::from(data.clone());
- let expected = Arc::new(StringArray::from(expected_data.clone())) as ArrayRef;
- let output = match limit {
- Some(_) => {
- sort_limit(&(Arc::new(output) as ArrayRef), options, limit).unwrap()
- }
- _ => sort(&(Arc::new(output) as ArrayRef), options).unwrap(),
- };
- assert_eq!(&output, &expected);
-
- let output = LargeStringArray::from(data);
- let expected = Arc::new(LargeStringArray::from(expected_data)) as ArrayRef;
- let output = match limit {
- Some(_) => {
- sort_limit(&(Arc::new(output) as ArrayRef), options, limit).unwrap()
- }
- _ => sort(&(Arc::new(output) as ArrayRef), options).unwrap(),
- };
- assert_eq!(&output, &expected)
- }
-
- fn test_sort_string_dict_arrays<T: ArrowDictionaryKeyType>(
- data: Vec<Option<&str>>,
- options: Option<SortOptions>,
- limit: Option<usize>,
- expected_data: Vec<Option<&str>>,
- ) {
- let array = DictionaryArray::<T>::from_iter(data.into_iter());
- let array_values = array.values().clone();
- let dict = array_values
- .as_any()
- .downcast_ref::<StringArray>()
- .expect("Unable to get dictionary values");
-
- let sorted = match limit {
- Some(_) => {
- sort_limit(&(Arc::new(array) as ArrayRef), options, limit).unwrap()
- }
- _ => sort(&(Arc::new(array) as ArrayRef), options).unwrap(),
- };
- let sorted = sorted
- .as_any()
- .downcast_ref::<DictionaryArray<T>>()
- .unwrap();
- let sorted_values = sorted.values();
- let sorted_dict = sorted_values
- .as_any()
- .downcast_ref::<StringArray>()
- .expect("Unable to get dictionary values");
- let sorted_keys = sorted.keys();
-
- assert_eq!(sorted_dict, dict);
-
- let sorted_strings = StringArray::try_from(
- (0..sorted.len())
- .map(|i| {
- if sorted.is_valid(i) {
- Some(sorted_dict.value(sorted_keys.value(i).to_usize().unwrap()))
- } else {
- None
- }
- })
- .collect::<Vec<Option<&str>>>(),
- )
- .expect("Unable to create string array from dictionary");
- let expected =
- StringArray::try_from(expected_data).expect("Unable to create string array");
-
- assert_eq!(sorted_strings, expected)
- }
-
- fn test_sort_list_arrays<T>(
- data: Vec<Option<Vec<Option<T::Native>>>>,
- options: Option<SortOptions>,
- limit: Option<usize>,
- expected_data: Vec<Option<Vec<Option<T::Native>>>>,
- fixed_length: Option<i32>,
- ) where
- T: ArrowPrimitiveType,
- PrimitiveArray<T>: From<Vec<Option<T::Native>>>,
- {
- // for FixedSizedList
- if let Some(length) = fixed_length {
- let input = Arc::new(build_fixed_size_list_nullable(data.clone(), length));
- let sorted = match limit {
- Some(_) => sort_limit(&(input as ArrayRef), options, limit).unwrap(),
- _ => sort(&(input as ArrayRef), options).unwrap(),
- };
- let expected = Arc::new(build_fixed_size_list_nullable(
- expected_data.clone(),
- length,
- )) as ArrayRef;
-
- assert_eq!(&sorted, &expected);
- }
-
- // for List
- let input = Arc::new(build_generic_list_nullable::<i32, T>(data.clone()));
- let sorted = match limit {
- Some(_) => sort_limit(&(input as ArrayRef), options, limit).unwrap(),
- _ => sort(&(input as ArrayRef), options).unwrap(),
- };
- let expected =
- Arc::new(build_generic_list_nullable::<i32, T>(expected_data.clone()))
- as ArrayRef;
-
- assert_eq!(&sorted, &expected);
-
- // for LargeList
- let input = Arc::new(build_generic_list_nullable::<i64, T>(data));
- let sorted = match limit {
- Some(_) => sort_limit(&(input as ArrayRef), options, limit).unwrap(),
- _ => sort(&(input as ArrayRef), options).unwrap(),
- };
- let expected =
- Arc::new(build_generic_list_nullable::<i64, T>(expected_data)) as ArrayRef;
-
- assert_eq!(&sorted, &expected);
- }
-
- fn test_lex_sort_arrays(
- input: Vec<SortColumn>,
- expected_output: Vec<ArrayRef>,
- limit: Option<usize>,
- ) {
- let sorted = lexsort(&input, limit).unwrap();
-
- for (result, expected) in sorted.iter().zip(expected_output.iter()) {
- assert_eq!(result, expected);
- }
- }
-
- #[test]
- fn test_sort_to_indices_primitives() {
- test_sort_to_indices_primitive_arrays::<Int8Type>(
- vec![None, Some(0), Some(2), Some(-1), Some(0), None],
- None,
- None,
- vec![0, 5, 3, 1, 4, 2],
- );
- test_sort_to_indices_primitive_arrays::<Int16Type>(
- vec![None, Some(0), Some(2), Some(-1), Some(0), None],
- None,
- None,
- vec![0, 5, 3, 1, 4, 2],
- );
- test_sort_to_indices_primitive_arrays::<Int32Type>(
- vec![None, Some(0), Some(2), Some(-1), Some(0), None],
- None,
- None,
- vec![0, 5, 3, 1, 4, 2],
- );
- test_sort_to_indices_primitive_arrays::<Int64Type>(
- vec![None, Some(0), Some(2), Some(-1), Some(0), None],
- None,
- None,
- vec![0, 5, 3, 1, 4, 2],
- );
- test_sort_to_indices_primitive_arrays::<Float32Type>(
- vec![
- None,
- Some(-0.05),
- Some(2.225),
- Some(-1.01),
- Some(-0.05),
- None,
- ],
- None,
- None,
- vec![0, 5, 3, 1, 4, 2],
- );
- test_sort_to_indices_primitive_arrays::<Float64Type>(
- vec![
- None,
- Some(-0.05),
- Some(2.225),
- Some(-1.01),
- Some(-0.05),
- None,
- ],
- None,
- None,
- vec![0, 5, 3, 1, 4, 2],
- );
-
- // descending
- test_sort_to_indices_primitive_arrays::<Int8Type>(
- vec![None, Some(0), Some(2), Some(-1), Some(0), None],
- Some(SortOptions {
- descending: true,
- nulls_first: false,
- }),
- None,
- vec![2, 1, 4, 3, 5, 0], // [2, 4, 1, 3, 5, 0]
- );
-
- test_sort_to_indices_primitive_arrays::<Int16Type>(
- vec![None, Some(0), Some(2), Some(-1), Some(0), None],
- Some(SortOptions {
- descending: true,
- nulls_first: false,
- }),
- None,
- vec![2, 1, 4, 3, 5, 0],
- );
-
- test_sort_to_indices_primitive_arrays::<Int32Type>(
- vec![None, Some(0), Some(2), Some(-1), Some(0), None],
- Some(SortOptions {
- descending: true,
- nulls_first: false,
- }),
- None,
- vec![2, 1, 4, 3, 5, 0],
- );
-
- test_sort_to_indices_primitive_arrays::<Int64Type>(
- vec![None, Some(0), Some(2), Some(-1), Some(0), None],
- Some(SortOptions {
- descending: true,
- nulls_first: false,
- }),
- None,
- vec![2, 1, 4, 3, 5, 0],
- );
-
- test_sort_to_indices_primitive_arrays::<Float32Type>(
- vec![
- None,
- Some(0.005),
- Some(20.22),
- Some(-10.3),
- Some(0.005),
- None,
- ],
- Some(SortOptions {
- descending: true,
- nulls_first: false,
- }),
- None,
- vec![2, 1, 4, 3, 5, 0],
- );
-
- test_sort_to_indices_primitive_arrays::<Float64Type>(
- vec![None, Some(0.0), Some(2.0), Some(-1.0), Some(0.0), None],
- Some(SortOptions {
- descending: true,
- nulls_first: false,
- }),
- None,
- vec![2, 1, 4, 3, 5, 0],
- );
-
- // descending, nulls first
- test_sort_to_indices_primitive_arrays::<Int8Type>(
- vec![None, Some(0), Some(2), Some(-1), Some(0), None],
- Some(SortOptions {
- descending: true,
- nulls_first: true,
- }),
- None,
- vec![5, 0, 2, 1, 4, 3], // [5, 0, 2, 4, 1, 3]
- );
-
- test_sort_to_indices_primitive_arrays::<Int16Type>(
- vec![None, Some(0), Some(2), Some(-1), Some(0), None],
- Some(SortOptions {
- descending: true,
- nulls_first: true,
- }),
- None,
- vec![5, 0, 2, 1, 4, 3], // [5, 0, 2, 4, 1, 3]
- );
-
- test_sort_to_indices_primitive_arrays::<Int32Type>(
- vec![None, Some(0), Some(2), Some(-1), Some(0), None],
- Some(SortOptions {
- descending: true,
- nulls_first: true,
- }),
- None,
- vec![5, 0, 2, 1, 4, 3],
- );
-
- test_sort_to_indices_primitive_arrays::<Int64Type>(
- vec![None, Some(0), Some(2), Some(-1), Some(0), None],
- Some(SortOptions {
- descending: true,
- nulls_first: true,
- }),
- None,
- vec![5, 0, 2, 1, 4, 3],
- );
-
- test_sort_to_indices_primitive_arrays::<Float32Type>(
- vec![None, Some(0.1), Some(0.2), Some(-1.3), Some(0.01), None],
- Some(SortOptions {
- descending: true,
- nulls_first: true,
- }),
- None,
- vec![5, 0, 2, 1, 4, 3],
- );
-
- test_sort_to_indices_primitive_arrays::<Float64Type>(
- vec![None, Some(10.1), Some(100.2), Some(-1.3), Some(10.01), None],
- Some(SortOptions {
- descending: true,
- nulls_first: true,
- }),
- None,
- vec![5, 0, 2, 1, 4, 3],
- );
-
- // valid values less than limit with extra nulls
- test_sort_to_indices_primitive_arrays::<Float64Type>(
- vec![Some(2.0), None, None, Some(1.0)],
- Some(SortOptions {
- descending: false,
- nulls_first: false,
- }),
- Some(3),
- vec![3, 0, 1],
- );
-
- test_sort_to_indices_primitive_arrays::<Float64Type>(
- vec![Some(2.0), None, None, Some(1.0)],
- Some(SortOptions {
- descending: false,
- nulls_first: true,
- }),
- Some(3),
- vec![1, 2, 3],
- );
-
- // more nulls than limit
- test_sort_to_indices_primitive_arrays::<Float64Type>(
- vec![Some(1.0), None, None, None],
- Some(SortOptions {
- descending: false,
- nulls_first: true,
- }),
- Some(2),
- vec![1, 2],
- );
-
- test_sort_to_indices_primitive_arrays::<Float64Type>(
- vec![Some(1.0), None, None, None],
- Some(SortOptions {
- descending: false,
- nulls_first: false,
- }),
- Some(2),
- vec![0, 1],
- );
- }
-
- #[test]
- fn test_sort_boolean() {
- // boolean
- test_sort_to_indices_boolean_arrays(
- vec![None, Some(false), Some(true), Some(true), Some(false), None],
- None,
- None,
- vec![0, 5, 1, 4, 2, 3],
- );
-
- // boolean, descending
- test_sort_to_indices_boolean_arrays(
- vec![None, Some(false), Some(true), Some(true), Some(false), None],
- Some(SortOptions {
- descending: true,
- nulls_first: false,
- }),
- None,
- vec![2, 3, 1, 4, 5, 0],
- );
-
- // boolean, descending, nulls first
- test_sort_to_indices_boolean_arrays(
- vec![None, Some(false), Some(true), Some(true), Some(false), None],
- Some(SortOptions {
- descending: true,
- nulls_first: true,
- }),
- None,
- vec![5, 0, 2, 3, 1, 4],
- );
-
- // boolean, descending, nulls first, limit
- test_sort_to_indices_boolean_arrays(
- vec![None, Some(false), Some(true), Some(true), Some(false), None],
- Some(SortOptions {
- descending: true,
- nulls_first: true,
- }),
- Some(3),
- vec![5, 0, 2],
- );
-
- // valid values less than limit with extra nulls
- test_sort_to_indices_boolean_arrays(
- vec![Some(true), None, None, Some(false)],
- Some(SortOptions {
- descending: false,
- nulls_first: false,
- }),
- Some(3),
- vec![3, 0, 1],
- );
-
- test_sort_to_indices_boolean_arrays(
- vec![Some(true), None, None, Some(false)],
- Some(SortOptions {
- descending: false,
- nulls_first: true,
- }),
- Some(3),
- vec![1, 2, 3],
- );
-
- // more nulls than limit
- test_sort_to_indices_boolean_arrays(
- vec![Some(true), None, None, None],
- Some(SortOptions {
- descending: false,
- nulls_first: true,
- }),
- Some(2),
- vec![1, 2],
- );
-
- test_sort_to_indices_boolean_arrays(
- vec![Some(true), None, None, None],
- Some(SortOptions {
- descending: false,
- nulls_first: false,
- }),
- Some(2),
- vec![0, 1],
- );
- }
-
- #[test]
- fn test_sort_primitives() {
- // default case
- test_sort_primitive_arrays::<UInt8Type>(
- vec![None, Some(3), Some(5), Some(2), Some(3), None],
- None,
- None,
- vec![None, None, Some(2), Some(3), Some(3), Some(5)],
- );
- test_sort_primitive_arrays::<UInt16Type>(
- vec![None, Some(3), Some(5), Some(2), Some(3), None],
- None,
- None,
- vec![None, None, Some(2), Some(3), Some(3), Some(5)],
- );
- test_sort_primitive_arrays::<UInt32Type>(
- vec![None, Some(3), Some(5), Some(2), Some(3), None],
- None,
- None,
- vec![None, None, Some(2), Some(3), Some(3), Some(5)],
- );
- test_sort_primitive_arrays::<UInt64Type>(
- vec![None, Some(3), Some(5), Some(2), Some(3), None],
- None,
- None,
- vec![None, None, Some(2), Some(3), Some(3), Some(5)],
- );
-
- // descending
- test_sort_primitive_arrays::<Int8Type>(
- vec![None, Some(0), Some(2), Some(-1), Some(0), None],
- Some(SortOptions {
- descending: true,
- nulls_first: false,
- }),
- None,
- vec![Some(2), Some(0), Some(0), Some(-1), None, None],
- );
- test_sort_primitive_arrays::<Int16Type>(
- vec![None, Some(0), Some(2), Some(-1), Some(0), None],
- Some(SortOptions {
- descending: true,
- nulls_first: false,
- }),
- None,
- vec![Some(2), Some(0), Some(0), Some(-1), None, None],
- );
- test_sort_primitive_arrays::<Int32Type>(
- vec![None, Some(0), Some(2), Some(-1), Some(0), None],
- Some(SortOptions {
- descending: true,
- nulls_first: false,
- }),
- None,
- vec![Some(2), Some(0), Some(0), Some(-1), None, None],
- );
- test_sort_primitive_arrays::<Int16Type>(
- vec![None, Some(0), Some(2), Some(-1), Some(0), None],
- Some(SortOptions {
- descending: true,
- nulls_first: false,
- }),
- None,
- vec![Some(2), Some(0), Some(0), Some(-1), None, None],
- );
-
- // descending, nulls first
- test_sort_primitive_arrays::<Int8Type>(
- vec![None, Some(0), Some(2), Some(-1), Some(0), None],
- Some(SortOptions {
- descending: true,
- nulls_first: true,
- }),
- None,
- vec![None, None, Some(2), Some(0), Some(0), Some(-1)],
- );
- test_sort_primitive_arrays::<Int16Type>(
- vec![None, Some(0), Some(2), Some(-1), Some(0), None],
- Some(SortOptions {
- descending: true,
- nulls_first: true,
- }),
- None,
- vec![None, None, Some(2), Some(0), Some(0), Some(-1)],
- );
- test_sort_primitive_arrays::<Int32Type>(
- vec![None, Some(0), Some(2), Some(-1), Some(0), None],
- Some(SortOptions {
- descending: true,
- nulls_first: true,
- }),
- None,
- vec![None, None, Some(2), Some(0), Some(0), Some(-1)],
- );
- test_sort_primitive_arrays::<Int64Type>(
- vec![None, Some(0), Some(2), Some(-1), Some(0), None],
- Some(SortOptions {
- descending: true,
- nulls_first: true,
- }),
- None,
- vec![None, None, Some(2), Some(0), Some(0), Some(-1)],
- );
-
- test_sort_primitive_arrays::<Int64Type>(
- vec![None, Some(0), Some(2), Some(-1), Some(0), None],
- Some(SortOptions {
- descending: true,
- nulls_first: true,
- }),
- Some(3),
- vec![None, None, Some(2)],
- );
-
- test_sort_primitive_arrays::<Float32Type>(
- vec![None, Some(0.0), Some(2.0), Some(-1.0), Some(0.0), None],
- Some(SortOptions {
- descending: true,
- nulls_first: true,
- }),
- None,
- vec![None, None, Some(2.0), Some(0.0), Some(0.0), Some(-1.0)],
- );
- test_sort_primitive_arrays::<Float64Type>(
- vec![None, Some(0.0), Some(2.0), Some(-1.0), Some(f64::NAN), None],
- Some(SortOptions {
- descending: true,
- nulls_first: true,
- }),
- None,
- vec![None, None, Some(f64::NAN), Some(2.0), Some(0.0), Some(-1.0)],
- );
- test_sort_primitive_arrays::<Float64Type>(
- vec![Some(f64::NAN), Some(f64::NAN), Some(f64::NAN), Some(1.0)],
- Some(SortOptions {
- descending: true,
- nulls_first: true,
- }),
- None,
- vec![Some(f64::NAN), Some(f64::NAN), Some(f64::NAN), Some(1.0)],
- );
-
- // int8 nulls first
- test_sort_primitive_arrays::<Int8Type>(
- vec![None, Some(0), Some(2), Some(-1), Some(0), None],
- Some(SortOptions {
- descending: false,
- nulls_first: true,
- }),
- None,
- vec![None, None, Some(-1), Some(0), Some(0), Some(2)],
- );
- test_sort_primitive_arrays::<Int16Type>(
- vec![None, Some(0), Some(2), Some(-1), Some(0), None],
- Some(SortOptions {
- descending: false,
- nulls_first: true,
- }),
- None,
- vec![None, None, Some(-1), Some(0), Some(0), Some(2)],
- );
- test_sort_primitive_arrays::<Int32Type>(
- vec![None, Some(0), Some(2), Some(-1), Some(0), None],
- Some(SortOptions {
- descending: false,
- nulls_first: true,
- }),
- None,
- vec![None, None, Some(-1), Some(0), Some(0), Some(2)],
- );
- test_sort_primitive_arrays::<Int64Type>(
- vec![None, Some(0), Some(2), Some(-1), Some(0), None],
- Some(SortOptions {
- descending: false,
- nulls_first: true,
- }),
- None,
- vec![None, None, Some(-1), Some(0), Some(0), Some(2)],
- );
- test_sort_primitive_arrays::<Float32Type>(
- vec![None, Some(0.0), Some(2.0), Some(-1.0), Some(0.0), None],
- Some(SortOptions {
- descending: false,
- nulls_first: true,
- }),
- None,
- vec![None, None, Some(-1.0), Some(0.0), Some(0.0), Some(2.0)],
- );
- test_sort_primitive_arrays::<Float64Type>(
- vec![None, Some(0.0), Some(2.0), Some(-1.0), Some(f64::NAN), None],
- Some(SortOptions {
- descending: false,
- nulls_first: true,
- }),
- None,
- vec![None, None, Some(-1.0), Some(0.0), Some(2.0), Some(f64::NAN)],
- );
- test_sort_primitive_arrays::<Float64Type>(
- vec![Some(f64::NAN), Some(f64::NAN), Some(f64::NAN), Some(1.0)],
- Some(SortOptions {
- descending: false,
- nulls_first: true,
- }),
- None,
- vec![Some(1.0), Some(f64::NAN), Some(f64::NAN), Some(f64::NAN)],
- );
-
- // limit
- test_sort_primitive_arrays::<Float64Type>(
- vec![Some(f64::NAN), Some(f64::NAN), Some(f64::NAN), Some(1.0)],
- Some(SortOptions {
- descending: false,
- nulls_first: true,
- }),
- Some(2),
- vec![Some(1.0), Some(f64::NAN)],
- );
-
- // limit with actual value
- test_sort_primitive_arrays::<Float64Type>(
- vec![Some(2.0), Some(4.0), Some(3.0), Some(1.0)],
- Some(SortOptions {
- descending: false,
- nulls_first: true,
- }),
- Some(3),
- vec![Some(1.0), Some(2.0), Some(3.0)],
- );
-
- // valid values less than limit with extra nulls
- test_sort_primitive_arrays::<Float64Type>(
- vec![Some(2.0), None, None, Some(1.0)],
- Some(SortOptions {
- descending: false,
- nulls_first: false,
- }),
- Some(3),
- vec![Some(1.0), Some(2.0), None],
- );
-
- test_sort_primitive_arrays::<Float64Type>(
- vec![Some(2.0), None, None, Some(1.0)],
- Some(SortOptions {
- descending: false,
- nulls_first: true,
- }),
- Some(3),
- vec![None, None, Some(1.0)],
- );
-
- // more nulls than limit
- test_sort_primitive_arrays::<Float64Type>(
- vec![Some(2.0), None, None, None],
- Some(SortOptions {
- descending: false,
- nulls_first: true,
- }),
- Some(2),
- vec![None, None],
- );
-
- test_sort_primitive_arrays::<Float64Type>(
- vec![Some(2.0), None, None, None],
- Some(SortOptions {
- descending: false,
- nulls_first: false,
- }),
- Some(2),
- vec![Some(2.0), None],
- );
- }
-
- #[test]
- fn test_sort_to_indices_strings() {
- test_sort_to_indices_string_arrays(
- vec![
- None,
- Some("bad"),
- Some("sad"),
- None,
- Some("glad"),
- Some("-ad"),
- ],
- None,
- None,
- vec![0, 3, 5, 1, 4, 2],
- );
-
- test_sort_to_indices_string_arrays(
- vec![
- None,
- Some("bad"),
- Some("sad"),
- None,
- Some("glad"),
- Some("-ad"),
- ],
- Some(SortOptions {
- descending: true,
- nulls_first: false,
- }),
- None,
- vec![2, 4, 1, 5, 3, 0],
- );
-
- test_sort_to_indices_string_arrays(
- vec![
- None,
- Some("bad"),
- Some("sad"),
- None,
- Some("glad"),
- Some("-ad"),
- ],
- Some(SortOptions {
- descending: false,
- nulls_first: true,
- }),
- None,
- vec![0, 3, 5, 1, 4, 2],
- );
-
- test_sort_to_indices_string_arrays(
- vec![
- None,
- Some("bad"),
- Some("sad"),
- None,
- Some("glad"),
- Some("-ad"),
- ],
- Some(SortOptions {
- descending: true,
- nulls_first: true,
- }),
- None,
- vec![3, 0, 2, 4, 1, 5],
- );
-
- test_sort_to_indices_string_arrays(
- vec![
- None,
- Some("bad"),
- Some("sad"),
- None,
- Some("glad"),
- Some("-ad"),
- ],
- Some(SortOptions {
- descending: true,
- nulls_first: true,
- }),
- Some(3),
- vec![3, 0, 2],
- );
-
- // valid values less than limit with extra nulls
- test_sort_to_indices_string_arrays(
- vec![Some("def"), None, None, Some("abc")],
- Some(SortOptions {
- descending: false,
- nulls_first: false,
- }),
- Some(3),
- vec![3, 0, 1],
- );
-
- test_sort_to_indices_string_arrays(
- vec![Some("def"), None, None, Some("abc")],
- Some(SortOptions {
- descending: false,
- nulls_first: true,
- }),
- Some(3),
- vec![1, 2, 3],
- );
-
- // more nulls than limit
- test_sort_to_indices_string_arrays(
- vec![Some("def"), None, None, None],
- Some(SortOptions {
- descending: false,
- nulls_first: true,
- }),
- Some(2),
- vec![1, 2],
- );
-
- test_sort_to_indices_string_arrays(
- vec![Some("def"), None, None, None],
- Some(SortOptions {
- descending: false,
- nulls_first: false,
- }),
- Some(2),
- vec![0, 1],
- );
- }
-
- #[test]
- fn test_sort_strings() {
- test_sort_string_arrays(
- vec![
- None,
- Some("bad"),
- Some("sad"),
- None,
- Some("glad"),
- Some("-ad"),
- ],
- None,
- None,
- vec![
- None,
- None,
- Some("-ad"),
- Some("bad"),
- Some("glad"),
- Some("sad"),
- ],
- );
-
- test_sort_string_arrays(
- vec![
- None,
- Some("bad"),
- Some("sad"),
- None,
- Some("glad"),
- Some("-ad"),
- ],
- Some(SortOptions {
- descending: true,
- nulls_first: false,
- }),
- None,
- vec![
- Some("sad"),
- Some("glad"),
- Some("bad"),
- Some("-ad"),
- None,
- None,
- ],
- );
-
- test_sort_string_arrays(
- vec![
- None,
- Some("bad"),
- Some("sad"),
- None,
- Some("glad"),
- Some("-ad"),
- ],
- Some(SortOptions {
- descending: false,
- nulls_first: true,
- }),
- None,
- vec![
- None,
- None,
- Some("-ad"),
- Some("bad"),
- Some("glad"),
- Some("sad"),
- ],
- );
-
- test_sort_string_arrays(
- vec![
- None,
- Some("bad"),
- Some("sad"),
- None,
- Some("glad"),
- Some("-ad"),
- ],
- Some(SortOptions {
- descending: true,
- nulls_first: true,
- }),
- None,
- vec![
- None,
- None,
- Some("sad"),
- Some("glad"),
- Some("bad"),
- Some("-ad"),
- ],
- );
-
- test_sort_string_arrays(
- vec![
- None,
- Some("bad"),
- Some("sad"),
- None,
- Some("glad"),
- Some("-ad"),
- ],
- Some(SortOptions {
- descending: true,
- nulls_first: true,
- }),
- Some(3),
- vec![None, None, Some("sad")],
- );
-
- // valid values less than limit with extra nulls
- test_sort_string_arrays(
- vec![Some("def"), None, None, Some("abc")],
- Some(SortOptions {
- descending: false,
- nulls_first: false,
- }),
- Some(3),
- vec![Some("abc"), Some("def"), None],
- );
-
- test_sort_string_arrays(
- vec![Some("def"), None, None, Some("abc")],
- Some(SortOptions {
- descending: false,
- nulls_first: true,
- }),
- Some(3),
- vec![None, None, Some("abc")],
- );
-
- // more nulls than limit
- test_sort_string_arrays(
- vec![Some("def"), None, None, None],
- Some(SortOptions {
- descending: false,
- nulls_first: true,
- }),
- Some(2),
- vec![None, None],
- );
-
- test_sort_string_arrays(
- vec![Some("def"), None, None, None],
- Some(SortOptions {
- descending: false,
- nulls_first: false,
- }),
- Some(2),
- vec![Some("def"), None],
- );
- }
-
- #[test]
- fn test_sort_string_dicts() {
- test_sort_string_dict_arrays::<Int8Type>(
- vec![
- None,
- Some("bad"),
- Some("sad"),
- None,
- Some("glad"),
- Some("-ad"),
- ],
- None,
- None,
- vec![
- None,
- None,
- Some("-ad"),
- Some("bad"),
- Some("glad"),
- Some("sad"),
- ],
- );
-
- test_sort_string_dict_arrays::<Int16Type>(
- vec![
- None,
- Some("bad"),
- Some("sad"),
- None,
- Some("glad"),
- Some("-ad"),
- ],
- Some(SortOptions {
- descending: true,
- nulls_first: false,
- }),
- None,
- vec![
- Some("sad"),
- Some("glad"),
- Some("bad"),
- Some("-ad"),
- None,
- None,
- ],
- );
-
- test_sort_string_dict_arrays::<Int32Type>(
- vec![
- None,
- Some("bad"),
- Some("sad"),
- None,
- Some("glad"),
- Some("-ad"),
- ],
- Some(SortOptions {
- descending: false,
- nulls_first: true,
- }),
- None,
- vec![
- None,
- None,
- Some("-ad"),
- Some("bad"),
- Some("glad"),
- Some("sad"),
- ],
- );
-
- test_sort_string_dict_arrays::<Int16Type>(
- vec![
- None,
- Some("bad"),
- Some("sad"),
- None,
- Some("glad"),
- Some("-ad"),
- ],
- Some(SortOptions {
- descending: true,
- nulls_first: true,
- }),
- None,
- vec![
- None,
- None,
- Some("sad"),
- Some("glad"),
- Some("bad"),
- Some("-ad"),
- ],
- );
-
- test_sort_string_dict_arrays::<Int16Type>(
- vec![
- None,
- Some("bad"),
- Some("sad"),
- None,
- Some("glad"),
- Some("-ad"),
- ],
- Some(SortOptions {
- descending: true,
- nulls_first: true,
- }),
- Some(3),
- vec![None, None, Some("sad")],
- );
-
- // valid values less than limit with extra nulls
- test_sort_string_dict_arrays::<Int16Type>(
- vec![Some("def"), None, None, Some("abc")],
- Some(SortOptions {
- descending: false,
- nulls_first: false,
- }),
- Some(3),
- vec![Some("abc"), Some("def"), None],
- );
-
- test_sort_string_dict_arrays::<Int16Type>(
- vec![Some("def"), None, None, Some("abc")],
- Some(SortOptions {
- descending: false,
- nulls_first: true,
- }),
- Some(3),
- vec![None, None, Some("abc")],
- );
-
- // more nulls than limit
- test_sort_string_dict_arrays::<Int16Type>(
- vec![Some("def"), None, None, None],
- Some(SortOptions {
- descending: false,
- nulls_first: true,
- }),
- Some(2),
- vec![None, None],
- );
-
- test_sort_string_dict_arrays::<Int16Type>(
- vec![Some("def"), None, None, None],
- Some(SortOptions {
- descending: false,
- nulls_first: false,
- }),
- Some(2),
- vec![Some("def"), None],
- );
- }
-
- #[test]
- fn test_sort_list() {
- test_sort_list_arrays::<Int8Type>(
- vec![
- Some(vec![Some(1)]),
- Some(vec![Some(4)]),
- Some(vec![Some(2)]),
- Some(vec![Some(3)]),
- ],
- Some(SortOptions {
- descending: false,
- nulls_first: false,
- }),
- None,
- vec![
- Some(vec![Some(1)]),
- Some(vec![Some(2)]),
- Some(vec![Some(3)]),
- Some(vec![Some(4)]),
- ],
- Some(1),
- );
-
- test_sort_list_arrays::<Float32Type>(
- vec![
- Some(vec![Some(1.0), Some(0.0)]),
- Some(vec![Some(4.0), Some(3.0), Some(2.0), Some(1.0)]),
- Some(vec![Some(2.0), Some(3.0), Some(4.0)]),
- Some(vec![Some(3.0), Some(3.0), Some(3.0), Some(3.0)]),
- Some(vec![Some(1.0), Some(1.0)]),
- ],
- Some(SortOptions {
- descending: false,
- nulls_first: false,
- }),
- None,
- vec![
- Some(vec![Some(1.0), Some(0.0)]),
- Some(vec![Some(1.0), Some(1.0)]),
- Some(vec![Some(2.0), Some(3.0), Some(4.0)]),
- Some(vec![Some(3.0), Some(3.0), Some(3.0), Some(3.0)]),
- Some(vec![Some(4.0), Some(3.0), Some(2.0), Some(1.0)]),
- ],
- None,
- );
-
- test_sort_list_arrays::<Float64Type>(
- vec![
- Some(vec![Some(1.0), Some(0.0)]),
- Some(vec![Some(4.0), Some(3.0), Some(2.0), Some(1.0)]),
- Some(vec![Some(2.0), Some(3.0), Some(4.0)]),
- Some(vec![Some(3.0), Some(3.0), Some(3.0), Some(3.0)]),
- Some(vec![Some(1.0), Some(1.0)]),
- ],
- Some(SortOptions {
- descending: false,
- nulls_first: false,
- }),
- None,
- vec![
- Some(vec![Some(1.0), Some(0.0)]),
- Some(vec![Some(1.0), Some(1.0)]),
- Some(vec![Some(2.0), Some(3.0), Some(4.0)]),
- Some(vec![Some(3.0), Some(3.0), Some(3.0), Some(3.0)]),
- Some(vec![Some(4.0), Some(3.0), Some(2.0), Some(1.0)]),
- ],
- None,
- );
-
- test_sort_list_arrays::<Int32Type>(
- vec![
- Some(vec![Some(1), Some(0)]),
- Some(vec![Some(4), Some(3), Some(2), Some(1)]),
- Some(vec![Some(2), Some(3), Some(4)]),
- Some(vec![Some(3), Some(3), Some(3), Some(3)]),
- Some(vec![Some(1), Some(1)]),
- ],
- Some(SortOptions {
- descending: false,
- nulls_first: false,
- }),
- None,
- vec![
- Some(vec![Some(1), Some(0)]),
- Some(vec![Some(1), Some(1)]),
- Some(vec![Some(2), Some(3), Some(4)]),
- Some(vec![Some(3), Some(3), Some(3), Some(3)]),
- Some(vec![Some(4), Some(3), Some(2), Some(1)]),
- ],
- None,
- );
-
- test_sort_list_arrays::<Int32Type>(
- vec![
- None,
- Some(vec![Some(4), None, Some(2)]),
- Some(vec![Some(2), Some(3), Some(4)]),
- None,
- Some(vec![Some(3), Some(3), None]),
- ],
- Some(SortOptions {
- descending: false,
- nulls_first: false,
- }),
- None,
- vec![
- Some(vec![Some(2), Some(3), Some(4)]),
- Some(vec![Some(3), Some(3), None]),
- Some(vec![Some(4), None, Some(2)]),
- None,
- None,
- ],
- Some(3),
- );
-
- test_sort_list_arrays::<Int32Type>(
- vec![
- Some(vec![Some(1), Some(0)]),
- Some(vec![Some(4), Some(3), Some(2), Some(1)]),
- Some(vec![Some(2), Some(3), Some(4)]),
- Some(vec![Some(3), Some(3), Some(3), Some(3)]),
- Some(vec![Some(1), Some(1)]),
- ],
- Some(SortOptions {
- descending: false,
- nulls_first: false,
- }),
- Some(2),
- vec![Some(vec![Some(1), Some(0)]), Some(vec![Some(1), Some(1)])],
- None,
- );
-
- // valid values less than limit with extra nulls
- test_sort_list_arrays::<Int32Type>(
- vec![Some(vec![Some(1)]), None, None, Some(vec![Some(2)])],
- Some(SortOptions {
- descending: false,
- nulls_first: false,
- }),
- Some(3),
- vec![Some(vec![Some(1)]), Some(vec![Some(2)]), None],
- None,
- );
-
- test_sort_list_arrays::<Int32Type>(
- vec![Some(vec![Some(1)]), None, None, Some(vec![Some(2)])],
- Some(SortOptions {
- descending: false,
- nulls_first: true,
- }),
- Some(3),
- vec![None, None, Some(vec![Some(2)])],
- None,
- );
-
- // more nulls than limit
- test_sort_list_arrays::<Int32Type>(
- vec![Some(vec![Some(1)]), None, None, None],
- Some(SortOptions {
- descending: false,
- nulls_first: true,
- }),
- Some(2),
- vec![None, None],
- None,
- );
-
- test_sort_list_arrays::<Int32Type>(
- vec![Some(vec![Some(1)]), None, None, None],
- Some(SortOptions {
- descending: false,
- nulls_first: false,
- }),
- Some(2),
- vec![Some(vec![Some(1)]), None],
- None,
- );
- }
-
- #[test]
- fn test_lex_sort_single_column() {
- let input = vec![SortColumn {
- values: Arc::new(PrimitiveArray::<Int64Type>::from(vec![
- Some(17),
- Some(2),
- Some(-1),
- Some(0),
- ])) as ArrayRef,
- options: None,
- }];
- let expected = vec![Arc::new(PrimitiveArray::<Int64Type>::from(vec![
- Some(-1),
- Some(0),
- Some(2),
- Some(17),
- ])) as ArrayRef];
- test_lex_sort_arrays(input.clone(), expected, None);
-
- let expected = vec![Arc::new(PrimitiveArray::<Int64Type>::from(vec![
- Some(-1),
- Some(0),
- Some(2),
- ])) as ArrayRef];
- test_lex_sort_arrays(input, expected, Some(3));
- }
-
- #[test]
- fn test_lex_sort_unaligned_rows() {
- let input = vec![
- SortColumn {
- values: Arc::new(PrimitiveArray::<Int64Type>::from(vec![None, Some(-1)]))
- as ArrayRef,
- options: None,
- },
- SortColumn {
- values: Arc::new(StringArray::from(vec![Some("foo")])) as ArrayRef,
- options: None,
- },
- ];
- assert!(
- lexsort(&input, None).is_err(),
- "lexsort should reject columns with different row counts"
- );
- }
-
- #[test]
- fn test_lex_sort_mixed_types() {
- let input = vec![
- SortColumn {
- values: Arc::new(PrimitiveArray::<Int64Type>::from(vec![
- Some(0),
- Some(2),
- Some(-1),
- Some(0),
- ])) as ArrayRef,
- options: None,
- },
- SortColumn {
- values: Arc::new(PrimitiveArray::<UInt32Type>::from(vec![
- Some(101),
- Some(8),
- Some(7),
- Some(102),
- ])) as ArrayRef,
- options: None,
- },
- SortColumn {
- values: Arc::new(PrimitiveArray::<Int64Type>::from(vec![
- Some(-1),
- Some(-2),
- Some(-3),
- Some(-4),
- ])) as ArrayRef,
- options: None,
- },
- ];
- let expected = vec![
- Arc::new(PrimitiveArray::<Int64Type>::from(vec![
- Some(-1),
- Some(0),
- Some(0),
- Some(2),
- ])) as ArrayRef,
- Arc::new(PrimitiveArray::<UInt32Type>::from(vec![
- Some(7),
- Some(101),
- Some(102),
- Some(8),
- ])) as ArrayRef,
- Arc::new(PrimitiveArray::<Int64Type>::from(vec![
- Some(-3),
- Some(-1),
- Some(-4),
- Some(-2),
- ])) as ArrayRef,
- ];
- test_lex_sort_arrays(input, expected, None);
-
- // test mix of string and in64 with option
- let input = vec![
- SortColumn {
- values: Arc::new(PrimitiveArray::<Int64Type>::from(vec![
- Some(0),
- Some(2),
- Some(-1),
- Some(0),
- ])) as ArrayRef,
- options: Some(SortOptions {
- descending: true,
- nulls_first: true,
- }),
- },
- SortColumn {
- values: Arc::new(StringArray::from(vec![
- Some("foo"),
- Some("9"),
- Some("7"),
- Some("bar"),
- ])) as ArrayRef,
- options: Some(SortOptions {
- descending: true,
- nulls_first: true,
- }),
- },
- ];
- let expected = vec![
- Arc::new(PrimitiveArray::<Int64Type>::from(vec![
- Some(2),
- Some(0),
- Some(0),
- Some(-1),
- ])) as ArrayRef,
- Arc::new(StringArray::from(vec![
- Some("9"),
- Some("foo"),
- Some("bar"),
- Some("7"),
- ])) as ArrayRef,
- ];
- test_lex_sort_arrays(input, expected, None);
-
- // test sort with nulls first
- let input = vec![
- SortColumn {
- values: Arc::new(PrimitiveArray::<Int64Type>::from(vec![
- None,
- Some(-1),
- Some(2),
- None,
- ])) as ArrayRef,
- options: Some(SortOptions {
- descending: true,
- nulls_first: true,
- }),
- },
- SortColumn {
- values: Arc::new(StringArray::from(vec![
- Some("foo"),
- Some("world"),
- Some("hello"),
- None,
- ])) as ArrayRef,
- options: Some(SortOptions {
- descending: true,
- nulls_first: true,
- }),
- },
- ];
- let expected = vec![
- Arc::new(PrimitiveArray::<Int64Type>::from(vec![
- None,
- None,
- Some(2),
- Some(-1),
- ])) as ArrayRef,
- Arc::new(StringArray::from(vec![
- None,
- Some("foo"),
- Some("hello"),
- Some("world"),
- ])) as ArrayRef,
- ];
- test_lex_sort_arrays(input, expected, None);
-
- // test sort with nulls last
- let input = vec![
- SortColumn {
- values: Arc::new(PrimitiveArray::<Int64Type>::from(vec![
- None,
- Some(-1),
- Some(2),
- None,
- ])) as ArrayRef,
- options: Some(SortOptions {
- descending: true,
- nulls_first: false,
- }),
- },
- SortColumn {
- values: Arc::new(StringArray::from(vec![
- Some("foo"),
- Some("world"),
- Some("hello"),
- None,
- ])) as ArrayRef,
- options: Some(SortOptions {
- descending: true,
- nulls_first: false,
- }),
- },
- ];
- let expected = vec![
- Arc::new(PrimitiveArray::<Int64Type>::from(vec![
- Some(2),
- Some(-1),
- None,
- None,
- ])) as ArrayRef,
- Arc::new(StringArray::from(vec![
- Some("hello"),
- Some("world"),
- Some("foo"),
- None,
- ])) as ArrayRef,
- ];
- test_lex_sort_arrays(input, expected, None);
-
- // test sort with opposite options
- let input = vec![
- SortColumn {
- values: Arc::new(PrimitiveArray::<Int64Type>::from(vec![
- None,
- Some(-1),
- Some(2),
- Some(-1),
- None,
- ])) as ArrayRef,
- options: Some(SortOptions {
- descending: false,
- nulls_first: false,
- }),
- },
- SortColumn {
- values: Arc::new(StringArray::from(vec![
- Some("foo"),
- Some("bar"),
- Some("world"),
- Some("hello"),
- None,
- ])) as ArrayRef,
- options: Some(SortOptions {
- descending: true,
- nulls_first: true,
- }),
- },
- ];
- let expected = vec![
- Arc::new(PrimitiveArray::<Int64Type>::from(vec![
- Some(-1),
- Some(-1),
- Some(2),
- None,
- None,
- ])) as ArrayRef,
- Arc::new(StringArray::from(vec![
- Some("hello"),
- Some("bar"),
- Some("world"),
- None,
- Some("foo"),
- ])) as ArrayRef,
- ];
- test_lex_sort_arrays(input, expected, None);
- }
-
- #[test]
- fn test_partial_sort() {
- let mut before: Vec<&str> = vec![
- "a", "cat", "mat", "on", "sat", "the", "xxx", "xxxx", "fdadfdsf",
- ];
- let mut d = before.clone();
- d.sort_unstable();
-
- for last in 0..before.len() {
- partial_sort(&mut before, last, |a, b| a.cmp(b));
- assert_eq!(&d[0..last], &before.as_slice()[0..last]);
- }
- }
-
- #[test]
- fn test_partial_rand_sort() {
- let size = 1000u32;
- let mut rng = StdRng::seed_from_u64(42);
- let mut before: Vec<u32> = (0..size).map(|_| rng.gen::<u32>()).collect();
- let mut d = before.clone();
- let last = (rng.next_u32() % size) as usize;
- d.sort_unstable();
-
- partial_sort(&mut before, last, |a, b| a.cmp(b));
- assert_eq!(&d[0..last], &before[0..last]);
- }
-}
diff --git a/arrow/src/compute/kernels/substring.rs b/arrow/src/compute/kernels/substring.rs
deleted file mode 100644
index d9956b8..0000000
--- a/arrow/src/compute/kernels/substring.rs
+++ /dev/null
@@ -1,269 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-//! Defines kernel to extract a substring of a \[Large\]StringArray
-
-use crate::{array::*, buffer::Buffer};
-use crate::{
- datatypes::DataType,
- error::{ArrowError, Result},
-};
-
-#[allow(clippy::unnecessary_wraps)]
-fn generic_substring<OffsetSize: StringOffsetSizeTrait>(
- array: &GenericStringArray<OffsetSize>,
- start: OffsetSize,
- length: &Option<OffsetSize>,
-) -> Result<ArrayRef> {
- // compute current offsets
- let offsets = array.data_ref().clone().buffers()[0].clone();
- let offsets: &[OffsetSize] = unsafe { offsets.typed_data::<OffsetSize>() };
-
- // compute null bitmap (copy)
- let null_bit_buffer = array.data_ref().null_buffer().cloned();
-
- // compute values
- let values = &array.data_ref().buffers()[1];
- let data = values.as_slice();
-
- let mut new_values = Vec::new(); // we have no way to estimate how much this will be.
- let mut new_offsets: Vec<OffsetSize> = Vec::with_capacity(array.len() + 1);
-
- let mut length_so_far = OffsetSize::zero();
- new_offsets.push(length_so_far);
- (0..array.len()).for_each(|i| {
- // the length of this entry
- let length_i: OffsetSize = offsets[i + 1] - offsets[i];
- // compute where we should start slicing this entry
- let start = offsets[i]
- + if start >= OffsetSize::zero() {
- start
- } else {
- length_i + start
- };
-
- let start = start.max(offsets[i]).min(offsets[i + 1]);
- // compute the length of the slice
- let length: OffsetSize = length
- .unwrap_or(length_i)
- // .max(0) is not needed as it is guaranteed
- .min(offsets[i + 1] - start); // so we do not go beyond this entry
-
- length_so_far += length;
-
- new_offsets.push(length_so_far);
-
- // we need usize for ranges
- let start = start.to_usize().unwrap();
- let length = length.to_usize().unwrap();
-
- new_values.extend_from_slice(&data[start..start + length]);
- });
-
- let data = ArrayData::new(
- <OffsetSize as StringOffsetSizeTrait>::DATA_TYPE,
- array.len(),
- None,
- null_bit_buffer,
- 0,
- vec![
- Buffer::from_slice_ref(&new_offsets),
- Buffer::from_slice_ref(&new_values),
- ],
- vec![],
- );
- Ok(make_array(data))
-}
-
-/// Returns an ArrayRef with a substring starting from `start` and with optional length `length` of each of the elements in `array`.
-/// `start` can be negative, in which case the start counts from the end of the string.
-/// this function errors when the passed array is not a \[Large\]String array.
-pub fn substring(array: &Array, start: i64, length: &Option<u64>) -> Result<ArrayRef> {
- match array.data_type() {
- DataType::LargeUtf8 => generic_substring(
- array
- .as_any()
- .downcast_ref::<LargeStringArray>()
- .expect("A large string is expected"),
- start,
- &length.map(|e| e as i64),
- ),
- DataType::Utf8 => generic_substring(
- array
- .as_any()
- .downcast_ref::<StringArray>()
- .expect("A string is expected"),
- start as i32,
- &length.map(|e| e as i32),
- ),
- _ => Err(ArrowError::ComputeError(format!(
- "substring does not support type {:?}",
- array.data_type()
- ))),
- }
-}
-
-#[cfg(test)]
-mod tests {
- use super::*;
-
- fn with_nulls<T: 'static + Array + PartialEq + From<Vec<Option<&'static str>>>>(
- ) -> Result<()> {
- let cases = vec![
- // identity
- (
- vec![Some("hello"), None, Some("word")],
- 0,
- None,
- vec![Some("hello"), None, Some("word")],
- ),
- // 0 length -> Nothing
- (
- vec![Some("hello"), None, Some("word")],
- 0,
- Some(0),
- vec![Some(""), None, Some("")],
- ),
- // high start -> Nothing
- (
- vec![Some("hello"), None, Some("word")],
- 1000,
- Some(0),
- vec![Some(""), None, Some("")],
- ),
- // high negative start -> identity
- (
- vec![Some("hello"), None, Some("word")],
- -1000,
- None,
- vec![Some("hello"), None, Some("word")],
- ),
- // high length -> identity
- (
- vec![Some("hello"), None, Some("word")],
- 0,
- Some(1000),
- vec![Some("hello"), None, Some("word")],
- ),
- ];
-
- cases.into_iter().try_for_each::<_, Result<()>>(
- |(array, start, length, expected)| {
- let array = T::from(array);
- let result: ArrayRef = substring(&array, start, &length)?;
- assert_eq!(array.len(), result.len());
-
- let result = result.as_any().downcast_ref::<T>().unwrap();
- let expected = T::from(expected);
- assert_eq!(&expected, result);
- Ok(())
- },
- )?;
-
- Ok(())
- }
-
- #[test]
- fn with_nulls_string() -> Result<()> {
- with_nulls::<StringArray>()
- }
-
- #[test]
- fn with_nulls_large_string() -> Result<()> {
- with_nulls::<LargeStringArray>()
- }
-
- fn without_nulls<T: 'static + Array + PartialEq + From<Vec<Option<&'static str>>>>(
- ) -> Result<()> {
- let cases = vec![
- // increase start
- (
- vec!["hello", "", "word"],
- 0,
- None,
- vec!["hello", "", "word"],
- ),
- (vec!["hello", "", "word"], 1, None, vec!["ello", "", "ord"]),
- (vec!["hello", "", "word"], 2, None, vec!["llo", "", "rd"]),
- (vec!["hello", "", "word"], 3, None, vec!["lo", "", "d"]),
- (vec!["hello", "", "word"], 10, None, vec!["", "", ""]),
- // increase start negatively
- (vec!["hello", "", "word"], -1, None, vec!["o", "", "d"]),
- (vec!["hello", "", "word"], -2, None, vec!["lo", "", "rd"]),
- (vec!["hello", "", "word"], -3, None, vec!["llo", "", "ord"]),
- (
- vec!["hello", "", "word"],
- -10,
- None,
- vec!["hello", "", "word"],
- ),
- // increase length
- (vec!["hello", "", "word"], 1, Some(1), vec!["e", "", "o"]),
- (vec!["hello", "", "word"], 1, Some(2), vec!["el", "", "or"]),
- (
- vec!["hello", "", "word"],
- 1,
- Some(3),
- vec!["ell", "", "ord"],
- ),
- (
- vec!["hello", "", "word"],
- 1,
- Some(4),
- vec!["ello", "", "ord"],
- ),
- (vec!["hello", "", "word"], -3, Some(1), vec!["l", "", "o"]),
- (vec!["hello", "", "word"], -3, Some(2), vec!["ll", "", "or"]),
- (
- vec!["hello", "", "word"],
- -3,
- Some(3),
- vec!["llo", "", "ord"],
- ),
- (
- vec!["hello", "", "word"],
- -3,
- Some(4),
- vec!["llo", "", "ord"],
- ),
- ];
-
- cases.into_iter().try_for_each::<_, Result<()>>(
- |(array, start, length, expected)| {
- let array = StringArray::from(array);
- let result = substring(&array, start, &length)?;
- assert_eq!(array.len(), result.len());
- let result = result.as_any().downcast_ref::<StringArray>().unwrap();
- let expected = StringArray::from(expected);
- assert_eq!(&expected, result,);
- Ok(())
- },
- )?;
-
- Ok(())
- }
-
- #[test]
- fn without_nulls_string() -> Result<()> {
- without_nulls::<StringArray>()
- }
-
- #[test]
- fn without_nulls_large_string() -> Result<()> {
- without_nulls::<LargeStringArray>()
- }
-}
diff --git a/arrow/src/compute/kernels/take.rs b/arrow/src/compute/kernels/take.rs
deleted file mode 100644
index 66bfd61..0000000
--- a/arrow/src/compute/kernels/take.rs
+++ /dev/null
@@ -1,1634 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-//! Defines take kernel for [Array]
-
-use std::{ops::AddAssign, sync::Arc};
-
-use crate::buffer::{Buffer, MutableBuffer};
-use crate::compute::util::{
- take_value_indices_from_fixed_size_list, take_value_indices_from_list,
-};
-use crate::datatypes::*;
-use crate::error::{ArrowError, Result};
-use crate::util::bit_util;
-use crate::{array::*, buffer::buffer_bin_and};
-
-use num::{ToPrimitive, Zero};
-use TimeUnit::*;
-
-macro_rules! downcast_take {
- ($type: ty, $values: expr, $indices: expr) => {{
- let values = $values
- .as_any()
- .downcast_ref::<PrimitiveArray<$type>>()
- .expect("Unable to downcast to a primitive array");
- Ok(Arc::new(take_primitive::<$type, _>(&values, $indices)?))
- }};
-}
-
-macro_rules! downcast_dict_take {
- ($type: ty, $values: expr, $indices: expr) => {{
- let values = $values
- .as_any()
- .downcast_ref::<DictionaryArray<$type>>()
- .expect("Unable to downcast to a dictionary array");
- Ok(Arc::new(take_dict::<$type, _>(values, $indices)?))
- }};
-}
-
-/// Take elements by index from [Array], creating a new [Array] from those indexes.
-///
-/// # Errors
-/// This function errors whenever:
-/// * An index cannot be casted to `usize` (typically 32 bit architectures)
-/// * An index is out of bounds and `options` is set to check bounds.
-/// # Safety
-/// When `options` is not set to check bounds (default), taking indexes after `len` is undefined behavior.
-/// # Examples
-/// ```
-/// use arrow::array::{StringArray, UInt32Array};
-/// use arrow::error::Result;
-/// use arrow::compute::take;
-/// # fn main() -> Result<()> {
-/// let values = StringArray::from(vec!["zero", "one", "two"]);
-///
-/// // Take items at index 2, and 1:
-/// let indices = UInt32Array::from(vec![2, 1]);
-/// let taken = take(&values, &indices, None)?;
-/// let taken = taken.as_any().downcast_ref::<StringArray>().unwrap();
-///
-/// assert_eq!(*taken, StringArray::from(vec!["two", "one"]));
-/// # Ok(())
-/// # }
-/// ```
-pub fn take<IndexType>(
- values: &Array,
- indices: &PrimitiveArray<IndexType>,
- options: Option<TakeOptions>,
-) -> Result<ArrayRef>
-where
- IndexType: ArrowNumericType,
- IndexType::Native: ToPrimitive,
-{
- take_impl(values, indices, options)
-}
-
-fn take_impl<IndexType>(
- values: &Array,
- indices: &PrimitiveArray<IndexType>,
- options: Option<TakeOptions>,
-) -> Result<ArrayRef>
-where
- IndexType: ArrowNumericType,
- IndexType::Native: ToPrimitive,
-{
- let options = options.unwrap_or_default();
- if options.check_bounds {
- let len = values.len();
- if indices.null_count() > 0 {
- indices.iter().flatten().try_for_each(|index| {
- let ix = ToPrimitive::to_usize(&index).ok_or_else(|| {
- ArrowError::ComputeError("Cast to usize failed".to_string())
- })?;
- if ix >= len {
- return Err(ArrowError::ComputeError(
- format!("Array index out of bounds, cannot get item at index {} from {} entries", ix, len))
- );
- }
- Ok(())
- })?;
- } else {
- indices.values().iter().try_for_each(|index| {
- let ix = ToPrimitive::to_usize(index).ok_or_else(|| {
- ArrowError::ComputeError("Cast to usize failed".to_string())
- })?;
- if ix >= len {
- return Err(ArrowError::ComputeError(
- format!("Array index out of bounds, cannot get item at index {} from {} entries", ix, len))
- );
- }
- Ok(())
- })?
- }
- }
- match values.data_type() {
- DataType::Boolean => {
- let values = values.as_any().downcast_ref::<BooleanArray>().unwrap();
- Ok(Arc::new(take_boolean(values, indices)?))
- }
- DataType::Int8 => downcast_take!(Int8Type, values, indices),
- DataType::Int16 => downcast_take!(Int16Type, values, indices),
- DataType::Int32 => downcast_take!(Int32Type, values, indices),
- DataType::Int64 => downcast_take!(Int64Type, values, indices),
- DataType::UInt8 => downcast_take!(UInt8Type, values, indices),
- DataType::UInt16 => downcast_take!(UInt16Type, values, indices),
- DataType::UInt32 => downcast_take!(UInt32Type, values, indices),
- DataType::UInt64 => downcast_take!(UInt64Type, values, indices),
- DataType::Float32 => downcast_take!(Float32Type, values, indices),
- DataType::Float64 => downcast_take!(Float64Type, values, indices),
- DataType::Date32 => downcast_take!(Date32Type, values, indices),
- DataType::Date64 => downcast_take!(Date64Type, values, indices),
- DataType::Time32(Second) => downcast_take!(Time32SecondType, values, indices),
- DataType::Time32(Millisecond) => {
- downcast_take!(Time32MillisecondType, values, indices)
- }
- DataType::Time64(Microsecond) => {
- downcast_take!(Time64MicrosecondType, values, indices)
- }
- DataType::Time64(Nanosecond) => {
- downcast_take!(Time64NanosecondType, values, indices)
- }
- DataType::Timestamp(Second, _) => {
- downcast_take!(TimestampSecondType, values, indices)
- }
- DataType::Timestamp(Millisecond, _) => {
- downcast_take!(TimestampMillisecondType, values, indices)
- }
- DataType::Timestamp(Microsecond, _) => {
- downcast_take!(TimestampMicrosecondType, values, indices)
- }
- DataType::Timestamp(Nanosecond, _) => {
- downcast_take!(TimestampNanosecondType, values, indices)
- }
- DataType::Interval(IntervalUnit::YearMonth) => {
- downcast_take!(IntervalYearMonthType, values, indices)
- }
- DataType::Interval(IntervalUnit::DayTime) => {
- downcast_take!(IntervalDayTimeType, values, indices)
- }
- DataType::Duration(TimeUnit::Second) => {
- downcast_take!(DurationSecondType, values, indices)
- }
- DataType::Duration(TimeUnit::Millisecond) => {
- downcast_take!(DurationMillisecondType, values, indices)
- }
- DataType::Duration(TimeUnit::Microsecond) => {
- downcast_take!(DurationMicrosecondType, values, indices)
- }
- DataType::Duration(TimeUnit::Nanosecond) => {
- downcast_take!(DurationNanosecondType, values, indices)
- }
- DataType::Utf8 => {
- let values = values
- .as_any()
- .downcast_ref::<GenericStringArray<i32>>()
- .unwrap();
- Ok(Arc::new(take_string::<i32, _>(values, indices)?))
- }
- DataType::LargeUtf8 => {
- let values = values
- .as_any()
- .downcast_ref::<GenericStringArray<i64>>()
- .unwrap();
- Ok(Arc::new(take_string::<i64, _>(values, indices)?))
- }
- DataType::List(_) => {
- let values = values
- .as_any()
- .downcast_ref::<GenericListArray<i32>>()
- .unwrap();
- Ok(Arc::new(take_list::<_, Int32Type>(values, indices)?))
- }
- DataType::LargeList(_) => {
- let values = values
- .as_any()
- .downcast_ref::<GenericListArray<i64>>()
- .unwrap();
- Ok(Arc::new(take_list::<_, Int64Type>(values, indices)?))
- }
- DataType::FixedSizeList(_, length) => {
- let values = values
- .as_any()
- .downcast_ref::<FixedSizeListArray>()
- .unwrap();
- Ok(Arc::new(take_fixed_size_list(
- values,
- indices,
- *length as u32,
- )?))
- }
- DataType::Struct(fields) => {
- let struct_: &StructArray =
- values.as_any().downcast_ref::<StructArray>().unwrap();
- let arrays: Result<Vec<ArrayRef>> = struct_
- .columns()
- .iter()
- .map(|a| take_impl(a.as_ref(), indices, Some(options.clone())))
- .collect();
- let arrays = arrays?;
- let pairs: Vec<(Field, ArrayRef)> =
- fields.clone().into_iter().zip(arrays).collect();
- Ok(Arc::new(StructArray::from(pairs)) as ArrayRef)
- }
- DataType::Dictionary(key_type, _) => match key_type.as_ref() {
- DataType::Int8 => downcast_dict_take!(Int8Type, values, indices),
- DataType::Int16 => downcast_dict_take!(Int16Type, values, indices),
- DataType::Int32 => downcast_dict_take!(Int32Type, values, indices),
- DataType::Int64 => downcast_dict_take!(Int64Type, values, indices),
- DataType::UInt8 => downcast_dict_take!(UInt8Type, values, indices),
- DataType::UInt16 => downcast_dict_take!(UInt16Type, values, indices),
- DataType::UInt32 => downcast_dict_take!(UInt32Type, values, indices),
- DataType::UInt64 => downcast_dict_take!(UInt64Type, values, indices),
- t => unimplemented!("Take not supported for dictionary key type {:?}", t),
- },
- t => unimplemented!("Take not supported for data type {:?}", t),
- }
-}
-
-/// Options that define how `take` should behave
-#[derive(Clone, Debug)]
-pub struct TakeOptions {
- /// Perform bounds check before taking indices from values.
- /// If enabled, an `ArrowError` is returned if the indices are out of bounds.
- /// If not enabled, and indices exceed bounds, the kernel will panic.
- pub check_bounds: bool,
-}
-
-impl Default for TakeOptions {
- fn default() -> Self {
- Self {
- check_bounds: false,
- }
- }
-}
-
-#[inline(always)]
-fn maybe_usize<I: ArrowPrimitiveType>(index: I::Native) -> Result<usize> {
- index
- .to_usize()
- .ok_or_else(|| ArrowError::ComputeError("Cast to usize failed".to_string()))
-}
-
-// take implementation when neither values nor indices contain nulls
-fn take_no_nulls<T, I>(
- values: &[T::Native],
- indices: &[I::Native],
-) -> Result<(Buffer, Option<Buffer>)>
-where
- T: ArrowPrimitiveType,
- I: ArrowNumericType,
-{
- let values = indices
- .iter()
- .map(|index| Result::Ok(values[maybe_usize::<I>(*index)?]));
- // Soundness: `slice.map` is `TrustedLen`.
- let buffer = unsafe { Buffer::try_from_trusted_len_iter(values)? };
-
- Ok((buffer, None))
-}
-
-// take implementation when only values contain nulls
-fn take_values_nulls<T, I>(
- values: &PrimitiveArray<T>,
- indices: &[I::Native],
-) -> Result<(Buffer, Option<Buffer>)>
-where
- T: ArrowPrimitiveType,
- I: ArrowNumericType,
- I::Native: ToPrimitive,
-{
- let num_bytes = bit_util::ceil(indices.len(), 8);
- let mut nulls = MutableBuffer::new(num_bytes).with_bitset(num_bytes, true);
- let null_slice = nulls.as_slice_mut();
- let mut null_count = 0;
-
- let values_values = values.values();
-
- let values = indices.iter().enumerate().map(|(i, index)| {
- let index = maybe_usize::<I>(*index)?;
- if values.is_null(index) {
- null_count += 1;
- bit_util::unset_bit(null_slice, i);
- }
- Result::Ok(values_values[index])
- });
- // Soundness: `slice.map` is `TrustedLen`.
- let buffer = unsafe { Buffer::try_from_trusted_len_iter(values)? };
-
- let nulls = if null_count == 0 {
- // if only non-null values were taken
- None
- } else {
- Some(nulls.into())
- };
-
- Ok((buffer, nulls))
-}
-
-// take implementation when only indices contain nulls
-fn take_indices_nulls<T, I>(
- values: &[T::Native],
- indices: &PrimitiveArray<I>,
-) -> Result<(Buffer, Option<Buffer>)>
-where
- T: ArrowPrimitiveType,
- I: ArrowNumericType,
- I::Native: ToPrimitive,
-{
- let values = indices.values().iter().map(|index| {
- let index = maybe_usize::<I>(*index)?;
- Result::Ok(match values.get(index) {
- Some(value) => *value,
- None => {
- if indices.is_null(index) {
- T::Native::default()
- } else {
- panic!("Out-of-bounds index {}", index)
- }
- }
- })
- });
-
- // Soundness: `slice.map` is `TrustedLen`.
- let buffer = unsafe { Buffer::try_from_trusted_len_iter(values)? };
-
- Ok((buffer, indices.data_ref().null_buffer().cloned()))
-}
-
-// take implementation when both values and indices contain nulls
-fn take_values_indices_nulls<T, I>(
- values: &PrimitiveArray<T>,
- indices: &PrimitiveArray<I>,
-) -> Result<(Buffer, Option<Buffer>)>
-where
- T: ArrowPrimitiveType,
- I: ArrowNumericType,
- I::Native: ToPrimitive,
-{
- let num_bytes = bit_util::ceil(indices.len(), 8);
- let mut nulls = MutableBuffer::new(num_bytes).with_bitset(num_bytes, true);
- let null_slice = nulls.as_slice_mut();
- let mut null_count = 0;
-
- let values_values = values.values();
- let values = indices.iter().enumerate().map(|(i, index)| match index {
- Some(index) => {
- let index = maybe_usize::<I>(index)?;
- if values.is_null(index) {
- null_count += 1;
- bit_util::unset_bit(null_slice, i);
- }
- Result::Ok(values_values[index])
- }
- None => {
- null_count += 1;
- bit_util::unset_bit(null_slice, i);
- Ok(T::Native::default())
- }
- });
- // Soundness: `slice.map` is `TrustedLen`.
- let buffer = unsafe { Buffer::try_from_trusted_len_iter(values)? };
-
- let nulls = if null_count == 0 {
- // if only non-null values were taken
- None
- } else {
- Some(nulls.into())
- };
-
- Ok((buffer, nulls))
-}
-
-/// `take` implementation for all primitive arrays
-///
-/// This checks if an `indices` slot is populated, and gets the value from `values`
-/// as the populated index.
-/// If the `indices` slot is null, a null value is returned.
-/// For example, given:
-/// values: [1, 2, 3, null, 5]
-/// indices: [0, null, 4, 3]
-/// The result is: [1 (slot 0), null (null slot), 5 (slot 4), null (slot 3)]
-fn take_primitive<T, I>(
- values: &PrimitiveArray<T>,
- indices: &PrimitiveArray<I>,
-) -> Result<PrimitiveArray<T>>
-where
- T: ArrowPrimitiveType,
- I: ArrowNumericType,
- I::Native: ToPrimitive,
-{
- let indices_has_nulls = indices.null_count() > 0;
- let values_has_nulls = values.null_count() > 0;
- // note: this function should only panic when "an index is not null and out of bounds".
- // if the index is null, its value is undefined and therefore we should not read from it.
-
- let (buffer, nulls) = match (values_has_nulls, indices_has_nulls) {
- (false, false) => {
- // * no nulls
- // * all `indices.values()` are valid
- take_no_nulls::<T, I>(values.values(), indices.values())?
- }
- (true, false) => {
- // * nulls come from `values` alone
- // * all `indices.values()` are valid
- take_values_nulls::<T, I>(values, indices.values())?
- }
- (false, true) => {
- // in this branch it is unsound to read and use `index.values()`,
- // as doing so is UB when they come from a null slot.
- take_indices_nulls::<T, I>(values.values(), indices)?
- }
- (true, true) => {
- // in this branch it is unsound to read and use `index.values()`,
- // as doing so is UB when they come from a null slot.
- take_values_indices_nulls::<T, I>(values, indices)?
- }
- };
-
- let data = ArrayData::new(
- T::DATA_TYPE,
- indices.len(),
- None,
- nulls,
- 0,
- vec![buffer],
- vec![],
- );
- Ok(PrimitiveArray::<T>::from(data))
-}
-
-/// `take` implementation for boolean arrays
-fn take_boolean<IndexType>(
- values: &BooleanArray,
- indices: &PrimitiveArray<IndexType>,
-) -> Result<BooleanArray>
-where
- IndexType: ArrowNumericType,
- IndexType::Native: ToPrimitive,
-{
- let data_len = indices.len();
-
- let num_byte = bit_util::ceil(data_len, 8);
- let mut val_buf = MutableBuffer::from_len_zeroed(num_byte);
-
- let val_slice = val_buf.as_slice_mut();
-
- let null_count = values.null_count();
-
- let nulls;
- if null_count == 0 {
- (0..data_len).try_for_each::<_, Result<()>>(|i| {
- let index = ToPrimitive::to_usize(&indices.value(i)).ok_or_else(|| {
- ArrowError::ComputeError("Cast to usize failed".to_string())
- })?;
-
- if values.value(index) {
- bit_util::set_bit(val_slice, i);
- }
-
- Ok(())
- })?;
-
- nulls = indices.data_ref().null_buffer().cloned();
- } else {
- let mut null_buf = MutableBuffer::new(num_byte).with_bitset(num_byte, true);
- let null_slice = null_buf.as_slice_mut();
-
- (0..data_len).try_for_each::<_, Result<()>>(|i| {
- let index = ToPrimitive::to_usize(&indices.value(i)).ok_or_else(|| {
- ArrowError::ComputeError("Cast to usize failed".to_string())
- })?;
-
- if values.is_null(index) {
- bit_util::unset_bit(null_slice, i);
- } else if values.value(index) {
- bit_util::set_bit(val_slice, i);
- }
-
- Ok(())
- })?;
-
- nulls = match indices.data_ref().null_buffer() {
- Some(buffer) => Some(buffer_bin_and(
- buffer,
- 0,
- &null_buf.into(),
- 0,
- indices.len(),
- )),
- None => Some(null_buf.into()),
- };
- }
-
- let data = ArrayData::new(
- DataType::Boolean,
- indices.len(),
- None,
- nulls,
- 0,
- vec![val_buf.into()],
- vec![],
- );
- Ok(BooleanArray::from(data))
-}
-
-/// `take` implementation for string arrays
-fn take_string<OffsetSize, IndexType>(
- array: &GenericStringArray<OffsetSize>,
- indices: &PrimitiveArray<IndexType>,
-) -> Result<GenericStringArray<OffsetSize>>
-where
- OffsetSize: Zero + AddAssign + StringOffsetSizeTrait,
- IndexType: ArrowNumericType,
- IndexType::Native: ToPrimitive,
-{
- let data_len = indices.len();
-
- let bytes_offset = (data_len + 1) * std::mem::size_of::<OffsetSize>();
- let mut offsets_buffer = MutableBuffer::from_len_zeroed(bytes_offset);
-
- let offsets = offsets_buffer.typed_data_mut();
- let mut values = MutableBuffer::new(0);
- let mut length_so_far = OffsetSize::zero();
- offsets[0] = length_so_far;
-
- let nulls;
- if array.null_count() == 0 && indices.null_count() == 0 {
- for (i, offset) in offsets.iter_mut().skip(1).enumerate() {
- let index = ToPrimitive::to_usize(&indices.value(i)).ok_or_else(|| {
- ArrowError::ComputeError("Cast to usize failed".to_string())
- })?;
-
- let s = array.value(index);
-
- length_so_far += OffsetSize::from_usize(s.len()).unwrap();
- values.extend_from_slice(s.as_bytes());
- *offset = length_so_far;
- }
- nulls = None
- } else if indices.null_count() == 0 {
- let num_bytes = bit_util::ceil(data_len, 8);
-
- let mut null_buf = MutableBuffer::new(num_bytes).with_bitset(num_bytes, true);
- let null_slice = null_buf.as_slice_mut();
-
- for (i, offset) in offsets.iter_mut().skip(1).enumerate() {
- let index = ToPrimitive::to_usize(&indices.value(i)).ok_or_else(|| {
- ArrowError::ComputeError("Cast to usize failed".to_string())
- })?;
-
- if array.is_valid(index) {
- let s = array.value(index);
-
- length_so_far += OffsetSize::from_usize(s.len()).unwrap();
- values.extend_from_slice(s.as_bytes());
- } else {
- bit_util::unset_bit(null_slice, i);
- }
- *offset = length_so_far;
- }
- nulls = Some(null_buf.into());
- } else if array.null_count() == 0 {
- for (i, offset) in offsets.iter_mut().skip(1).enumerate() {
- if indices.is_valid(i) {
- let index =
- ToPrimitive::to_usize(&indices.value(i)).ok_or_else(|| {
- ArrowError::ComputeError("Cast to usize failed".to_string())
- })?;
-
- let s = array.value(index);
-
- length_so_far += OffsetSize::from_usize(s.len()).unwrap();
- values.extend_from_slice(s.as_bytes());
- }
- *offset = length_so_far;
- }
- nulls = indices.data_ref().null_buffer().cloned();
- } else {
- let num_bytes = bit_util::ceil(data_len, 8);
-
- let mut null_buf = MutableBuffer::new(num_bytes).with_bitset(num_bytes, true);
- let null_slice = null_buf.as_slice_mut();
-
- for (i, offset) in offsets.iter_mut().skip(1).enumerate() {
- let index = ToPrimitive::to_usize(&indices.value(i)).ok_or_else(|| {
- ArrowError::ComputeError("Cast to usize failed".to_string())
- })?;
-
- if array.is_valid(index) && indices.is_valid(i) {
- let s = array.value(index);
-
- length_so_far += OffsetSize::from_usize(s.len()).unwrap();
- values.extend_from_slice(s.as_bytes());
- } else {
- // set null bit
- bit_util::unset_bit(null_slice, i);
- }
- *offset = length_so_far;
- }
-
- nulls = match indices.data_ref().null_buffer() {
- Some(buffer) => {
- Some(buffer_bin_and(buffer, 0, &null_buf.into(), 0, data_len))
- }
- None => Some(null_buf.into()),
- };
- }
-
- let mut data = ArrayData::builder(<OffsetSize as StringOffsetSizeTrait>::DATA_TYPE)
- .len(data_len)
- .add_buffer(offsets_buffer.into())
- .add_buffer(values.into());
- if let Some(null_buffer) = nulls {
- data = data.null_bit_buffer(null_buffer);
- }
- Ok(GenericStringArray::<OffsetSize>::from(data.build()))
-}
-
-/// `take` implementation for list arrays
-///
-/// Calculates the index and indexed offset for the inner array,
-/// applying `take` on the inner array, then reconstructing a list array
-/// with the indexed offsets
-fn take_list<IndexType, OffsetType>(
- values: &GenericListArray<OffsetType::Native>,
- indices: &PrimitiveArray<IndexType>,
-) -> Result<GenericListArray<OffsetType::Native>>
-where
- IndexType: ArrowNumericType,
- IndexType::Native: ToPrimitive,
- OffsetType: ArrowNumericType,
- OffsetType::Native: ToPrimitive + OffsetSizeTrait,
- PrimitiveArray<OffsetType>: From<Vec<Option<OffsetType::Native>>>,
-{
- // TODO: Some optimizations can be done here such as if it is
- // taking the whole list or a contiguous sublist
- let (list_indices, offsets) =
- take_value_indices_from_list::<IndexType, OffsetType>(values, indices)?;
-
- let taken = take_impl::<OffsetType>(values.values().as_ref(), &list_indices, None)?;
- // determine null count and null buffer, which are a function of `values` and `indices`
- let mut null_count = 0;
- let num_bytes = bit_util::ceil(indices.len(), 8);
- let mut null_buf = MutableBuffer::new(num_bytes).with_bitset(num_bytes, true);
- {
- let null_slice = null_buf.as_slice_mut();
- offsets[..].windows(2).enumerate().for_each(
- |(i, window): (usize, &[OffsetType::Native])| {
- if window[0] == window[1] {
- // offsets are equal, slot is null
- bit_util::unset_bit(null_slice, i);
- null_count += 1;
- }
- },
- );
- }
- let value_offsets = Buffer::from_slice_ref(&offsets);
- // create a new list with taken data and computed null information
- let list_data = ArrayDataBuilder::new(values.data_type().clone())
- .len(indices.len())
- .null_bit_buffer(null_buf.into())
- .offset(0)
- .add_child_data(taken.data().clone())
- .add_buffer(value_offsets)
- .build();
- Ok(GenericListArray::<OffsetType::Native>::from(list_data))
-}
-
-/// `take` implementation for `FixedSizeListArray`
-///
-/// Calculates the index and indexed offset for the inner array,
-/// applying `take` on the inner array, then reconstructing a list array
-/// with the indexed offsets
-fn take_fixed_size_list<IndexType>(
- values: &FixedSizeListArray,
- indices: &PrimitiveArray<IndexType>,
- length: <UInt32Type as ArrowPrimitiveType>::Native,
-) -> Result<FixedSizeListArray>
-where
- IndexType: ArrowNumericType,
- IndexType::Native: ToPrimitive,
-{
- let list_indices = take_value_indices_from_fixed_size_list(values, indices, length)?;
- let taken = take_impl::<UInt32Type>(values.values().as_ref(), &list_indices, None)?;
-
- // determine null count and null buffer, which are a function of `values` and `indices`
- let num_bytes = bit_util::ceil(indices.len(), 8);
- let mut null_buf = MutableBuffer::new(num_bytes).with_bitset(num_bytes, true);
- let null_slice = null_buf.as_slice_mut();
-
- for i in 0..indices.len() {
- let index = ToPrimitive::to_usize(&indices.value(i)).ok_or_else(|| {
- ArrowError::ComputeError("Cast to usize failed".to_string())
- })?;
- if !indices.is_valid(i) || values.is_null(index) {
- bit_util::unset_bit(null_slice, i);
- }
- }
-
- let list_data = ArrayDataBuilder::new(values.data_type().clone())
- .len(indices.len())
- .null_bit_buffer(null_buf.into())
- .offset(0)
- .add_child_data(taken.data().clone())
- .build();
-
- Ok(FixedSizeListArray::from(list_data))
-}
-
-/// `take` implementation for dictionary arrays
-///
-/// applies `take` to the keys of the dictionary array and returns a new dictionary array
-/// with the same dictionary values and reordered keys
-fn take_dict<T, I>(
- values: &DictionaryArray<T>,
- indices: &PrimitiveArray<I>,
-) -> Result<DictionaryArray<T>>
-where
- T: ArrowPrimitiveType,
- T::Native: num::Num,
- I: ArrowNumericType,
- I::Native: ToPrimitive,
-{
- let new_keys = take_primitive::<T, I>(values.keys(), indices)?;
- let new_keys_data = new_keys.data_ref();
-
- let data = ArrayData::new(
- values.data_type().clone(),
- new_keys.len(),
- Some(new_keys_data.null_count()),
- new_keys_data.null_buffer().cloned(),
- 0,
- new_keys_data.buffers().to_vec(),
- values.data().child_data().to_vec(),
- );
-
- Ok(DictionaryArray::<T>::from(data))
-}
-
-#[cfg(test)]
-mod tests {
- use super::*;
- use crate::compute::util::tests::build_fixed_size_list_nullable;
-
- fn test_take_boolean_arrays(
- data: Vec<Option<bool>>,
- index: &UInt32Array,
- options: Option<TakeOptions>,
- expected_data: Vec<Option<bool>>,
- ) {
- let output = BooleanArray::from(data);
- let expected = Arc::new(BooleanArray::from(expected_data)) as ArrayRef;
- let output = take(&output, index, options).unwrap();
- assert_eq!(&output, &expected)
- }
-
- fn test_take_primitive_arrays<T>(
- data: Vec<Option<T::Native>>,
- index: &UInt32Array,
- options: Option<TakeOptions>,
- expected_data: Vec<Option<T::Native>>,
- ) -> Result<()>
- where
- T: ArrowPrimitiveType,
- PrimitiveArray<T>: From<Vec<Option<T::Native>>>,
- {
- let output = PrimitiveArray::<T>::from(data);
- let expected = Arc::new(PrimitiveArray::<T>::from(expected_data)) as ArrayRef;
- let output = take(&output, index, options)?;
- assert_eq!(&output, &expected);
- Ok(())
- }
-
- fn test_take_impl_primitive_arrays<T, I>(
- data: Vec<Option<T::Native>>,
- index: &PrimitiveArray<I>,
- options: Option<TakeOptions>,
- expected_data: Vec<Option<T::Native>>,
- ) where
- T: ArrowPrimitiveType,
- PrimitiveArray<T>: From<Vec<Option<T::Native>>>,
- I: ArrowNumericType,
- I::Native: ToPrimitive,
- {
- let output = PrimitiveArray::<T>::from(data);
- let expected = PrimitiveArray::<T>::from(expected_data);
- let output = take_impl(&output, index, options).unwrap();
- let output = output.as_any().downcast_ref::<PrimitiveArray<T>>().unwrap();
- assert_eq!(output, &expected)
- }
-
- // create a simple struct for testing purposes
- fn create_test_struct() -> StructArray {
- let boolean_data = BooleanArray::from(vec![true, false, false, true])
- .data()
- .clone();
- let int_data = Int32Array::from(vec![42, 28, 19, 31]).data().clone();
- let mut field_types = vec![];
- field_types.push(Field::new("a", DataType::Boolean, true));
- field_types.push(Field::new("b", DataType::Int32, true));
- let struct_array_data = ArrayData::builder(DataType::Struct(field_types))
- .len(4)
- .add_child_data(boolean_data)
- .add_child_data(int_data)
- .build();
- StructArray::from(struct_array_data)
- }
-
- #[test]
- fn test_take_primitive_non_null_indices() {
- let index = UInt32Array::from(vec![0, 5, 3, 1, 4, 2]);
- test_take_primitive_arrays::<Int8Type>(
- vec![None, Some(3), Some(5), Some(2), Some(3), None],
- &index,
- None,
- vec![None, None, Some(2), Some(3), Some(3), Some(5)],
- )
- .unwrap();
- }
-
- #[test]
- fn test_take_primitive_non_null_values() {
- let index = UInt32Array::from(vec![Some(3), None, Some(1), Some(3), Some(2)]);
- test_take_primitive_arrays::<Int8Type>(
- vec![Some(0), Some(1), Some(2), Some(3), Some(4)],
- &index,
- None,
- vec![Some(3), None, Some(1), Some(3), Some(2)],
- )
- .unwrap();
- }
-
- #[test]
- fn test_take_primitive_non_null() {
- let index = UInt32Array::from(vec![0, 5, 3, 1, 4, 2]);
- test_take_primitive_arrays::<Int8Type>(
- vec![Some(0), Some(3), Some(5), Some(2), Some(3), Some(1)],
- &index,
- None,
- vec![Some(0), Some(1), Some(2), Some(3), Some(3), Some(5)],
- )
- .unwrap();
- }
-
- #[test]
- fn test_take_primitive() {
- let index = UInt32Array::from(vec![Some(3), None, Some(1), Some(3), Some(2)]);
-
- // int8
- test_take_primitive_arrays::<Int8Type>(
- vec![Some(0), None, Some(2), Some(3), None],
- &index,
- None,
- vec![Some(3), None, None, Some(3), Some(2)],
- )
- .unwrap();
-
- // int16
- test_take_primitive_arrays::<Int16Type>(
- vec![Some(0), None, Some(2), Some(3), None],
- &index,
- None,
- vec![Some(3), None, None, Some(3), Some(2)],
- )
- .unwrap();
-
- // int32
- test_take_primitive_arrays::<Int32Type>(
- vec![Some(0), None, Some(2), Some(3), None],
- &index,
- None,
- vec![Some(3), None, None, Some(3), Some(2)],
- )
- .unwrap();
-
- // int64
- test_take_primitive_arrays::<Int64Type>(
- vec![Some(0), None, Some(2), Some(3), None],
- &index,
- None,
- vec![Some(3), None, None, Some(3), Some(2)],
- )
- .unwrap();
-
- // uint8
- test_take_primitive_arrays::<UInt8Type>(
- vec![Some(0), None, Some(2), Some(3), None],
- &index,
- None,
- vec![Some(3), None, None, Some(3), Some(2)],
- )
- .unwrap();
-
- // uint16
- test_take_primitive_arrays::<UInt16Type>(
- vec![Some(0), None, Some(2), Some(3), None],
- &index,
- None,
- vec![Some(3), None, None, Some(3), Some(2)],
- )
- .unwrap();
-
- // uint32
- test_take_primitive_arrays::<UInt32Type>(
- vec![Some(0), None, Some(2), Some(3), None],
- &index,
- None,
- vec![Some(3), None, None, Some(3), Some(2)],
- )
- .unwrap();
-
- // int64
- test_take_primitive_arrays::<Int64Type>(
- vec![Some(0), None, Some(2), Some(-15), None],
- &index,
- None,
- vec![Some(-15), None, None, Some(-15), Some(2)],
- )
- .unwrap();
-
- // interval_year_month
- test_take_primitive_arrays::<IntervalYearMonthType>(
- vec![Some(0), None, Some(2), Some(-15), None],
- &index,
- None,
- vec![Some(-15), None, None, Some(-15), Some(2)],
- )
- .unwrap();
-
- // interval_day_time
- test_take_primitive_arrays::<IntervalDayTimeType>(
- vec![Some(0), None, Some(2), Some(-15), None],
- &index,
- None,
- vec![Some(-15), None, None, Some(-15), Some(2)],
- )
- .unwrap();
-
- // duration_second
- test_take_primitive_arrays::<DurationSecondType>(
- vec![Some(0), None, Some(2), Some(-15), None],
- &index,
- None,
- vec![Some(-15), None, None, Some(-15), Some(2)],
- )
- .unwrap();
-
- // duration_millisecond
- test_take_primitive_arrays::<DurationMillisecondType>(
- vec![Some(0), None, Some(2), Some(-15), None],
- &index,
- None,
- vec![Some(-15), None, None, Some(-15), Some(2)],
- )
- .unwrap();
-
- // duration_microsecond
- test_take_primitive_arrays::<DurationMicrosecondType>(
- vec![Some(0), None, Some(2), Some(-15), None],
- &index,
- None,
- vec![Some(-15), None, None, Some(-15), Some(2)],
- )
- .unwrap();
-
- // duration_nanosecond
- test_take_primitive_arrays::<DurationNanosecondType>(
- vec![Some(0), None, Some(2), Some(-15), None],
- &index,
- None,
- vec![Some(-15), None, None, Some(-15), Some(2)],
- )
- .unwrap();
-
- // float32
- test_take_primitive_arrays::<Float32Type>(
- vec![Some(0.0), None, Some(2.21), Some(-3.1), None],
- &index,
- None,
- vec![Some(-3.1), None, None, Some(-3.1), Some(2.21)],
- )
- .unwrap();
-
- // float64
- test_take_primitive_arrays::<Float64Type>(
- vec![Some(0.0), None, Some(2.21), Some(-3.1), None],
- &index,
- None,
- vec![Some(-3.1), None, None, Some(-3.1), Some(2.21)],
- )
- .unwrap();
- }
-
- #[test]
- fn test_take_impl_primitive_with_int64_indices() {
- let index = Int64Array::from(vec![Some(3), None, Some(1), Some(3), Some(2)]);
-
- // int16
- test_take_impl_primitive_arrays::<Int16Type, Int64Type>(
- vec![Some(0), None, Some(2), Some(3), None],
- &index,
- None,
- vec![Some(3), None, None, Some(3), Some(2)],
- );
-
- // int64
- test_take_impl_primitive_arrays::<Int64Type, Int64Type>(
- vec![Some(0), None, Some(2), Some(-15), None],
- &index,
- None,
- vec![Some(-15), None, None, Some(-15), Some(2)],
- );
-
- // uint64
- test_take_impl_primitive_arrays::<UInt64Type, Int64Type>(
- vec![Some(0), None, Some(2), Some(3), None],
- &index,
- None,
- vec![Some(3), None, None, Some(3), Some(2)],
- );
-
- // duration_millisecond
- test_take_impl_primitive_arrays::<DurationMillisecondType, Int64Type>(
- vec![Some(0), None, Some(2), Some(-15), None],
- &index,
- None,
- vec![Some(-15), None, None, Some(-15), Some(2)],
- );
-
- // float32
- test_take_impl_primitive_arrays::<Float32Type, Int64Type>(
- vec![Some(0.0), None, Some(2.21), Some(-3.1), None],
- &index,
- None,
- vec![Some(-3.1), None, None, Some(-3.1), Some(2.21)],
- );
- }
-
- #[test]
- fn test_take_impl_primitive_with_uint8_indices() {
- let index = UInt8Array::from(vec![Some(3), None, Some(1), Some(3), Some(2)]);
-
- // int16
- test_take_impl_primitive_arrays::<Int16Type, UInt8Type>(
- vec![Some(0), None, Some(2), Some(3), None],
- &index,
- None,
- vec![Some(3), None, None, Some(3), Some(2)],
- );
-
- // duration_millisecond
- test_take_impl_primitive_arrays::<DurationMillisecondType, UInt8Type>(
- vec![Some(0), None, Some(2), Some(-15), None],
- &index,
- None,
- vec![Some(-15), None, None, Some(-15), Some(2)],
- );
-
- // float32
- test_take_impl_primitive_arrays::<Float32Type, UInt8Type>(
- vec![Some(0.0), None, Some(2.21), Some(-3.1), None],
- &index,
- None,
- vec![Some(-3.1), None, None, Some(-3.1), Some(2.21)],
- );
- }
-
- #[test]
- fn test_take_primitive_bool() {
- let index = UInt32Array::from(vec![Some(3), None, Some(1), Some(3), Some(2)]);
- // boolean
- test_take_boolean_arrays(
- vec![Some(false), None, Some(true), Some(false), None],
- &index,
- None,
- vec![Some(false), None, None, Some(false), Some(true)],
- );
- }
-
- fn _test_take_string<'a, K: 'static>()
- where
- K: Array + PartialEq + From<Vec<Option<&'a str>>>,
- {
- let index = UInt32Array::from(vec![Some(3), None, Some(1), Some(3), Some(4)]);
-
- let array = K::from(vec![
- Some("one"),
- None,
- Some("three"),
- Some("four"),
- Some("five"),
- ]);
- let actual = take(&array, &index, None).unwrap();
- assert_eq!(actual.len(), index.len());
-
- let actual = actual.as_any().downcast_ref::<K>().unwrap();
-
- let expected =
- K::from(vec![Some("four"), None, None, Some("four"), Some("five")]);
-
- assert_eq!(actual, &expected);
- }
-
- #[test]
- fn test_take_string() {
- _test_take_string::<StringArray>()
- }
-
- #[test]
- fn test_take_large_string() {
- _test_take_string::<LargeStringArray>()
- }
-
- macro_rules! test_take_list {
- ($offset_type:ty, $list_data_type:ident, $list_array_type:ident) => {{
- // Construct a value array, [[0,0,0], [-1,-2,-1], [2,3]]
- let value_data = Int32Array::from(vec![0, 0, 0, -1, -2, -1, 2, 3])
- .data()
- .clone();
- // Construct offsets
- let value_offsets: [$offset_type; 4] = [0, 3, 6, 8];
- let value_offsets = Buffer::from_slice_ref(&value_offsets);
- // Construct a list array from the above two
- let list_data_type = DataType::$list_data_type(Box::new(Field::new(
- "item",
- DataType::Int32,
- false,
- )));
- let list_data = ArrayData::builder(list_data_type.clone())
- .len(3)
- .add_buffer(value_offsets)
- .add_child_data(value_data)
- .build();
- let list_array = $list_array_type::from(list_data);
-
- // index returns: [[2,3], null, [-1,-2,-1], [2,3], [0,0,0]]
- let index = UInt32Array::from(vec![Some(2), None, Some(1), Some(2), Some(0)]);
-
- let a = take(&list_array, &index, None).unwrap();
- let a: &$list_array_type =
- a.as_any().downcast_ref::<$list_array_type>().unwrap();
-
- // construct a value array with expected results:
- // [[2,3], null, [-1,-2,-1], [2,3], [0,0,0]]
- let expected_data = Int32Array::from(vec![
- Some(2),
- Some(3),
- Some(-1),
- Some(-2),
- Some(-1),
- Some(2),
- Some(3),
- Some(0),
- Some(0),
- Some(0),
- ])
- .data()
- .clone();
- // construct offsets
- let expected_offsets: [$offset_type; 6] = [0, 2, 2, 5, 7, 10];
- let expected_offsets = Buffer::from_slice_ref(&expected_offsets);
- // construct list array from the two
- let expected_list_data = ArrayData::builder(list_data_type)
- .len(5)
- // null buffer remains the same as only the indices have nulls
- .null_bit_buffer(
- index.data().null_bitmap().as_ref().unwrap().bits.clone(),
- )
- .add_buffer(expected_offsets)
- .add_child_data(expected_data)
- .build();
- let expected_list_array = $list_array_type::from(expected_list_data);
-
- assert_eq!(a, &expected_list_array);
- }};
- }
-
- macro_rules! test_take_list_with_value_nulls {
- ($offset_type:ty, $list_data_type:ident, $list_array_type:ident) => {{
- // Construct a value array, [[0,null,0], [-1,-2,3], [null], [5,null]]
- let value_data = Int32Array::from(vec![
- Some(0),
- None,
- Some(0),
- Some(-1),
- Some(-2),
- Some(3),
- None,
- Some(5),
- None,
- ])
- .data()
- .clone();
- // Construct offsets
- let value_offsets: [$offset_type; 5] = [0, 3, 6, 7, 9];
- let value_offsets = Buffer::from_slice_ref(&value_offsets);
- // Construct a list array from the above two
- let list_data_type = DataType::$list_data_type(Box::new(Field::new(
- "item",
- DataType::Int32,
- false,
- )));
- let list_data = ArrayData::builder(list_data_type.clone())
- .len(4)
- .add_buffer(value_offsets)
- .null_bit_buffer(Buffer::from([0b10111101, 0b00000000]))
- .add_child_data(value_data)
- .build();
- let list_array = $list_array_type::from(list_data);
-
- // index returns: [[null], null, [-1,-2,3], [2,null], [0,null,0]]
- let index = UInt32Array::from(vec![Some(2), None, Some(1), Some(3), Some(0)]);
-
- let a = take(&list_array, &index, None).unwrap();
- let a: &$list_array_type =
- a.as_any().downcast_ref::<$list_array_type>().unwrap();
-
- // construct a value array with expected results:
- // [[null], null, [-1,-2,3], [5,null], [0,null,0]]
- let expected_data = Int32Array::from(vec![
- None,
- Some(-1),
- Some(-2),
- Some(3),
- Some(5),
- None,
- Some(0),
- None,
- Some(0),
- ])
- .data()
- .clone();
- // construct offsets
- let expected_offsets: [$offset_type; 6] = [0, 1, 1, 4, 6, 9];
- let expected_offsets = Buffer::from_slice_ref(&expected_offsets);
- // construct list array from the two
- let expected_list_data = ArrayData::builder(list_data_type)
- .len(5)
- // null buffer remains the same as only the indices have nulls
- .null_bit_buffer(
- index.data().null_bitmap().as_ref().unwrap().bits.clone(),
- )
- .add_buffer(expected_offsets)
- .add_child_data(expected_data)
- .build();
- let expected_list_array = $list_array_type::from(expected_list_data);
-
- assert_eq!(a, &expected_list_array);
- }};
- }
-
- macro_rules! test_take_list_with_nulls {
- ($offset_type:ty, $list_data_type:ident, $list_array_type:ident) => {{
- // Construct a value array, [[0,null,0], [-1,-2,3], null, [5,null]]
- let value_data = Int32Array::from(vec![
- Some(0),
- None,
- Some(0),
- Some(-1),
- Some(-2),
- Some(3),
- Some(5),
- None,
- ])
- .data()
- .clone();
- // Construct offsets
- let value_offsets: [$offset_type; 5] = [0, 3, 6, 6, 8];
- let value_offsets = Buffer::from_slice_ref(&value_offsets);
- // Construct a list array from the above two
- let list_data_type = DataType::$list_data_type(Box::new(Field::new(
- "item",
- DataType::Int32,
- false,
- )));
- let list_data = ArrayData::builder(list_data_type.clone())
- .len(4)
- .add_buffer(value_offsets)
- .null_bit_buffer(Buffer::from([0b01111101]))
- .add_child_data(value_data)
- .build();
- let list_array = $list_array_type::from(list_data);
-
- // index returns: [null, null, [-1,-2,3], [5,null], [0,null,0]]
- let index = UInt32Array::from(vec![Some(2), None, Some(1), Some(3), Some(0)]);
-
- let a = take(&list_array, &index, None).unwrap();
- let a: &$list_array_type =
- a.as_any().downcast_ref::<$list_array_type>().unwrap();
-
- // construct a value array with expected results:
- // [null, null, [-1,-2,3], [5,null], [0,null,0]]
- let expected_data = Int32Array::from(vec![
- Some(-1),
- Some(-2),
- Some(3),
- Some(5),
- None,
- Some(0),
- None,
- Some(0),
- ])
- .data()
- .clone();
- // construct offsets
- let expected_offsets: [$offset_type; 6] = [0, 0, 0, 3, 5, 8];
- let expected_offsets = Buffer::from_slice_ref(&expected_offsets);
- // construct list array from the two
- let mut null_bits: [u8; 1] = [0; 1];
- bit_util::set_bit(&mut null_bits, 2);
- bit_util::set_bit(&mut null_bits, 3);
- bit_util::set_bit(&mut null_bits, 4);
- let expected_list_data = ArrayData::builder(list_data_type)
- .len(5)
- // null buffer must be recalculated as both values and indices have nulls
- .null_bit_buffer(Buffer::from(null_bits))
- .add_buffer(expected_offsets)
- .add_child_data(expected_data)
- .build();
- let expected_list_array = $list_array_type::from(expected_list_data);
-
- assert_eq!(a, &expected_list_array);
- }};
- }
-
- fn do_take_fixed_size_list_test<T>(
- length: <Int32Type as ArrowPrimitiveType>::Native,
- input_data: Vec<Option<Vec<Option<T::Native>>>>,
- indices: Vec<<UInt32Type as ArrowPrimitiveType>::Native>,
- expected_data: Vec<Option<Vec<Option<T::Native>>>>,
- ) where
- T: ArrowPrimitiveType,
- PrimitiveArray<T>: From<Vec<Option<T::Native>>>,
- {
- let indices = UInt32Array::from(indices);
-
- let input_array = build_fixed_size_list_nullable::<T>(input_data, length);
-
- let output = take_fixed_size_list(&input_array, &indices, length as u32).unwrap();
-
- let expected = build_fixed_size_list_nullable::<T>(expected_data, length);
-
- assert_eq!(&output, &expected)
- }
-
- #[test]
- fn test_take_list() {
- test_take_list!(i32, List, ListArray);
- }
-
- #[test]
- fn test_take_large_list() {
- test_take_list!(i64, LargeList, LargeListArray);
- }
-
- #[test]
- fn test_take_list_with_value_nulls() {
- test_take_list_with_value_nulls!(i32, List, ListArray);
- }
-
- #[test]
- fn test_take_large_list_with_value_nulls() {
- test_take_list_with_value_nulls!(i64, LargeList, LargeListArray);
- }
-
- #[test]
- fn test_test_take_list_with_nulls() {
- test_take_list_with_nulls!(i32, List, ListArray);
- }
-
- #[test]
- fn test_test_take_large_list_with_nulls() {
- test_take_list_with_nulls!(i64, LargeList, LargeListArray);
- }
-
- #[test]
- fn test_take_fixed_size_list() {
- do_take_fixed_size_list_test::<Int32Type>(
- 3,
- vec![
- Some(vec![None, Some(1), Some(2)]),
- Some(vec![Some(3), Some(4), None]),
- Some(vec![Some(6), Some(7), Some(8)]),
- ],
- vec![2, 1, 0],
- vec![
- Some(vec![Some(6), Some(7), Some(8)]),
- Some(vec![Some(3), Some(4), None]),
- Some(vec![None, Some(1), Some(2)]),
- ],
- );
-
- do_take_fixed_size_list_test::<UInt8Type>(
- 1,
- vec![
- Some(vec![Some(1)]),
- Some(vec![Some(2)]),
- Some(vec![Some(3)]),
- Some(vec![Some(4)]),
- Some(vec![Some(5)]),
- Some(vec![Some(6)]),
- Some(vec![Some(7)]),
- Some(vec![Some(8)]),
- ],
- vec![2, 7, 0],
- vec![
- Some(vec![Some(3)]),
- Some(vec![Some(8)]),
- Some(vec![Some(1)]),
- ],
- );
-
- do_take_fixed_size_list_test::<UInt64Type>(
- 3,
- vec![
- Some(vec![Some(10), Some(11), Some(12)]),
- Some(vec![Some(13), Some(14), Some(15)]),
- None,
- Some(vec![Some(16), Some(17), Some(18)]),
- ],
- vec![3, 2, 1, 2, 0],
- vec![
- Some(vec![Some(16), Some(17), Some(18)]),
- None,
- Some(vec![Some(13), Some(14), Some(15)]),
- None,
- Some(vec![Some(10), Some(11), Some(12)]),
- ],
- );
- }
-
- #[test]
- #[should_panic(expected = "index out of bounds: the len is 4 but the index is 1000")]
- fn test_take_list_out_of_bounds() {
- // Construct a value array, [[0,0,0], [-1,-2,-1], [2,3]]
- let value_data = Int32Array::from(vec![0, 0, 0, -1, -2, -1, 2, 3])
- .data()
- .clone();
- // Construct offsets
- let value_offsets = Buffer::from_slice_ref(&[0, 3, 6, 8]);
- // Construct a list array from the above two
- let list_data_type =
- DataType::List(Box::new(Field::new("item", DataType::Int32, false)));
- let list_data = ArrayData::builder(list_data_type)
- .len(3)
- .add_buffer(value_offsets)
- .add_child_data(value_data)
- .build();
- let list_array = ListArray::from(list_data);
-
- let index = UInt32Array::from(vec![1000]);
-
- // A panic is expected here since we have not supplied the check_bounds
- // option.
- take(&list_array, &index, None).unwrap();
- }
-
- #[test]
- fn test_take_struct() {
- let array = create_test_struct();
-
- let index = UInt32Array::from(vec![0, 3, 1, 0, 2]);
- let a = take(&array, &index, None).unwrap();
- let a: &StructArray = a.as_any().downcast_ref::<StructArray>().unwrap();
- assert_eq!(index.len(), a.len());
- assert_eq!(0, a.null_count());
-
- let expected_bool_data = BooleanArray::from(vec![true, true, false, true, false])
- .data()
- .clone();
- let expected_int_data = Int32Array::from(vec![42, 31, 28, 42, 19]).data().clone();
- let mut field_types = vec![];
- field_types.push(Field::new("a", DataType::Boolean, true));
- field_types.push(Field::new("b", DataType::Int32, true));
- let struct_array_data = ArrayData::builder(DataType::Struct(field_types))
- .len(5)
- .add_child_data(expected_bool_data)
- .add_child_data(expected_int_data)
- .build();
- let struct_array = StructArray::from(struct_array_data);
-
- assert_eq!(a, &struct_array);
- }
-
- #[test]
- fn test_take_struct_with_nulls() {
- let array = create_test_struct();
-
- let index = UInt32Array::from(vec![None, Some(3), Some(1), None, Some(0)]);
- let a = take(&array, &index, None).unwrap();
- let a: &StructArray = a.as_any().downcast_ref::<StructArray>().unwrap();
- assert_eq!(index.len(), a.len());
- assert_eq!(0, a.null_count());
-
- let expected_bool_data =
- BooleanArray::from(vec![None, Some(true), Some(false), None, Some(true)])
- .data()
- .clone();
- let expected_int_data =
- Int32Array::from(vec![None, Some(31), Some(28), None, Some(42)])
- .data()
- .clone();
-
- let mut field_types = vec![];
- field_types.push(Field::new("a", DataType::Boolean, true));
- field_types.push(Field::new("b", DataType::Int32, true));
- let struct_array_data = ArrayData::builder(DataType::Struct(field_types))
- .len(5)
- // TODO: see https://issues.apache.org/jira/browse/ARROW-5408 for why count != 2
- .add_child_data(expected_bool_data)
- .add_child_data(expected_int_data)
- .build();
- let struct_array = StructArray::from(struct_array_data);
- assert_eq!(a, &struct_array);
- }
-
- #[test]
- fn test_take_out_of_bounds() {
- let index = UInt32Array::from(vec![Some(3), None, Some(1), Some(3), Some(6)]);
- let take_opt = TakeOptions { check_bounds: true };
-
- // int64
- let result = test_take_primitive_arrays::<Int64Type>(
- vec![Some(0), None, Some(2), Some(3), None],
- &index,
- Some(take_opt),
- vec![None],
- );
- assert!(result.is_err());
- }
-
- #[test]
- #[should_panic(expected = "index out of bounds: the len is 4 but the index is 1000")]
- fn test_take_out_of_bounds_panic() {
- let index = UInt32Array::from(vec![Some(1000)]);
-
- test_take_primitive_arrays::<Int64Type>(
- vec![Some(0), Some(1), Some(2), Some(3)],
- &index,
- None,
- vec![None],
- )
- .unwrap();
- }
-
- #[test]
- fn test_take_dict() {
- let keys_builder = Int16Builder::new(8);
- let values_builder = StringBuilder::new(4);
-
- let mut dict_builder = StringDictionaryBuilder::new(keys_builder, values_builder);
-
- dict_builder.append("foo").unwrap();
- dict_builder.append("bar").unwrap();
- dict_builder.append("").unwrap();
- dict_builder.append_null().unwrap();
- dict_builder.append("foo").unwrap();
- dict_builder.append("bar").unwrap();
- dict_builder.append("bar").unwrap();
- dict_builder.append("foo").unwrap();
-
- let array = dict_builder.finish();
- let dict_values = array.values().clone();
- let dict_values = dict_values.as_any().downcast_ref::<StringArray>().unwrap();
-
- let indices = UInt32Array::from(vec![
- Some(0), // first "foo"
- Some(7), // last "foo"
- None, // null index should return null
- Some(5), // second "bar"
- Some(6), // another "bar"
- Some(2), // empty string
- Some(3), // input is null at this index
- ]);
-
- let result = take(&array, &indices, None).unwrap();
- let result = result
- .as_any()
- .downcast_ref::<DictionaryArray<Int16Type>>()
- .unwrap();
-
- let result_values: StringArray = result.values().data().clone().into();
-
- // dictionary values should stay the same
- let expected_values = StringArray::from(vec!["foo", "bar", ""]);
- assert_eq!(&expected_values, dict_values);
- assert_eq!(&expected_values, &result_values);
-
- let expected_keys = Int16Array::from(vec![
- Some(0),
- Some(0),
- None,
- Some(1),
- Some(1),
- Some(2),
- None,
- ]);
- assert_eq!(result.keys(), &expected_keys);
- }
-}
diff --git a/arrow/src/compute/kernels/temporal.rs b/arrow/src/compute/kernels/temporal.rs
deleted file mode 100644
index 63e4129..0000000
--- a/arrow/src/compute/kernels/temporal.rs
+++ /dev/null
@@ -1,187 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-//! Defines temporal kernels for time and date related functions.
-
-use chrono::{Datelike, Timelike};
-
-use crate::array::*;
-use crate::datatypes::*;
-use crate::error::{ArrowError, Result};
-/// Extracts the hours of a given temporal array as an array of integers
-pub fn hour<T>(array: &PrimitiveArray<T>) -> Result<Int32Array>
-where
- T: ArrowTemporalType + ArrowNumericType,
- i64: std::convert::From<T::Native>,
-{
- let mut b = Int32Builder::new(array.len());
- match array.data_type() {
- &DataType::Time32(_) | &DataType::Time64(_) => {
- for i in 0..array.len() {
- if array.is_null(i) {
- b.append_null()?;
- } else {
- match array.value_as_time(i) {
- Some(time) => b.append_value(time.hour() as i32)?,
- None => b.append_null()?,
- };
- }
- }
- }
- &DataType::Date32 | &DataType::Date64 | &DataType::Timestamp(_, _) => {
- for i in 0..array.len() {
- if array.is_null(i) {
- b.append_null()?;
- } else {
- match array.value_as_datetime(i) {
- Some(dt) => b.append_value(dt.hour() as i32)?,
- None => b.append_null()?,
- }
- }
- }
- }
- dt => {
- return {
- Err(ArrowError::ComputeError(format!(
- "hour does not support type {:?}",
- dt
- )))
- }
- }
- }
-
- Ok(b.finish())
-}
-
-/// Extracts the years of a given temporal array as an array of integers
-pub fn year<T>(array: &PrimitiveArray<T>) -> Result<Int32Array>
-where
- T: ArrowTemporalType + ArrowNumericType,
- i64: std::convert::From<T::Native>,
-{
- let mut b = Int32Builder::new(array.len());
- match array.data_type() {
- &DataType::Date32 | &DataType::Date64 | &DataType::Timestamp(_, _) => {
- for i in 0..array.len() {
- if array.is_null(i) {
- b.append_null()?;
- } else {
- match array.value_as_datetime(i) {
- Some(dt) => b.append_value(dt.year() as i32)?,
- None => b.append_null()?,
- }
- }
- }
- }
- dt => {
- return {
- Err(ArrowError::ComputeError(format!(
- "year does not support type {:?}",
- dt
- )))
- }
- }
- }
-
- Ok(b.finish())
-}
-
-#[cfg(test)]
-mod tests {
- use super::*;
-
- #[test]
- fn test_temporal_array_date64_hour() {
- let a: PrimitiveArray<Date64Type> =
- vec![Some(1514764800000), None, Some(1550636625000)].into();
-
- let b = hour(&a).unwrap();
- assert_eq!(0, b.value(0));
- assert_eq!(false, b.is_valid(1));
- assert_eq!(4, b.value(2));
- }
-
- #[test]
- fn test_temporal_array_date32_hour() {
- let a: PrimitiveArray<Date32Type> = vec![Some(15147), None, Some(15148)].into();
-
- let b = hour(&a).unwrap();
- assert_eq!(0, b.value(0));
- assert_eq!(false, b.is_valid(1));
- assert_eq!(0, b.value(2));
- }
-
- #[test]
- fn test_temporal_array_time32_second_hour() {
- let a: PrimitiveArray<Time32SecondType> = vec![37800, 86339].into();
-
- let b = hour(&a).unwrap();
- assert_eq!(10, b.value(0));
- assert_eq!(23, b.value(1));
- }
-
- #[test]
- fn test_temporal_array_time64_micro_hour() {
- let a: PrimitiveArray<Time64MicrosecondType> =
- vec![37800000000, 86339000000].into();
-
- let b = hour(&a).unwrap();
- assert_eq!(10, b.value(0));
- assert_eq!(23, b.value(1));
- }
-
- #[test]
- fn test_temporal_array_timestamp_micro_hour() {
- let a: TimestampMicrosecondArray = vec![37800000000, 86339000000].into();
-
- let b = hour(&a).unwrap();
- assert_eq!(10, b.value(0));
- assert_eq!(23, b.value(1));
- }
-
- #[test]
- fn test_temporal_array_date64_year() {
- let a: PrimitiveArray<Date64Type> =
- vec![Some(1514764800000), None, Some(1550636625000)].into();
-
- let b = year(&a).unwrap();
- assert_eq!(2018, b.value(0));
- assert_eq!(false, b.is_valid(1));
- assert_eq!(2019, b.value(2));
- }
-
- #[test]
- fn test_temporal_array_date32_year() {
- let a: PrimitiveArray<Date32Type> = vec![Some(15147), None, Some(15448)].into();
-
- let b = year(&a).unwrap();
- assert_eq!(2011, b.value(0));
- assert_eq!(false, b.is_valid(1));
- assert_eq!(2012, b.value(2));
- }
-
- #[test]
- fn test_temporal_array_timestamp_micro_year() {
- let a: TimestampMicrosecondArray =
- vec![Some(1612025847000000), None, Some(1722015847000000)].into();
-
- let b = year(&a).unwrap();
- assert_eq!(2021, b.value(0));
- assert_eq!(false, b.is_valid(1));
- assert_eq!(2024, b.value(2));
- }
-}
diff --git a/arrow/src/compute/kernels/window.rs b/arrow/src/compute/kernels/window.rs
deleted file mode 100644
index a537cbc..0000000
--- a/arrow/src/compute/kernels/window.rs
+++ /dev/null
@@ -1,195 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-//! Defines windowing functions, like `shift`ing
-
-use crate::array::{Array, ArrayRef};
-use crate::error::Result;
-use crate::{
- array::{make_array, new_null_array},
- compute::concat,
-};
-use num::{abs, clamp};
-
-/// Shifts array by defined number of items (to left or right)
-/// A positive value for `offset` shifts the array to the right
-/// a negative value shifts the array to the left.
-/// # Examples
-/// ```
-/// use arrow::array::Int32Array;
-/// use arrow::error::Result;
-/// use arrow::compute::shift;
-///
-/// let a: Int32Array = vec![Some(1), None, Some(4)].into();
-///
-/// // shift array 1 element to the right
-/// let res = shift(&a, 1).unwrap();
-/// let expected: Int32Array = vec![None, Some(1), None].into();
-/// assert_eq!(res.as_ref(), &expected);
-///
-/// // shift array 1 element to the left
-/// let res = shift(&a, -1).unwrap();
-/// let expected: Int32Array = vec![None, Some(4), None].into();
-/// assert_eq!(res.as_ref(), &expected);
-///
-/// // shift array 0 element, although not recommended
-/// let res = shift(&a, 0).unwrap();
-/// let expected: Int32Array = vec![Some(1), None, Some(4)].into();
-/// assert_eq!(res.as_ref(), &expected);
-///
-/// // shift array 3 element tot he right
-/// let res = shift(&a, 3).unwrap();
-/// let expected: Int32Array = vec![None, None, None].into();
-/// assert_eq!(res.as_ref(), &expected);
-/// ```
-pub fn shift(array: &Array, offset: i64) -> Result<ArrayRef> {
- let value_len = array.len() as i64;
- if offset == 0 {
- Ok(make_array(array.data_ref().clone()))
- } else if offset == i64::MIN || abs(offset) >= value_len {
- Ok(new_null_array(array.data_type(), array.len()))
- } else {
- let slice_offset = clamp(-offset, 0, value_len) as usize;
- let length = array.len() - abs(offset) as usize;
- let slice = array.slice(slice_offset, length);
-
- // Generate array with remaining `null` items
- let nulls = abs(offset) as usize;
- let null_arr = new_null_array(array.data_type(), nulls);
-
- // Concatenate both arrays, add nulls after if shift > 0 else before
- if offset > 0 {
- concat(&[null_arr.as_ref(), slice.as_ref()])
- } else {
- concat(&[slice.as_ref(), null_arr.as_ref()])
- }
- }
-}
-
-#[cfg(test)]
-mod tests {
- use super::*;
- use crate::array::{Float64Array, Int32Array, Int32DictionaryArray};
-
- #[test]
- fn test_shift_neg() {
- let a: Int32Array = vec![Some(1), None, Some(4)].into();
- let res = shift(&a, -1).unwrap();
- let expected: Int32Array = vec![None, Some(4), None].into();
- assert_eq!(res.as_ref(), &expected);
- }
-
- #[test]
- fn test_shift_pos() {
- let a: Int32Array = vec![Some(1), None, Some(4)].into();
- let res = shift(&a, 1).unwrap();
- let expected: Int32Array = vec![None, Some(1), None].into();
- assert_eq!(res.as_ref(), &expected);
- }
-
- #[test]
- fn test_shift_neg_float64() {
- let a: Float64Array = vec![Some(1.), None, Some(4.)].into();
- let res = shift(&a, -1).unwrap();
- let expected: Float64Array = vec![None, Some(4.), None].into();
- assert_eq!(res.as_ref(), &expected);
- }
-
- #[test]
- fn test_shift_pos_float64() {
- let a: Float64Array = vec![Some(1.), None, Some(4.)].into();
- let res = shift(&a, 1).unwrap();
- let expected: Float64Array = vec![None, Some(1.), None].into();
- assert_eq!(res.as_ref(), &expected);
- }
-
- #[test]
- fn test_shift_neg_int32_dict() {
- let a: Int32DictionaryArray = [Some("alpha"), None, Some("beta"), Some("alpha")]
- .iter()
- .copied()
- .collect();
- let res = shift(&a, -1).unwrap();
- let expected: Int32DictionaryArray = [None, Some("beta"), Some("alpha"), None]
- .iter()
- .copied()
- .collect();
- assert_eq!(res.as_ref(), &expected);
- }
-
- #[test]
- fn test_shift_pos_int32_dict() {
- let a: Int32DictionaryArray = [Some("alpha"), None, Some("beta"), Some("alpha")]
- .iter()
- .copied()
- .collect();
- let res = shift(&a, 1).unwrap();
- let expected: Int32DictionaryArray = [None, Some("alpha"), None, Some("beta")]
- .iter()
- .copied()
- .collect();
- assert_eq!(res.as_ref(), &expected);
- }
-
- #[test]
- fn test_shift_nil() {
- let a: Int32Array = vec![Some(1), None, Some(4)].into();
- let res = shift(&a, 0).unwrap();
- let expected: Int32Array = vec![Some(1), None, Some(4)].into();
- assert_eq!(res.as_ref(), &expected);
- }
-
- #[test]
- fn test_shift_boundary_pos() {
- let a: Int32Array = vec![Some(1), None, Some(4)].into();
- let res = shift(&a, 3).unwrap();
- let expected: Int32Array = vec![None, None, None].into();
- assert_eq!(res.as_ref(), &expected);
- }
-
- #[test]
- fn test_shift_boundary_neg() {
- let a: Int32Array = vec![Some(1), None, Some(4)].into();
- let res = shift(&a, -3).unwrap();
- let expected: Int32Array = vec![None, None, None].into();
- assert_eq!(res.as_ref(), &expected);
- }
-
- #[test]
- fn test_shift_boundary_neg_min() {
- let a: Int32Array = vec![Some(1), None, Some(4)].into();
- let res = shift(&a, i64::MIN).unwrap();
- let expected: Int32Array = vec![None, None, None].into();
- assert_eq!(res.as_ref(), &expected);
- }
-
- #[test]
- fn test_shift_large_pos() {
- let a: Int32Array = vec![Some(1), None, Some(4)].into();
- let res = shift(&a, 1000).unwrap();
- let expected: Int32Array = vec![None, None, None].into();
- assert_eq!(res.as_ref(), &expected);
- }
-
- #[test]
- fn test_shift_large_neg() {
- let a: Int32Array = vec![Some(1), None, Some(4)].into();
- let res = shift(&a, -1000).unwrap();
- let expected: Int32Array = vec![None, None, None].into();
- assert_eq!(res.as_ref(), &expected);
- }
-}
diff --git a/arrow/src/compute/kernels/zip.rs b/arrow/src/compute/kernels/zip.rs
deleted file mode 100644
index 0ee8e47..0000000
--- a/arrow/src/compute/kernels/zip.rs
+++ /dev/null
@@ -1,87 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-use crate::array::*;
-use crate::compute::SlicesIterator;
-use crate::error::{ArrowError, Result};
-
-/// Zip two arrays by some boolean mask. Where the mask evaluates `true` values of `truthy`
-/// are taken, where the mask evaluates `false` values of `falsy` are taken.
-///
-/// # Arguments
-/// * `mask` - Boolean values used to determine from which array to take the values.
-/// * `truthy` - Values of this array are taken if mask evaluates `true`
-/// * `falsy` - Values of this array are taken if mask evaluates `false`
-pub fn zip(
- mask: &BooleanArray,
- truthy: &dyn Array,
- falsy: &dyn Array,
-) -> Result<ArrayRef> {
- if truthy.data_type() != falsy.data_type() {
- return Err(ArrowError::InvalidArgumentError(
- "arguments need to have the same data type".into(),
- ));
- }
- if truthy.len() != falsy.len() || falsy.len() != mask.len() {
- return Err(ArrowError::InvalidArgumentError(
- "all arrays should have the same length".into(),
- ));
- }
- let falsy = falsy.data();
- let truthy = truthy.data();
-
- let mut mutable = MutableArrayData::new(vec![&*truthy, &*falsy], false, truthy.len());
-
- // the SlicesIterator slices only the true values. So the gaps left by this iterator we need to
- // fill with falsy values
-
- // keep track of how much is filled
- let mut filled = 0;
-
- SlicesIterator::new(mask).for_each(|(start, end)| {
- // the gap needs to be filled with falsy values
- if start > filled {
- mutable.extend(1, filled, start);
- }
- // fill with truthy values
- mutable.extend(0, start, end);
- filled = end;
- });
- // the remaining part is falsy
- if filled < truthy.len() {
- mutable.extend(1, filled, truthy.len());
- }
-
- let data = mutable.freeze();
- Ok(make_array(data))
-}
-
-#[cfg(test)]
-mod test {
- use super::*;
-
- #[test]
- fn test_zip_kernel() {
- let a = Int32Array::from(vec![Some(5), None, Some(7), None, Some(1)]);
- let b = Int32Array::from(vec![None, Some(3), Some(6), Some(7), Some(3)]);
- let mask = BooleanArray::from(vec![true, true, false, false, true]);
- let out = zip(&mask, &a, &b).unwrap();
- let actual = out.as_any().downcast_ref::<Int32Array>().unwrap();
- let expected = Int32Array::from(vec![Some(5), None, Some(6), Some(7), Some(1)]);
- assert_eq!(actual, &expected);
- }
-}
diff --git a/arrow/src/compute/mod.rs b/arrow/src/compute/mod.rs
deleted file mode 100644
index 2b3b9a7..0000000
--- a/arrow/src/compute/mod.rs
+++ /dev/null
@@ -1,38 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-//! Computation kernels on Arrow Arrays
-
-pub mod kernels;
-
-mod util;
-
-pub use self::kernels::aggregate::*;
-pub use self::kernels::arithmetic::*;
-pub use self::kernels::arity::*;
-pub use self::kernels::boolean::*;
-pub use self::kernels::cast::*;
-pub use self::kernels::comparison::*;
-pub use self::kernels::concat::*;
-pub use self::kernels::filter::*;
-pub use self::kernels::limit::*;
-pub use self::kernels::partition::*;
-pub use self::kernels::regexp::*;
-pub use self::kernels::sort::*;
-pub use self::kernels::take::*;
-pub use self::kernels::temporal::*;
-pub use self::kernels::window::*;
diff --git a/arrow/src/compute/util.rs b/arrow/src/compute/util.rs
deleted file mode 100644
index 56de594..0000000
--- a/arrow/src/compute/util.rs
+++ /dev/null
@@ -1,463 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-//! Common utilities for computation kernels.
-
-use crate::array::*;
-use crate::buffer::{buffer_bin_and, buffer_bin_or, Buffer};
-use crate::datatypes::*;
-use crate::error::{ArrowError, Result};
-use num::{One, ToPrimitive, Zero};
-use std::ops::Add;
-
-/// Combines the null bitmaps of two arrays using a bitwise `and` operation.
-///
-/// This function is useful when implementing operations on higher level arrays.
-#[allow(clippy::unnecessary_wraps)]
-pub(super) fn combine_option_bitmap(
- left_data: &ArrayData,
- right_data: &ArrayData,
- len_in_bits: usize,
-) -> Result<Option<Buffer>> {
- let left_offset_in_bits = left_data.offset();
- let right_offset_in_bits = right_data.offset();
-
- let left = left_data.null_buffer();
- let right = right_data.null_buffer();
-
- match left {
- None => match right {
- None => Ok(None),
- Some(r) => Ok(Some(r.bit_slice(right_offset_in_bits, len_in_bits))),
- },
- Some(l) => match right {
- None => Ok(Some(l.bit_slice(left_offset_in_bits, len_in_bits))),
-
- Some(r) => Ok(Some(buffer_bin_and(
- &l,
- left_offset_in_bits,
- &r,
- right_offset_in_bits,
- len_in_bits,
- ))),
- },
- }
-}
-
-/// Compares the null bitmaps of two arrays using a bitwise `or` operation.
-///
-/// This function is useful when implementing operations on higher level arrays.
-#[allow(clippy::unnecessary_wraps)]
-pub(super) fn compare_option_bitmap(
- left_data: &ArrayData,
- right_data: &ArrayData,
- len_in_bits: usize,
-) -> Result<Option<Buffer>> {
- let left_offset_in_bits = left_data.offset();
- let right_offset_in_bits = right_data.offset();
-
- let left = left_data.null_buffer();
- let right = right_data.null_buffer();
-
- match left {
- None => match right {
- None => Ok(None),
- Some(r) => Ok(Some(r.bit_slice(right_offset_in_bits, len_in_bits))),
- },
- Some(l) => match right {
- None => Ok(Some(l.bit_slice(left_offset_in_bits, len_in_bits))),
-
- Some(r) => Ok(Some(buffer_bin_or(
- &l,
- left_offset_in_bits,
- &r,
- right_offset_in_bits,
- len_in_bits,
- ))),
- },
- }
-}
-
-/// Takes/filters a list array's inner data using the offsets of the list array.
-///
-/// Where a list array has indices `[0,2,5,10]`, taking indices of `[2,0]` returns
-/// an array of the indices `[5..10, 0..2]` and offsets `[0,5,7]` (5 elements and 2
-/// elements)
-pub(super) fn take_value_indices_from_list<IndexType, OffsetType>(
- list: &GenericListArray<OffsetType::Native>,
- indices: &PrimitiveArray<IndexType>,
-) -> Result<(PrimitiveArray<OffsetType>, Vec<OffsetType::Native>)>
-where
- IndexType: ArrowNumericType,
- IndexType::Native: ToPrimitive,
- OffsetType: ArrowNumericType,
- OffsetType::Native: OffsetSizeTrait + Add + Zero + One,
- PrimitiveArray<OffsetType>: From<Vec<Option<OffsetType::Native>>>,
-{
- // TODO: benchmark this function, there might be a faster unsafe alternative
- let offsets: &[OffsetType::Native] = list.value_offsets();
-
- let mut new_offsets = Vec::with_capacity(indices.len());
- let mut values = Vec::new();
- let mut current_offset = OffsetType::Native::zero();
- // add first offset
- new_offsets.push(OffsetType::Native::zero());
- // compute the value indices, and set offsets accordingly
- for i in 0..indices.len() {
- if indices.is_valid(i) {
- let ix = ToPrimitive::to_usize(&indices.value(i)).ok_or_else(|| {
- ArrowError::ComputeError("Cast to usize failed".to_string())
- })?;
- let start = offsets[ix];
- let end = offsets[ix + 1];
- current_offset += end - start;
- new_offsets.push(current_offset);
-
- let mut curr = start;
-
- // if start == end, this slot is empty
- while curr < end {
- values.push(Some(curr));
- curr += OffsetType::Native::one();
- }
- } else {
- new_offsets.push(current_offset);
- }
- }
-
- Ok((PrimitiveArray::<OffsetType>::from(values), new_offsets))
-}
-
-/// Takes/filters a fixed size list array's inner data using the offsets of the list array.
-pub(super) fn take_value_indices_from_fixed_size_list<IndexType>(
- list: &FixedSizeListArray,
- indices: &PrimitiveArray<IndexType>,
- length: <UInt32Type as ArrowPrimitiveType>::Native,
-) -> Result<PrimitiveArray<UInt32Type>>
-where
- IndexType: ArrowNumericType,
- IndexType::Native: ToPrimitive,
-{
- let mut values = vec![];
-
- for i in 0..indices.len() {
- if indices.is_valid(i) {
- let index = ToPrimitive::to_usize(&indices.value(i)).ok_or_else(|| {
- ArrowError::ComputeError("Cast to usize failed".to_string())
- })?;
- let start =
- list.value_offset(index) as <UInt32Type as ArrowPrimitiveType>::Native;
-
- values.extend(start..start + length);
- }
- }
-
- Ok(PrimitiveArray::<UInt32Type>::from(values))
-}
-
-#[cfg(test)]
-pub(super) mod tests {
- use super::*;
-
- use std::sync::Arc;
-
- use crate::datatypes::DataType;
- use crate::util::bit_util;
- use crate::{array::ArrayData, buffer::MutableBuffer};
-
- fn make_data_with_null_bit_buffer(
- len: usize,
- offset: usize,
- null_bit_buffer: Option<Buffer>,
- ) -> Arc<ArrayData> {
- // empty vec for buffers and children is not really correct, but for these tests we only care about the null bitmap
- Arc::new(ArrayData::new(
- DataType::UInt8,
- len,
- None,
- null_bit_buffer,
- offset,
- vec![],
- vec![],
- ))
- }
-
- #[test]
- fn test_combine_option_bitmap() {
- let none_bitmap = make_data_with_null_bit_buffer(8, 0, None);
- let some_bitmap =
- make_data_with_null_bit_buffer(8, 0, Some(Buffer::from([0b01001010])));
- let inverse_bitmap =
- make_data_with_null_bit_buffer(8, 0, Some(Buffer::from([0b10110101])));
- assert_eq!(
- None,
- combine_option_bitmap(&none_bitmap, &none_bitmap, 8).unwrap()
- );
- assert_eq!(
- Some(Buffer::from([0b01001010])),
- combine_option_bitmap(&some_bitmap, &none_bitmap, 8).unwrap()
- );
- assert_eq!(
- Some(Buffer::from([0b01001010])),
- combine_option_bitmap(&none_bitmap, &some_bitmap, 8,).unwrap()
- );
- assert_eq!(
- Some(Buffer::from([0b01001010])),
- combine_option_bitmap(&some_bitmap, &some_bitmap, 8,).unwrap()
- );
- assert_eq!(
- Some(Buffer::from([0b0])),
- combine_option_bitmap(&some_bitmap, &inverse_bitmap, 8,).unwrap()
- );
- }
-
- #[test]
- fn test_compare_option_bitmap() {
- let none_bitmap = make_data_with_null_bit_buffer(8, 0, None);
- let some_bitmap =
- make_data_with_null_bit_buffer(8, 0, Some(Buffer::from([0b01001010])));
- let inverse_bitmap =
- make_data_with_null_bit_buffer(8, 0, Some(Buffer::from([0b10110101])));
- assert_eq!(
- None,
- compare_option_bitmap(&none_bitmap, &none_bitmap, 8).unwrap()
- );
- assert_eq!(
- Some(Buffer::from([0b01001010])),
- compare_option_bitmap(&some_bitmap, &none_bitmap, 8).unwrap()
- );
- assert_eq!(
- Some(Buffer::from([0b01001010])),
- compare_option_bitmap(&none_bitmap, &some_bitmap, 8,).unwrap()
- );
- assert_eq!(
- Some(Buffer::from([0b01001010])),
- compare_option_bitmap(&some_bitmap, &some_bitmap, 8,).unwrap()
- );
- assert_eq!(
- Some(Buffer::from([0b11111111])),
- compare_option_bitmap(&some_bitmap, &inverse_bitmap, 8,).unwrap()
- );
- }
-
- pub(crate) fn build_generic_list<S, T>(
- data: Vec<Option<Vec<T::Native>>>,
- ) -> GenericListArray<S>
- where
- S: OffsetSizeTrait + 'static,
- T: ArrowPrimitiveType,
- PrimitiveArray<T>: From<Vec<Option<T::Native>>>,
- {
- let data = data
- .into_iter()
- .map(|subarray| {
- subarray.map(|item| {
- item.into_iter()
- .map(Some)
- .collect::<Vec<Option<T::Native>>>()
- })
- })
- .collect();
- build_generic_list_nullable(data)
- }
-
- pub(crate) fn build_generic_list_nullable<S, T>(
- data: Vec<Option<Vec<Option<T::Native>>>>,
- ) -> GenericListArray<S>
- where
- S: OffsetSizeTrait + 'static,
- T: ArrowPrimitiveType,
- PrimitiveArray<T>: From<Vec<Option<T::Native>>>,
- {
- use std::any::TypeId;
-
- let mut offset = vec![0];
- let mut values = vec![];
-
- let list_len = data.len();
- let num_bytes = bit_util::ceil(list_len, 8);
- let mut list_null_count = 0;
- let mut list_bitmap = MutableBuffer::new(num_bytes).with_bitset(num_bytes, true);
- for (idx, array) in data.into_iter().enumerate() {
- if let Some(mut array) = array {
- values.append(&mut array);
- } else {
- list_null_count += 1;
- bit_util::unset_bit(&mut list_bitmap.as_slice_mut(), idx);
- }
- offset.push(values.len() as i64);
- }
-
- let value_data = PrimitiveArray::<T>::from(values).data().clone();
- let (list_data_type, value_offsets) = if TypeId::of::<S>() == TypeId::of::<i32>()
- {
- (
- DataType::List(Box::new(Field::new(
- "item",
- T::DATA_TYPE,
- list_null_count == 0,
- ))),
- Buffer::from_slice_ref(
- &offset.into_iter().map(|x| x as i32).collect::<Vec<i32>>(),
- ),
- )
- } else if TypeId::of::<S>() == TypeId::of::<i64>() {
- (
- DataType::LargeList(Box::new(Field::new(
- "item",
- T::DATA_TYPE,
- list_null_count == 0,
- ))),
- Buffer::from_slice_ref(&offset),
- )
- } else {
- unreachable!()
- };
-
- let list_data = ArrayData::builder(list_data_type)
- .len(list_len)
- .null_bit_buffer(list_bitmap.into())
- .add_buffer(value_offsets)
- .add_child_data(value_data)
- .build();
-
- GenericListArray::<S>::from(list_data)
- }
-
- pub(crate) fn build_fixed_size_list<T>(
- data: Vec<Option<Vec<T::Native>>>,
- length: <Int32Type as ArrowPrimitiveType>::Native,
- ) -> FixedSizeListArray
- where
- T: ArrowPrimitiveType,
- PrimitiveArray<T>: From<Vec<Option<T::Native>>>,
- {
- let data = data
- .into_iter()
- .map(|subarray| {
- subarray.map(|item| {
- item.into_iter()
- .map(Some)
- .collect::<Vec<Option<T::Native>>>()
- })
- })
- .collect();
- build_fixed_size_list_nullable(data, length)
- }
-
- pub(crate) fn build_fixed_size_list_nullable<T>(
- list_values: Vec<Option<Vec<Option<T::Native>>>>,
- length: <Int32Type as ArrowPrimitiveType>::Native,
- ) -> FixedSizeListArray
- where
- T: ArrowPrimitiveType,
- PrimitiveArray<T>: From<Vec<Option<T::Native>>>,
- {
- let mut values = vec![];
- let mut list_null_count = 0;
- let list_len = list_values.len();
-
- let num_bytes = bit_util::ceil(list_len, 8);
- let mut list_bitmap = MutableBuffer::new(num_bytes).with_bitset(num_bytes, true);
- for (idx, list_element) in list_values.into_iter().enumerate() {
- if let Some(items) = list_element {
- // every sub-array should have the same length
- debug_assert_eq!(length as usize, items.len());
-
- values.extend(items.into_iter());
- } else {
- list_null_count += 1;
- bit_util::unset_bit(&mut list_bitmap.as_slice_mut(), idx);
- values.extend(vec![None; length as usize].into_iter());
- }
- }
-
- let list_data_type = DataType::FixedSizeList(
- Box::new(Field::new("item", T::DATA_TYPE, list_null_count == 0)),
- length,
- );
-
- let child_data = PrimitiveArray::<T>::from(values).data().clone();
-
- let list_data = ArrayData::builder(list_data_type)
- .len(list_len)
- .null_bit_buffer(list_bitmap.into())
- .add_child_data(child_data)
- .build();
-
- FixedSizeListArray::from(list_data)
- }
-
- #[test]
- fn test_take_value_index_from_list() {
- let list = build_generic_list::<i32, Int32Type>(vec![
- Some(vec![0, 1]),
- Some(vec![2, 3, 4]),
- Some(vec![5, 6, 7, 8, 9]),
- ]);
- let indices = UInt32Array::from(vec![2, 0]);
-
- let (indexed, offsets) = take_value_indices_from_list(&list, &indices).unwrap();
-
- assert_eq!(indexed, Int32Array::from(vec![5, 6, 7, 8, 9, 0, 1]));
- assert_eq!(offsets, vec![0, 5, 7]);
- }
-
- #[test]
- fn test_take_value_index_from_large_list() {
- let list = build_generic_list::<i64, Int32Type>(vec![
- Some(vec![0, 1]),
- Some(vec![2, 3, 4]),
- Some(vec![5, 6, 7, 8, 9]),
- ]);
- let indices = UInt32Array::from(vec![2, 0]);
-
- let (indexed, offsets) =
- take_value_indices_from_list::<_, Int64Type>(&list, &indices).unwrap();
-
- assert_eq!(indexed, Int64Array::from(vec![5, 6, 7, 8, 9, 0, 1]));
- assert_eq!(offsets, vec![0, 5, 7]);
- }
-
- #[test]
- fn test_take_value_index_from_fixed_list() {
- let list = build_fixed_size_list_nullable::<Int32Type>(
- vec![
- Some(vec![Some(1), Some(2), None]),
- Some(vec![Some(4), None, Some(6)]),
- None,
- Some(vec![None, Some(8), Some(9)]),
- ],
- 3,
- );
-
- let indices = UInt32Array::from(vec![2, 1, 0]);
- let indexed =
- take_value_indices_from_fixed_size_list(&list, &indices, 3).unwrap();
-
- assert_eq!(indexed, UInt32Array::from(vec![6, 7, 8, 3, 4, 5, 0, 1, 2]));
-
- let indices = UInt32Array::from(vec![3, 2, 1, 2, 0]);
- let indexed =
- take_value_indices_from_fixed_size_list(&list, &indices, 3).unwrap();
-
- assert_eq!(
- indexed,
- UInt32Array::from(vec![9, 10, 11, 6, 7, 8, 3, 4, 5, 6, 7, 8, 0, 1, 2])
- );
- }
-}
diff --git a/arrow/src/csv/mod.rs b/arrow/src/csv/mod.rs
deleted file mode 100644
index ffe82f3..0000000
--- a/arrow/src/csv/mod.rs
+++ /dev/null
@@ -1,27 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-//! Transfer data between the Arrow memory format and CSV (comma-separated values).
-
-pub mod reader;
-pub mod writer;
-
-pub use self::reader::infer_schema_from_files;
-pub use self::reader::Reader;
-pub use self::reader::ReaderBuilder;
-pub use self::writer::Writer;
-pub use self::writer::WriterBuilder;
diff --git a/arrow/src/csv/reader.rs b/arrow/src/csv/reader.rs
deleted file mode 100644
index b2f7a7a..0000000
--- a/arrow/src/csv/reader.rs
+++ /dev/null
@@ -1,1599 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-//! CSV Reader
-//!
-//! This CSV reader allows CSV files to be read into the Arrow memory model. Records are
-//! loaded in batches and are then converted from row-based data to columnar data.
-//!
-//! Example:
-//!
-//! ```
-//! use arrow::csv;
-//! use arrow::datatypes::{DataType, Field, Schema};
-//! use std::fs::File;
-//! use std::sync::Arc;
-//!
-//! let schema = Schema::new(vec![
-//! Field::new("city", DataType::Utf8, false),
-//! Field::new("lat", DataType::Float64, false),
-//! Field::new("lng", DataType::Float64, false),
-//! ]);
-//!
-//! let file = File::open("test/data/uk_cities.csv").unwrap();
-//!
-//! let mut csv = csv::Reader::new(file, Arc::new(schema), false, None, 1024, None, None);
-//! let batch = csv.next().unwrap().unwrap();
-//! ```
-
-use core::cmp::min;
-use lazy_static::lazy_static;
-use regex::{Regex, RegexBuilder};
-use std::collections::HashSet;
-use std::fmt;
-use std::fs::File;
-use std::io::{Read, Seek, SeekFrom};
-use std::sync::Arc;
-
-use crate::array::{
- ArrayRef, BooleanArray, DictionaryArray, PrimitiveArray, StringArray,
-};
-use crate::datatypes::*;
-use crate::error::{ArrowError, Result};
-use crate::record_batch::RecordBatch;
-
-use csv_crate::{ByteRecord, StringRecord};
-
-lazy_static! {
- static ref DECIMAL_RE: Regex = Regex::new(r"^-?(\d+\.\d+)$").unwrap();
- static ref INTEGER_RE: Regex = Regex::new(r"^-?(\d+)$").unwrap();
- static ref BOOLEAN_RE: Regex = RegexBuilder::new(r"^(true)$|^(false)$")
- .case_insensitive(true)
- .build()
- .unwrap();
- static ref DATE_RE: Regex = Regex::new(r"^\d{4}-\d\d-\d\d$").unwrap();
- static ref DATETIME_RE: Regex =
- Regex::new(r"^\d{4}-\d\d-\d\dT\d\d:\d\d:\d\d$").unwrap();
-}
-
-/// Infer the data type of a record
-fn infer_field_schema(string: &str) -> DataType {
- // when quoting is enabled in the reader, these quotes aren't escaped, we default to
- // Utf8 for them
- if string.starts_with('"') {
- return DataType::Utf8;
- }
- // match regex in a particular order
- if BOOLEAN_RE.is_match(string) {
- DataType::Boolean
- } else if DECIMAL_RE.is_match(string) {
- DataType::Float64
- } else if INTEGER_RE.is_match(string) {
- DataType::Int64
- } else if DATETIME_RE.is_match(string) {
- DataType::Date64
- } else if DATE_RE.is_match(string) {
- DataType::Date32
- } else {
- DataType::Utf8
- }
-}
-
-/// Infer the schema of a CSV file by reading through the first n records of the file,
-/// with `max_read_records` controlling the maximum number of records to read.
-///
-/// If `max_read_records` is not set, the whole file is read to infer its schema.
-///
-/// Return infered schema and number of records used for inference. This function does not change
-/// reader cursor offset.
-pub fn infer_file_schema<R: Read + Seek>(
- reader: &mut R,
- delimiter: u8,
- max_read_records: Option<usize>,
- has_header: bool,
-) -> Result<(Schema, usize)> {
- infer_file_schema_with_csv_options(
- reader,
- delimiter,
- max_read_records,
- has_header,
- None,
- None,
- None,
- )
-}
-
-fn infer_file_schema_with_csv_options<R: Read + Seek>(
- reader: &mut R,
- delimiter: u8,
- max_read_records: Option<usize>,
- has_header: bool,
- escape: Option<u8>,
- quote: Option<u8>,
- terminator: Option<u8>,
-) -> Result<(Schema, usize)> {
- let saved_offset = reader.seek(SeekFrom::Current(0))?;
-
- let (schema, records_count) = infer_reader_schema_with_csv_options(
- reader,
- delimiter,
- max_read_records,
- has_header,
- escape,
- quote,
- terminator,
- )?;
-
- // return the reader seek back to the start
- reader.seek(SeekFrom::Start(saved_offset))?;
-
- Ok((schema, records_count))
-}
-
-/// Infer schema of CSV records provided by struct that implements `Read` trait.
-///
-/// `max_read_records` controlling the maximum number of records to read. If `max_read_records` is
-/// not set, all records are read to infer the schema.
-///
-/// Return infered schema and number of records used for inference.
-pub fn infer_reader_schema<R: Read>(
- reader: &mut R,
- delimiter: u8,
- max_read_records: Option<usize>,
- has_header: bool,
-) -> Result<(Schema, usize)> {
- infer_reader_schema_with_csv_options(
- reader,
- delimiter,
- max_read_records,
- has_header,
- None,
- None,
- None,
- )
-}
-
-fn infer_reader_schema_with_csv_options<R: Read>(
- reader: &mut R,
- delimiter: u8,
- max_read_records: Option<usize>,
- has_header: bool,
- escape: Option<u8>,
- quote: Option<u8>,
- terminator: Option<u8>,
-) -> Result<(Schema, usize)> {
- let mut csv_reader = Reader::build_csv_reader(
- reader,
- has_header,
- Some(delimiter),
- escape,
- quote,
- terminator,
- );
-
- // get or create header names
- // when has_header is false, creates default column names with column_ prefix
- let headers: Vec<String> = if has_header {
- let headers = &csv_reader.headers()?.clone();
- headers.iter().map(|s| s.to_string()).collect()
- } else {
- let first_record_count = &csv_reader.headers()?.len();
- (0..*first_record_count)
- .map(|i| format!("column_{}", i + 1))
- .collect()
- };
-
- let header_length = headers.len();
- // keep track of inferred field types
- let mut column_types: Vec<HashSet<DataType>> = vec![HashSet::new(); header_length];
- // keep track of columns with nulls
- let mut nulls: Vec<bool> = vec![false; header_length];
-
- let mut records_count = 0;
- let mut fields = vec![];
-
- let mut record = StringRecord::new();
- let max_records = max_read_records.unwrap_or(usize::MAX);
- while records_count < max_records {
- if !csv_reader.read_record(&mut record)? {
- break;
- }
- records_count += 1;
-
- for i in 0..header_length {
- if let Some(string) = record.get(i) {
- if string.is_empty() {
- nulls[i] = true;
- } else {
- column_types[i].insert(infer_field_schema(string));
- }
- }
- }
- }
-
- // build schema from inference results
- for i in 0..header_length {
- let possibilities = &column_types[i];
- let has_nulls = nulls[i];
- let field_name = &headers[i];
-
- // determine data type based on possible types
- // if there are incompatible types, use DataType::Utf8
- match possibilities.len() {
- 1 => {
- for dtype in possibilities.iter() {
- fields.push(Field::new(&field_name, dtype.clone(), has_nulls));
- }
- }
- 2 => {
- if possibilities.contains(&DataType::Int64)
- && possibilities.contains(&DataType::Float64)
- {
- // we have an integer and double, fall down to double
- fields.push(Field::new(&field_name, DataType::Float64, has_nulls));
- } else {
- // default to Utf8 for conflicting datatypes (e.g bool and int)
- fields.push(Field::new(&field_name, DataType::Utf8, has_nulls));
- }
- }
- _ => fields.push(Field::new(&field_name, DataType::Utf8, has_nulls)),
- }
- }
-
- Ok((Schema::new(fields), records_count))
-}
-
-/// Infer schema from a list of CSV files by reading through first n records
-/// with `max_read_records` controlling the maximum number of records to read.
-///
-/// Files will be read in the given order untill n records have been reached.
-///
-/// If `max_read_records` is not set, all files will be read fully to infer the schema.
-pub fn infer_schema_from_files(
- files: &[String],
- delimiter: u8,
- max_read_records: Option<usize>,
- has_header: bool,
-) -> Result<Schema> {
- let mut schemas = vec![];
- let mut records_to_read = max_read_records.unwrap_or(std::usize::MAX);
-
- for fname in files.iter() {
- let (schema, records_read) = infer_file_schema(
- &mut File::open(fname)?,
- delimiter,
- Some(records_to_read),
- has_header,
- )?;
- if records_read == 0 {
- continue;
- }
- schemas.push(schema.clone());
- records_to_read -= records_read;
- if records_to_read == 0 {
- break;
- }
- }
-
- Schema::try_merge(schemas)
-}
-
-// optional bounds of the reader, of the form (min line, max line).
-type Bounds = Option<(usize, usize)>;
-
-/// CSV file reader
-pub struct Reader<R: Read> {
- /// Explicit schema for the CSV file
- schema: SchemaRef,
- /// Optional projection for which columns to load (zero-based column indices)
- projection: Option<Vec<usize>>,
- /// File reader
- reader: csv_crate::Reader<R>,
- /// Current line number
- line_number: usize,
- /// Maximum number of rows to read
- end: usize,
- /// Number of records per batch
- batch_size: usize,
- /// Vector that can hold the `StringRecord`s of the batches
- batch_records: Vec<StringRecord>,
-}
-
-impl<R> fmt::Debug for Reader<R>
-where
- R: Read,
-{
- fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
- f.debug_struct("Reader")
- .field("schema", &self.schema)
- .field("projection", &self.projection)
- .field("line_number", &self.line_number)
- .finish()
- }
-}
-
-impl<R: Read> Reader<R> {
- /// Create a new CsvReader from any value that implements the `Read` trait.
- ///
- /// If reading a `File` or an input that supports `std::io::Read` and `std::io::Seek`;
- /// you can customise the Reader, such as to enable schema inference, use
- /// `ReaderBuilder`.
- pub fn new(
- reader: R,
- schema: SchemaRef,
- has_header: bool,
- delimiter: Option<u8>,
- batch_size: usize,
- bounds: Bounds,
- projection: Option<Vec<usize>>,
- ) -> Self {
- Self::from_reader(
- reader, schema, has_header, delimiter, batch_size, bounds, projection,
- )
- }
-
- /// Returns the schema of the reader, useful for getting the schema without reading
- /// record batches
- pub fn schema(&self) -> SchemaRef {
- match &self.projection {
- Some(projection) => {
- let fields = self.schema.fields();
- let projected_fields: Vec<Field> =
- projection.iter().map(|i| fields[*i].clone()).collect();
-
- Arc::new(Schema::new(projected_fields))
- }
- None => self.schema.clone(),
- }
- }
-
- /// Create a new CsvReader from a Reader
- ///
- /// This constructor allows you more flexibility in what records are processed by the
- /// csv reader.
- pub fn from_reader(
- reader: R,
- schema: SchemaRef,
- has_header: bool,
- delimiter: Option<u8>,
- batch_size: usize,
- bounds: Bounds,
- projection: Option<Vec<usize>>,
- ) -> Self {
- let csv_reader =
- Self::build_csv_reader(reader, has_header, delimiter, None, None, None);
- Self::from_csv_reader(
- csv_reader, schema, has_header, batch_size, bounds, projection,
- )
- }
-
- fn build_csv_reader(
- reader: R,
- has_header: bool,
- delimiter: Option<u8>,
- escape: Option<u8>,
- quote: Option<u8>,
- terminator: Option<u8>,
- ) -> csv_crate::Reader<R> {
- let mut reader_builder = csv_crate::ReaderBuilder::new();
- reader_builder.has_headers(has_header);
-
- if let Some(c) = delimiter {
- reader_builder.delimiter(c);
- }
- reader_builder.escape(escape);
- if let Some(c) = quote {
- reader_builder.quote(c);
- }
- if let Some(t) = terminator {
- reader_builder.terminator(csv_crate::Terminator::Any(t));
- }
- reader_builder.from_reader(reader)
- }
-
- fn from_csv_reader(
- mut csv_reader: csv_crate::Reader<R>,
- schema: SchemaRef,
- has_header: bool,
- batch_size: usize,
- bounds: Bounds,
- projection: Option<Vec<usize>>,
- ) -> Self {
- let (start, end) = match bounds {
- None => (0, usize::MAX),
- Some((start, end)) => (start, end),
- };
-
- // First we will skip `start` rows
- // note that this skips by iteration. This is because in general it is not possible
- // to seek in CSV. However, skiping still saves the burden of creating arrow arrays,
- // which is a slow operation that scales with the number of columns
-
- let mut record = ByteRecord::new();
- // Skip first start items
- for _ in 0..start {
- let res = csv_reader.read_byte_record(&mut record);
- if !res.unwrap_or(false) {
- break;
- }
- }
-
- // Initialize batch_records with StringRecords so they
- // can be reused across batches
- let mut batch_records = Vec::with_capacity(batch_size);
- batch_records.resize_with(batch_size, Default::default);
-
- Self {
- schema,
- projection,
- reader: csv_reader,
- line_number: if has_header { start + 1 } else { start },
- batch_size,
- end,
- batch_records,
- }
- }
-}
-
-impl<R: Read> Iterator for Reader<R> {
- type Item = Result<RecordBatch>;
-
- fn next(&mut self) -> Option<Self::Item> {
- let remaining = self.end - self.line_number;
-
- let mut read_records = 0;
- for i in 0..min(self.batch_size, remaining) {
- match self.reader.read_record(&mut self.batch_records[i]) {
- Ok(true) => {
- read_records += 1;
- }
- Ok(false) => break,
- Err(e) => {
- return Some(Err(ArrowError::ParseError(format!(
- "Error parsing line {}: {:?}",
- self.line_number + i,
- e
- ))));
- }
- }
- }
-
- // return early if no data was loaded
- if read_records == 0 {
- return None;
- }
-
- // parse the batches into a RecordBatch
- let result = parse(
- &self.batch_records[..read_records],
- &self.schema.fields(),
- Some(self.schema.metadata.clone()),
- &self.projection,
- self.line_number,
- );
-
- self.line_number += read_records;
-
- Some(result)
- }
-}
-
-/// parses a slice of [csv_crate::StringRecord] into a [array::record_batch::RecordBatch].
-fn parse(
- rows: &[StringRecord],
- fields: &[Field],
- metadata: Option<std::collections::HashMap<String, String>>,
- projection: &Option<Vec<usize>>,
- line_number: usize,
-) -> Result<RecordBatch> {
- let projection: Vec<usize> = match projection {
- Some(ref v) => v.clone(),
- None => fields.iter().enumerate().map(|(i, _)| i).collect(),
- };
-
- let arrays: Result<Vec<ArrayRef>> = projection
- .iter()
- .map(|i| {
- let i = *i;
- let field = &fields[i];
- match field.data_type() {
- DataType::Boolean => build_boolean_array(line_number, rows, i),
- DataType::Int8 => build_primitive_array::<Int8Type>(line_number, rows, i),
- DataType::Int16 => {
- build_primitive_array::<Int16Type>(line_number, rows, i)
- }
- DataType::Int32 => {
- build_primitive_array::<Int32Type>(line_number, rows, i)
- }
- DataType::Int64 => {
- build_primitive_array::<Int64Type>(line_number, rows, i)
- }
- DataType::UInt8 => {
- build_primitive_array::<UInt8Type>(line_number, rows, i)
- }
- DataType::UInt16 => {
- build_primitive_array::<UInt16Type>(line_number, rows, i)
- }
- DataType::UInt32 => {
- build_primitive_array::<UInt32Type>(line_number, rows, i)
- }
- DataType::UInt64 => {
- build_primitive_array::<UInt64Type>(line_number, rows, i)
- }
- DataType::Float32 => {
- build_primitive_array::<Float32Type>(line_number, rows, i)
- }
- DataType::Float64 => {
- build_primitive_array::<Float64Type>(line_number, rows, i)
- }
- DataType::Date32 => {
- build_primitive_array::<Date32Type>(line_number, rows, i)
- }
- DataType::Date64 => {
- build_primitive_array::<Date64Type>(line_number, rows, i)
- }
- DataType::Timestamp(TimeUnit::Microsecond, _) => build_primitive_array::<
- TimestampMicrosecondType,
- >(
- line_number, rows, i
- ),
- DataType::Timestamp(TimeUnit::Nanosecond, _) => {
- build_primitive_array::<TimestampNanosecondType>(line_number, rows, i)
- }
- DataType::Utf8 => Ok(Arc::new(
- rows.iter().map(|row| row.get(i)).collect::<StringArray>(),
- ) as ArrayRef),
- DataType::Dictionary(key_type, value_type)
- if value_type.as_ref() == &DataType::Utf8 =>
- {
- match key_type.as_ref() {
- DataType::Int8 => Ok(Arc::new(
- rows.iter()
- .map(|row| row.get(i))
- .collect::<DictionaryArray<Int8Type>>(),
- ) as ArrayRef),
- DataType::Int16 => Ok(Arc::new(
- rows.iter()
- .map(|row| row.get(i))
- .collect::<DictionaryArray<Int16Type>>(),
- ) as ArrayRef),
- DataType::Int32 => Ok(Arc::new(
- rows.iter()
- .map(|row| row.get(i))
- .collect::<DictionaryArray<Int32Type>>(),
- ) as ArrayRef),
- DataType::Int64 => Ok(Arc::new(
- rows.iter()
- .map(|row| row.get(i))
- .collect::<DictionaryArray<Int64Type>>(),
- ) as ArrayRef),
- DataType::UInt8 => Ok(Arc::new(
- rows.iter()
- .map(|row| row.get(i))
- .collect::<DictionaryArray<UInt8Type>>(),
- ) as ArrayRef),
- DataType::UInt16 => Ok(Arc::new(
- rows.iter()
- .map(|row| row.get(i))
- .collect::<DictionaryArray<UInt16Type>>(),
- ) as ArrayRef),
- DataType::UInt32 => Ok(Arc::new(
- rows.iter()
- .map(|row| row.get(i))
- .collect::<DictionaryArray<UInt32Type>>(),
- ) as ArrayRef),
- DataType::UInt64 => Ok(Arc::new(
- rows.iter()
- .map(|row| row.get(i))
- .collect::<DictionaryArray<UInt64Type>>(),
- ) as ArrayRef),
- _ => Err(ArrowError::ParseError(format!(
- "Unsupported dictionary key type {:?}",
- key_type
- ))),
- }
- }
- other => Err(ArrowError::ParseError(format!(
- "Unsupported data type {:?}",
- other
- ))),
- }
- })
- .collect();
-
- let projected_fields: Vec<Field> =
- projection.iter().map(|i| fields[*i].clone()).collect();
-
- let projected_schema = Arc::new(match metadata {
- None => Schema::new(projected_fields),
- Some(metadata) => Schema::new_with_metadata(projected_fields, metadata),
- });
-
- arrays.and_then(|arr| RecordBatch::try_new(projected_schema, arr))
-}
-
-/// Specialized parsing implementations
-trait Parser: ArrowPrimitiveType {
- fn parse(string: &str) -> Option<Self::Native> {
- string.parse::<Self::Native>().ok()
- }
-}
-
-impl Parser for Float32Type {
- fn parse(string: &str) -> Option<f32> {
- lexical_core::parse(string.as_bytes()).ok()
- }
-}
-
-impl Parser for Float64Type {
- fn parse(string: &str) -> Option<f64> {
- lexical_core::parse(string.as_bytes()).ok()
- }
-}
-
-impl Parser for UInt64Type {}
-
-impl Parser for UInt32Type {}
-
-impl Parser for UInt16Type {}
-
-impl Parser for UInt8Type {}
-
-impl Parser for Int64Type {}
-
-impl Parser for Int32Type {}
-
-impl Parser for Int16Type {}
-
-impl Parser for Int8Type {}
-
-/// Number of days between 0001-01-01 and 1970-01-01
-const EPOCH_DAYS_FROM_CE: i32 = 719_163;
-
-impl Parser for Date32Type {
- fn parse(string: &str) -> Option<i32> {
- use chrono::Datelike;
-
- match Self::DATA_TYPE {
- DataType::Date32 => {
- let date = string.parse::<chrono::NaiveDate>().ok()?;
- Self::Native::from_i32(date.num_days_from_ce() - EPOCH_DAYS_FROM_CE)
- }
- _ => None,
- }
- }
-}
-
-impl Parser for Date64Type {
- fn parse(string: &str) -> Option<i64> {
- match Self::DATA_TYPE {
- DataType::Date64 => {
- let date_time = string.parse::<chrono::NaiveDateTime>().ok()?;
- Self::Native::from_i64(date_time.timestamp_millis())
- }
- _ => None,
- }
- }
-}
-
-impl Parser for TimestampNanosecondType {
- fn parse(string: &str) -> Option<i64> {
- match Self::DATA_TYPE {
- DataType::Timestamp(TimeUnit::Nanosecond, None) => {
- let date_time = string.parse::<chrono::NaiveDateTime>().ok()?;
- Self::Native::from_i64(date_time.timestamp_nanos())
- }
- _ => None,
- }
- }
-}
-
-impl Parser for TimestampMicrosecondType {
- fn parse(string: &str) -> Option<i64> {
- match Self::DATA_TYPE {
- DataType::Timestamp(TimeUnit::Microsecond, None) => {
- let date_time = string.parse::<chrono::NaiveDateTime>().ok()?;
- Self::Native::from_i64(date_time.timestamp_nanos() / 1000)
- }
- _ => None,
- }
- }
-}
-
-fn parse_item<T: Parser>(string: &str) -> Option<T::Native> {
- T::parse(string)
-}
-
-fn parse_bool(string: &str) -> Option<bool> {
- if string.eq_ignore_ascii_case("false") {
- Some(false)
- } else if string.eq_ignore_ascii_case("true") {
- Some(true)
- } else {
- None
- }
-}
-
-// parses a specific column (col_idx) into an Arrow Array.
-fn build_primitive_array<T: ArrowPrimitiveType + Parser>(
- line_number: usize,
- rows: &[StringRecord],
- col_idx: usize,
-) -> Result<ArrayRef> {
- rows.iter()
- .enumerate()
- .map(|(row_index, row)| {
- match row.get(col_idx) {
- Some(s) => {
- if s.is_empty() {
- return Ok(None);
- }
-
- let parsed = parse_item::<T>(s);
- match parsed {
- Some(e) => Ok(Some(e)),
- None => Err(ArrowError::ParseError(format!(
- // TODO: we should surface the underlying error here.
- "Error while parsing value {} for column {} at line {}",
- s,
- col_idx,
- line_number + row_index
- ))),
- }
- }
- None => Ok(None),
- }
- })
- .collect::<Result<PrimitiveArray<T>>>()
- .map(|e| Arc::new(e) as ArrayRef)
-}
-
-// parses a specific column (col_idx) into an Arrow Array.
-fn build_boolean_array(
- line_number: usize,
- rows: &[StringRecord],
- col_idx: usize,
-) -> Result<ArrayRef> {
- rows.iter()
- .enumerate()
- .map(|(row_index, row)| {
- match row.get(col_idx) {
- Some(s) => {
- if s.is_empty() {
- return Ok(None);
- }
-
- let parsed = parse_bool(s);
- match parsed {
- Some(e) => Ok(Some(e)),
- None => Err(ArrowError::ParseError(format!(
- // TODO: we should surface the underlying error here.
- "Error while parsing value {} for column {} at line {}",
- s,
- col_idx,
- line_number + row_index
- ))),
- }
- }
- None => Ok(None),
- }
- })
- .collect::<Result<BooleanArray>>()
- .map(|e| Arc::new(e) as ArrayRef)
-}
-
-/// CSV file reader builder
-#[derive(Debug)]
-pub struct ReaderBuilder {
- /// Optional schema for the CSV file
- ///
- /// If the schema is not supplied, the reader will try to infer the schema
- /// based on the CSV structure.
- schema: Option<SchemaRef>,
- /// Whether the file has headers or not
- ///
- /// If schema inference is run on a file with no headers, default column names
- /// are created.
- has_header: bool,
- /// An optional column delimiter. Defaults to `b','`
- delimiter: Option<u8>,
- /// An optional escape charactor. Defaults None
- escape: Option<u8>,
- /// An optional quote charactor. Defaults b'\"'
- quote: Option<u8>,
- /// An optional record terminator. Defaults CRLF
- terminator: Option<u8>,
- /// Optional maximum number of records to read during schema inference
- ///
- /// If a number is not provided, all the records are read.
- max_records: Option<usize>,
- /// Batch size (number of records to load each time)
- ///
- /// The default batch size when using the `ReaderBuilder` is 1024 records
- batch_size: usize,
- /// The bounds over which to scan the reader. `None` starts from 0 and runs until EOF.
- bounds: Bounds,
- /// Optional projection for which columns to load (zero-based column indices)
- projection: Option<Vec<usize>>,
-}
-
-impl Default for ReaderBuilder {
- fn default() -> Self {
- Self {
- schema: None,
- has_header: false,
- delimiter: None,
- escape: None,
- quote: None,
- terminator: None,
- max_records: None,
- batch_size: 1024,
- bounds: None,
- projection: None,
- }
- }
-}
-
-impl ReaderBuilder {
- /// Create a new builder for configuring CSV parsing options.
- ///
- /// To convert a builder into a reader, call `ReaderBuilder::build`
- ///
- /// # Example
- ///
- /// ```
- /// extern crate arrow;
- ///
- /// use arrow::csv;
- /// use std::fs::File;
- ///
- /// fn example() -> csv::Reader<File> {
- /// let file = File::open("test/data/uk_cities_with_headers.csv").unwrap();
- ///
- /// // create a builder, inferring the schema with the first 100 records
- /// let builder = csv::ReaderBuilder::new().infer_schema(Some(100));
- ///
- /// let reader = builder.build(file).unwrap();
- ///
- /// reader
- /// }
- /// ```
- pub fn new() -> ReaderBuilder {
- ReaderBuilder::default()
- }
-
- /// Set the CSV file's schema
- pub fn with_schema(mut self, schema: SchemaRef) -> Self {
- self.schema = Some(schema);
- self
- }
-
- /// Set whether the CSV file has headers
- pub fn has_header(mut self, has_header: bool) -> Self {
- self.has_header = has_header;
- self
- }
-
- /// Set the CSV file's column delimiter as a byte character
- pub fn with_delimiter(mut self, delimiter: u8) -> Self {
- self.delimiter = Some(delimiter);
- self
- }
-
- pub fn with_escape(mut self, escape: u8) -> Self {
- self.escape = Some(escape);
- self
- }
-
- pub fn with_quote(mut self, quote: u8) -> Self {
- self.quote = Some(quote);
- self
- }
-
- pub fn with_terminator(mut self, terminator: u8) -> Self {
- self.terminator = Some(terminator);
- self
- }
-
- /// Set the CSV reader to infer the schema of the file
- pub fn infer_schema(mut self, max_records: Option<usize>) -> Self {
- // remove any schema that is set
- self.schema = None;
- self.max_records = max_records;
- self
- }
-
- /// Set the batch size (number of records to load at one time)
- pub fn with_batch_size(mut self, batch_size: usize) -> Self {
- self.batch_size = batch_size;
- self
- }
-
- /// Set the reader's column projection
- pub fn with_projection(mut self, projection: Vec<usize>) -> Self {
- self.projection = Some(projection);
- self
- }
-
- /// Create a new `Reader` from the `ReaderBuilder`
- pub fn build<R: Read + Seek>(self, mut reader: R) -> Result<Reader<R>> {
- // check if schema should be inferred
- let delimiter = self.delimiter.unwrap_or(b',');
- let schema = match self.schema {
- Some(schema) => schema,
- None => {
- let (inferred_schema, _) = infer_file_schema_with_csv_options(
- &mut reader,
- delimiter,
- self.max_records,
- self.has_header,
- self.escape,
- self.quote,
- self.terminator,
- )?;
-
- Arc::new(inferred_schema)
- }
- };
- let csv_reader = Reader::build_csv_reader(
- reader,
- self.has_header,
- self.delimiter,
- self.escape,
- self.quote,
- self.terminator,
- );
- Ok(Reader::from_csv_reader(
- csv_reader,
- schema,
- self.has_header,
- self.batch_size,
- None,
- self.projection.clone(),
- ))
- }
-}
-
-#[cfg(test)]
-mod tests {
- use super::*;
-
- use std::fs::File;
- use std::io::{Cursor, Write};
- use tempfile::NamedTempFile;
-
- use crate::array::*;
- use crate::compute::cast;
- use crate::datatypes::Field;
-
- #[test]
- fn test_csv() {
- let schema = Schema::new(vec![
- Field::new("city", DataType::Utf8, false),
- Field::new("lat", DataType::Float64, false),
- Field::new("lng", DataType::Float64, false),
- ]);
-
- let file = File::open("test/data/uk_cities.csv").unwrap();
-
- let mut csv = Reader::new(
- file,
- Arc::new(schema.clone()),
- false,
- None,
- 1024,
- None,
- None,
- );
- assert_eq!(Arc::new(schema), csv.schema());
- let batch = csv.next().unwrap().unwrap();
- assert_eq!(37, batch.num_rows());
- assert_eq!(3, batch.num_columns());
-
- // access data from a primitive array
- let lat = batch
- .column(1)
- .as_any()
- .downcast_ref::<Float64Array>()
- .unwrap();
- assert!(57.653484 - lat.value(0) < f64::EPSILON);
-
- // access data from a string array (ListArray<u8>)
- let city = batch
- .column(0)
- .as_any()
- .downcast_ref::<StringArray>()
- .unwrap();
-
- assert_eq!("Aberdeen, Aberdeen City, UK", city.value(13));
- }
-
- #[test]
- fn test_csv_schema_metadata() {
- let mut metadata = std::collections::HashMap::new();
- metadata.insert("foo".to_owned(), "bar".to_owned());
- let schema = Schema::new_with_metadata(
- vec![
- Field::new("city", DataType::Utf8, false),
- Field::new("lat", DataType::Float64, false),
- Field::new("lng", DataType::Float64, false),
- ],
- metadata.clone(),
- );
-
- let file = File::open("test/data/uk_cities.csv").unwrap();
-
- let mut csv = Reader::new(
- file,
- Arc::new(schema.clone()),
- false,
- None,
- 1024,
- None,
- None,
- );
- assert_eq!(Arc::new(schema), csv.schema());
- let batch = csv.next().unwrap().unwrap();
- assert_eq!(37, batch.num_rows());
- assert_eq!(3, batch.num_columns());
-
- assert_eq!(&metadata, batch.schema().metadata());
- }
-
- #[test]
- fn test_csv_from_buf_reader() {
- let schema = Schema::new(vec![
- Field::new("city", DataType::Utf8, false),
- Field::new("lat", DataType::Float64, false),
- Field::new("lng", DataType::Float64, false),
- ]);
-
- let file_with_headers =
- File::open("test/data/uk_cities_with_headers.csv").unwrap();
- let file_without_headers = File::open("test/data/uk_cities.csv").unwrap();
- let both_files = file_with_headers
- .chain(Cursor::new("\n".to_string()))
- .chain(file_without_headers);
- let mut csv = Reader::from_reader(
- both_files,
- Arc::new(schema),
- true,
- None,
- 1024,
- None,
- None,
- );
- let batch = csv.next().unwrap().unwrap();
- assert_eq!(74, batch.num_rows());
- assert_eq!(3, batch.num_columns());
- }
-
- #[test]
- fn test_csv_with_schema_inference() {
- let file = File::open("test/data/uk_cities_with_headers.csv").unwrap();
-
- let builder = ReaderBuilder::new().has_header(true).infer_schema(None);
-
- let mut csv = builder.build(file).unwrap();
- let expected_schema = Schema::new(vec![
- Field::new("city", DataType::Utf8, false),
- Field::new("lat", DataType::Float64, false),
- Field::new("lng", DataType::Float64, false),
- ]);
- assert_eq!(Arc::new(expected_schema), csv.schema());
- let batch = csv.next().unwrap().unwrap();
- assert_eq!(37, batch.num_rows());
- assert_eq!(3, batch.num_columns());
-
- // access data from a primitive array
- let lat = batch
- .column(1)
- .as_any()
- .downcast_ref::<Float64Array>()
- .unwrap();
- assert!(57.653484 - lat.value(0) < f64::EPSILON);
-
- // access data from a string array (ListArray<u8>)
- let city = batch
- .column(0)
- .as_any()
- .downcast_ref::<StringArray>()
- .unwrap();
-
- assert_eq!("Aberdeen, Aberdeen City, UK", city.value(13));
- }
-
- #[test]
- fn test_csv_with_schema_inference_no_headers() {
- let file = File::open("test/data/uk_cities.csv").unwrap();
-
- let builder = ReaderBuilder::new().infer_schema(None);
-
- let mut csv = builder.build(file).unwrap();
-
- // csv field names should be 'column_{number}'
- let schema = csv.schema();
- assert_eq!("column_1", schema.field(0).name());
- assert_eq!("column_2", schema.field(1).name());
- assert_eq!("column_3", schema.field(2).name());
- let batch = csv.next().unwrap().unwrap();
- let batch_schema = batch.schema();
-
- assert_eq!(schema, batch_schema);
- assert_eq!(37, batch.num_rows());
- assert_eq!(3, batch.num_columns());
-
- // access data from a primitive array
- let lat = batch
- .column(1)
- .as_any()
- .downcast_ref::<Float64Array>()
- .unwrap();
- assert!(57.653484 - lat.value(0) < f64::EPSILON);
-
- // access data from a string array (ListArray<u8>)
- let city = batch
- .column(0)
- .as_any()
- .downcast_ref::<StringArray>()
- .unwrap();
-
- assert_eq!("Aberdeen, Aberdeen City, UK", city.value(13));
- }
-
- #[test]
- fn test_csv_with_projection() {
- let schema = Schema::new(vec![
- Field::new("city", DataType::Utf8, false),
- Field::new("lat", DataType::Float64, false),
- Field::new("lng", DataType::Float64, false),
- ]);
-
- let file = File::open("test/data/uk_cities.csv").unwrap();
-
- let mut csv = Reader::new(
- file,
- Arc::new(schema),
- false,
- None,
- 1024,
- None,
- Some(vec![0, 1]),
- );
- let projected_schema = Arc::new(Schema::new(vec![
- Field::new("city", DataType::Utf8, false),
- Field::new("lat", DataType::Float64, false),
- ]));
- assert_eq!(projected_schema, csv.schema());
- let batch = csv.next().unwrap().unwrap();
- assert_eq!(projected_schema, batch.schema());
- assert_eq!(37, batch.num_rows());
- assert_eq!(2, batch.num_columns());
- }
-
- #[test]
- fn test_csv_with_dictionary() {
- let schema = Schema::new(vec![
- Field::new(
- "city",
- DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
- false,
- ),
- Field::new("lat", DataType::Float64, false),
- Field::new("lng", DataType::Float64, false),
- ]);
-
- let file = File::open("test/data/uk_cities.csv").unwrap();
-
- let mut csv = Reader::new(
- file,
- Arc::new(schema),
- false,
- None,
- 1024,
- None,
- Some(vec![0, 1]),
- );
- let projected_schema = Arc::new(Schema::new(vec![
- Field::new(
- "city",
- DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
- false,
- ),
- Field::new("lat", DataType::Float64, false),
- ]));
- assert_eq!(projected_schema, csv.schema());
- let batch = csv.next().unwrap().unwrap();
- assert_eq!(projected_schema, batch.schema());
- assert_eq!(37, batch.num_rows());
- assert_eq!(2, batch.num_columns());
-
- let strings = cast(batch.column(0), &DataType::Utf8).unwrap();
- let strings = strings.as_any().downcast_ref::<StringArray>().unwrap();
-
- assert_eq!(strings.value(0), "Elgin, Scotland, the UK");
- assert_eq!(strings.value(4), "Eastbourne, East Sussex, UK");
- assert_eq!(strings.value(29), "Uckfield, East Sussex, UK");
- }
-
- #[test]
- fn test_nulls() {
- let schema = Schema::new(vec![
- Field::new("c_int", DataType::UInt64, false),
- Field::new("c_float", DataType::Float32, false),
- Field::new("c_string", DataType::Utf8, false),
- ]);
-
- let file = File::open("test/data/null_test.csv").unwrap();
-
- let mut csv = Reader::new(file, Arc::new(schema), true, None, 1024, None, None);
- let batch = csv.next().unwrap().unwrap();
-
- assert_eq!(false, batch.column(1).is_null(0));
- assert_eq!(false, batch.column(1).is_null(1));
- assert_eq!(true, batch.column(1).is_null(2));
- assert_eq!(false, batch.column(1).is_null(3));
- assert_eq!(false, batch.column(1).is_null(4));
- }
-
- #[test]
- fn test_nulls_with_inference() {
- let file = File::open("test/data/various_types.csv").unwrap();
-
- let builder = ReaderBuilder::new()
- .infer_schema(None)
- .has_header(true)
- .with_delimiter(b'|')
- .with_batch_size(512)
- .with_projection(vec![0, 1, 2, 3, 4, 5]);
-
- let mut csv = builder.build(file).unwrap();
- let batch = csv.next().unwrap().unwrap();
-
- assert_eq!(5, batch.num_rows());
- assert_eq!(6, batch.num_columns());
-
- let schema = batch.schema();
-
- assert_eq!(&DataType::Int64, schema.field(0).data_type());
- assert_eq!(&DataType::Float64, schema.field(1).data_type());
- assert_eq!(&DataType::Float64, schema.field(2).data_type());
- assert_eq!(&DataType::Boolean, schema.field(3).data_type());
- assert_eq!(&DataType::Date32, schema.field(4).data_type());
- assert_eq!(&DataType::Date64, schema.field(5).data_type());
-
- let names: Vec<&str> =
- schema.fields().iter().map(|x| x.name().as_str()).collect();
- assert_eq!(
- names,
- vec![
- "c_int",
- "c_float",
- "c_string",
- "c_bool",
- "c_date",
- "c_datetime"
- ]
- );
-
- assert_eq!(false, schema.field(0).is_nullable());
- assert_eq!(true, schema.field(1).is_nullable());
- assert_eq!(true, schema.field(2).is_nullable());
- assert_eq!(false, schema.field(3).is_nullable());
- assert_eq!(true, schema.field(4).is_nullable());
- assert_eq!(true, schema.field(5).is_nullable());
-
- assert_eq!(false, batch.column(1).is_null(0));
- assert_eq!(false, batch.column(1).is_null(1));
- assert_eq!(true, batch.column(1).is_null(2));
- assert_eq!(false, batch.column(1).is_null(3));
- assert_eq!(false, batch.column(1).is_null(4));
- }
-
- #[test]
- fn test_parse_invalid_csv() {
- let file = File::open("test/data/various_types_invalid.csv").unwrap();
-
- let schema = Schema::new(vec![
- Field::new("c_int", DataType::UInt64, false),
- Field::new("c_float", DataType::Float32, false),
- Field::new("c_string", DataType::Utf8, false),
- Field::new("c_bool", DataType::Boolean, false),
- ]);
-
- let builder = ReaderBuilder::new()
- .with_schema(Arc::new(schema))
- .has_header(true)
- .with_delimiter(b'|')
- .with_batch_size(512)
- .with_projection(vec![0, 1, 2, 3]);
-
- let mut csv = builder.build(file).unwrap();
- match csv.next() {
- Some(e) => match e {
- Err(e) => assert_eq!(
- "ParseError(\"Error while parsing value 4.x4 for column 1 at line 4\")",
- format!("{:?}", e)
- ),
- Ok(_) => panic!("should have failed"),
- },
- None => panic!("should have failed"),
- }
- }
-
- #[test]
- fn test_infer_field_schema() {
- assert_eq!(infer_field_schema("A"), DataType::Utf8);
- assert_eq!(infer_field_schema("\"123\""), DataType::Utf8);
- assert_eq!(infer_field_schema("10"), DataType::Int64);
- assert_eq!(infer_field_schema("10.2"), DataType::Float64);
- assert_eq!(infer_field_schema("true"), DataType::Boolean);
- assert_eq!(infer_field_schema("false"), DataType::Boolean);
- assert_eq!(infer_field_schema("2020-11-08"), DataType::Date32);
- assert_eq!(infer_field_schema("2020-11-08T14:20:01"), DataType::Date64);
- }
-
- #[test]
- fn parse_date32() {
- assert_eq!(parse_item::<Date32Type>("1970-01-01").unwrap(), 0);
- assert_eq!(parse_item::<Date32Type>("2020-03-15").unwrap(), 18336);
- assert_eq!(parse_item::<Date32Type>("1945-05-08").unwrap(), -9004);
- }
-
- #[test]
- fn parse_date64() {
- assert_eq!(parse_item::<Date64Type>("1970-01-01T00:00:00").unwrap(), 0);
- assert_eq!(
- parse_item::<Date64Type>("2018-11-13T17:11:10").unwrap(),
- 1542129070000
- );
- assert_eq!(
- parse_item::<Date64Type>("2018-11-13T17:11:10.011").unwrap(),
- 1542129070011
- );
- assert_eq!(
- parse_item::<Date64Type>("1900-02-28T12:34:56").unwrap(),
- -2203932304000
- );
- }
-
- #[test]
- fn test_infer_schema_from_multiple_files() -> Result<()> {
- let mut csv1 = NamedTempFile::new()?;
- let mut csv2 = NamedTempFile::new()?;
- let csv3 = NamedTempFile::new()?; // empty csv file should be skipped
- let mut csv4 = NamedTempFile::new()?;
- writeln!(csv1, "c1,c2,c3")?;
- writeln!(csv1, "1,\"foo\",0.5")?;
- writeln!(csv1, "3,\"bar\",1")?;
- // reading csv2 will set c2 to optional
- writeln!(csv2, "c1,c2,c3,c4")?;
- writeln!(csv2, "10,,3.14,true")?;
- // reading csv4 will set c3 to optional
- writeln!(csv4, "c1,c2,c3")?;
- writeln!(csv4, "10,\"foo\",")?;
-
- let schema = infer_schema_from_files(
- &[
- csv3.path().to_str().unwrap().to_string(),
- csv1.path().to_str().unwrap().to_string(),
- csv2.path().to_str().unwrap().to_string(),
- csv4.path().to_str().unwrap().to_string(),
- ],
- b',',
- Some(3), // only csv1 and csv2 should be read
- true,
- )?;
-
- assert_eq!(schema.fields().len(), 4);
- assert_eq!(false, schema.field(0).is_nullable());
- assert_eq!(true, schema.field(1).is_nullable());
- assert_eq!(false, schema.field(2).is_nullable());
- assert_eq!(false, schema.field(3).is_nullable());
-
- assert_eq!(&DataType::Int64, schema.field(0).data_type());
- assert_eq!(&DataType::Utf8, schema.field(1).data_type());
- assert_eq!(&DataType::Float64, schema.field(2).data_type());
- assert_eq!(&DataType::Boolean, schema.field(3).data_type());
-
- Ok(())
- }
-
- #[test]
- fn test_bounded() {
- let schema = Schema::new(vec![Field::new("int", DataType::UInt32, false)]);
- let data = vec![
- vec!["0"],
- vec!["1"],
- vec!["2"],
- vec!["3"],
- vec!["4"],
- vec!["5"],
- vec!["6"],
- ];
-
- let data = data
- .iter()
- .map(|x| x.join(","))
- .collect::<Vec<_>>()
- .join("\n");
- let data = data.as_bytes();
-
- let reader = std::io::Cursor::new(data);
-
- let mut csv = Reader::new(
- reader,
- Arc::new(schema),
- false,
- None,
- 2,
- // starting at row 2 and up to row 6.
- Some((2, 6)),
- Some(vec![0]),
- );
-
- let batch = csv.next().unwrap().unwrap();
- let a = batch.column(0);
- let a = a.as_any().downcast_ref::<UInt32Array>().unwrap();
- assert_eq!(a, &UInt32Array::from(vec![2, 3]));
-
- let batch = csv.next().unwrap().unwrap();
- let a = batch.column(0);
- let a = a.as_any().downcast_ref::<UInt32Array>().unwrap();
- assert_eq!(a, &UInt32Array::from(vec![4, 5]));
-
- assert!(csv.next().is_none());
- }
-
- #[test]
- fn test_parsing_bool() {
- // Encode the expected behavior of boolean parsing
- assert_eq!(Some(true), parse_bool("true"));
- assert_eq!(Some(true), parse_bool("tRUe"));
- assert_eq!(Some(true), parse_bool("True"));
- assert_eq!(Some(true), parse_bool("TRUE"));
- assert_eq!(None, parse_bool("t"));
- assert_eq!(None, parse_bool("T"));
- assert_eq!(None, parse_bool(""));
-
- assert_eq!(Some(false), parse_bool("false"));
- assert_eq!(Some(false), parse_bool("fALse"));
- assert_eq!(Some(false), parse_bool("False"));
- assert_eq!(Some(false), parse_bool("FALSE"));
- assert_eq!(None, parse_bool("f"));
- assert_eq!(None, parse_bool("F"));
- assert_eq!(None, parse_bool(""));
- }
-
- #[test]
- fn test_parsing_float() {
- assert_eq!(Some(12.34), parse_item::<Float64Type>("12.34"));
- assert_eq!(Some(-12.34), parse_item::<Float64Type>("-12.34"));
- assert_eq!(Some(12.0), parse_item::<Float64Type>("12"));
- assert_eq!(Some(0.0), parse_item::<Float64Type>("0"));
- assert!(parse_item::<Float64Type>("nan").unwrap().is_nan());
- assert!(parse_item::<Float64Type>("NaN").unwrap().is_nan());
- assert!(parse_item::<Float64Type>("inf").unwrap().is_infinite());
- assert!(parse_item::<Float64Type>("inf").unwrap().is_sign_positive());
- assert!(parse_item::<Float64Type>("-inf").unwrap().is_infinite());
- assert!(parse_item::<Float64Type>("-inf")
- .unwrap()
- .is_sign_negative());
- assert_eq!(None, parse_item::<Float64Type>(""));
- assert_eq!(None, parse_item::<Float64Type>("dd"));
- assert_eq!(None, parse_item::<Float64Type>("12.34.56"));
- }
-
- #[test]
- fn test_non_std_quote() {
- let schema = Schema::new(vec![
- Field::new("text1", DataType::Utf8, false),
- Field::new("text2", DataType::Utf8, false),
- ]);
- let builder = ReaderBuilder::new()
- .with_schema(Arc::new(schema))
- .has_header(false)
- .with_quote(b'~'); // default is ", change to ~
-
- let mut csv_text = Vec::new();
- let mut csv_writer = std::io::Cursor::new(&mut csv_text);
- for index in 0..10 {
- let text1 = format!("id{:}", index);
- let text2 = format!("value{:}", index);
- csv_writer
- .write_fmt(format_args!("~{}~,~{}~\r\n", text1, text2))
- .unwrap();
- }
- let mut csv_reader = std::io::Cursor::new(&csv_text);
- let mut reader = builder.build(&mut csv_reader).unwrap();
- let batch = reader.next().unwrap().unwrap();
- let col0 = batch.column(0);
- assert_eq!(col0.len(), 10);
- let col0_arr = col0.as_any().downcast_ref::<StringArray>().unwrap();
- assert_eq!(col0_arr.value(0), "id0");
- let col1 = batch.column(1);
- assert_eq!(col1.len(), 10);
- let col1_arr = col1.as_any().downcast_ref::<StringArray>().unwrap();
- assert_eq!(col1_arr.value(5), "value5");
- }
-
- #[test]
- fn test_non_std_escape() {
- let schema = Schema::new(vec![
- Field::new("text1", DataType::Utf8, false),
- Field::new("text2", DataType::Utf8, false),
- ]);
- let builder = ReaderBuilder::new()
- .with_schema(Arc::new(schema))
- .has_header(false)
- .with_escape(b'\\'); // default is None, change to \
-
- let mut csv_text = Vec::new();
- let mut csv_writer = std::io::Cursor::new(&mut csv_text);
- for index in 0..10 {
- let text1 = format!("id{:}", index);
- let text2 = format!("value\\\"{:}", index);
- csv_writer
- .write_fmt(format_args!("\"{}\",\"{}\"\r\n", text1, text2))
- .unwrap();
- }
- let mut csv_reader = std::io::Cursor::new(&csv_text);
- let mut reader = builder.build(&mut csv_reader).unwrap();
- let batch = reader.next().unwrap().unwrap();
- let col0 = batch.column(0);
- assert_eq!(col0.len(), 10);
- let col0_arr = col0.as_any().downcast_ref::<StringArray>().unwrap();
- assert_eq!(col0_arr.value(0), "id0");
- let col1 = batch.column(1);
- assert_eq!(col1.len(), 10);
- let col1_arr = col1.as_any().downcast_ref::<StringArray>().unwrap();
- assert_eq!(col1_arr.value(5), "value\"5");
- }
-
- #[test]
- fn test_non_std_terminator() {
- let schema = Schema::new(vec![
- Field::new("text1", DataType::Utf8, false),
- Field::new("text2", DataType::Utf8, false),
- ]);
- let builder = ReaderBuilder::new()
- .with_schema(Arc::new(schema))
- .has_header(false)
- .with_terminator(b'\n'); // default is CRLF, change to LF
-
- let mut csv_text = Vec::new();
- let mut csv_writer = std::io::Cursor::new(&mut csv_text);
- for index in 0..10 {
- let text1 = format!("id{:}", index);
- let text2 = format!("value{:}", index);
- csv_writer
- .write_fmt(format_args!("\"{}\",\"{}\"\n", text1, text2))
- .unwrap();
- }
- let mut csv_reader = std::io::Cursor::new(&csv_text);
- let mut reader = builder.build(&mut csv_reader).unwrap();
- let batch = reader.next().unwrap().unwrap();
- let col0 = batch.column(0);
- assert_eq!(col0.len(), 10);
- let col0_arr = col0.as_any().downcast_ref::<StringArray>().unwrap();
- assert_eq!(col0_arr.value(0), "id0");
- let col1 = batch.column(1);
- assert_eq!(col1.len(), 10);
- let col1_arr = col1.as_any().downcast_ref::<StringArray>().unwrap();
- assert_eq!(col1_arr.value(5), "value5");
- }
-}
diff --git a/arrow/src/csv/writer.rs b/arrow/src/csv/writer.rs
deleted file mode 100644
index b3b8838..0000000
--- a/arrow/src/csv/writer.rs
+++ /dev/null
@@ -1,678 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-//! CSV Writer
-//!
-//! This CSV writer allows Arrow data (in record batches) to be written as CSV files.
-//! The writer does not support writing `ListArray` and `StructArray`.
-//!
-//! Example:
-//!
-//! ```
-//! use arrow::array::*;
-//! use arrow::csv;
-//! use arrow::datatypes::*;
-//! use arrow::record_batch::RecordBatch;
-//! use arrow::util::test_util::get_temp_file;
-//! use std::fs::File;
-//! use std::sync::Arc;
-//!
-//! let schema = Schema::new(vec![
-//! Field::new("c1", DataType::Utf8, false),
-//! Field::new("c2", DataType::Float64, true),
-//! Field::new("c3", DataType::UInt32, false),
-//! Field::new("c3", DataType::Boolean, true),
-//! ]);
-//! let c1 = StringArray::from(vec![
-//! "Lorem ipsum dolor sit amet",
-//! "consectetur adipiscing elit",
-//! "sed do eiusmod tempor",
-//! ]);
-//! let c2 = PrimitiveArray::<Float64Type>::from(vec![
-//! Some(123.564532),
-//! None,
-//! Some(-556132.25),
-//! ]);
-//! let c3 = PrimitiveArray::<UInt32Type>::from(vec![3, 2, 1]);
-//! let c4 = BooleanArray::from(vec![Some(true), Some(false), None]);
-//!
-//! let batch = RecordBatch::try_new(
-//! Arc::new(schema),
-//! vec![Arc::new(c1), Arc::new(c2), Arc::new(c3), Arc::new(c4)],
-//! )
-//! .unwrap();
-//!
-//! let file = get_temp_file("out.csv", &[]);
-//!
-//! let mut writer = csv::Writer::new(file);
-//! let batches = vec![&batch, &batch];
-//! for batch in batches {
-//! writer.write(batch).unwrap();
-//! }
-//! ```
-
-use std::io::Write;
-
-use crate::datatypes::*;
-use crate::error::{ArrowError, Result};
-use crate::record_batch::RecordBatch;
-use crate::util::display::make_string_from_decimal;
-use crate::{array::*, util::serialization::lexical_to_string};
-const DEFAULT_DATE_FORMAT: &str = "%F";
-const DEFAULT_TIME_FORMAT: &str = "%T";
-const DEFAULT_TIMESTAMP_FORMAT: &str = "%FT%H:%M:%S.%9f";
-
-fn write_primitive_value<T>(array: &ArrayRef, i: usize) -> String
-where
- T: ArrowNumericType,
- T::Native: lexical_core::ToLexical,
-{
- let c = array.as_any().downcast_ref::<PrimitiveArray<T>>().unwrap();
- lexical_to_string(c.value(i))
-}
-
-/// A CSV writer
-#[derive(Debug)]
-pub struct Writer<W: Write> {
- /// The object to write to
- writer: csv_crate::Writer<W>,
- /// Column delimiter. Defaults to `b','`
- delimiter: u8,
- /// Whether file should be written with headers. Defaults to `true`
- has_headers: bool,
- /// The date format for date arrays
- date_format: String,
- /// The datetime format for datetime arrays
- datetime_format: String,
- /// The timestamp format for timestamp arrays
- timestamp_format: String,
- /// The time format for time arrays
- time_format: String,
- /// Is the beginning-of-writer
- beginning: bool,
-}
-
-impl<W: Write> Writer<W> {
- /// Create a new CsvWriter from a writable object, with default options
- pub fn new(writer: W) -> Self {
- let delimiter = b',';
- let mut builder = csv_crate::WriterBuilder::new();
- let writer = builder.delimiter(delimiter).from_writer(writer);
- Writer {
- writer,
- delimiter,
- has_headers: true,
- date_format: DEFAULT_DATE_FORMAT.to_string(),
- datetime_format: DEFAULT_TIMESTAMP_FORMAT.to_string(),
- time_format: DEFAULT_TIME_FORMAT.to_string(),
- timestamp_format: DEFAULT_TIMESTAMP_FORMAT.to_string(),
- beginning: true,
- }
- }
-
- /// Convert a record to a string vector
- fn convert(
- &self,
- batch: &[ArrayRef],
- row_index: usize,
- buffer: &mut [String],
- ) -> Result<()> {
- // TODO: it'd be more efficient if we could create `record: Vec<&[u8]>
- for (col_index, item) in buffer.iter_mut().enumerate() {
- let col = &batch[col_index];
- if col.is_null(row_index) {
- // write an empty value
- *item = "".to_string();
- continue;
- }
- let string = match col.data_type() {
- DataType::Float64 => write_primitive_value::<Float64Type>(col, row_index),
- DataType::Float32 => write_primitive_value::<Float32Type>(col, row_index),
- DataType::Int8 => write_primitive_value::<Int8Type>(col, row_index),
- DataType::Int16 => write_primitive_value::<Int16Type>(col, row_index),
- DataType::Int32 => write_primitive_value::<Int32Type>(col, row_index),
- DataType::Int64 => write_primitive_value::<Int64Type>(col, row_index),
- DataType::UInt8 => write_primitive_value::<UInt8Type>(col, row_index),
- DataType::UInt16 => write_primitive_value::<UInt16Type>(col, row_index),
- DataType::UInt32 => write_primitive_value::<UInt32Type>(col, row_index),
- DataType::UInt64 => write_primitive_value::<UInt64Type>(col, row_index),
- DataType::Boolean => {
- let c = col.as_any().downcast_ref::<BooleanArray>().unwrap();
- c.value(row_index).to_string()
- }
- DataType::Utf8 => {
- let c = col.as_any().downcast_ref::<StringArray>().unwrap();
- c.value(row_index).to_owned()
- }
- DataType::LargeUtf8 => {
- let c = col.as_any().downcast_ref::<LargeStringArray>().unwrap();
- c.value(row_index).to_owned()
- }
- DataType::Date32 => {
- let c = col.as_any().downcast_ref::<Date32Array>().unwrap();
- c.value_as_date(row_index)
- .unwrap()
- .format(&self.date_format)
- .to_string()
- }
- DataType::Date64 => {
- let c = col.as_any().downcast_ref::<Date64Array>().unwrap();
- c.value_as_datetime(row_index)
- .unwrap()
- .format(&self.datetime_format)
- .to_string()
- }
- DataType::Time32(TimeUnit::Second) => {
- let c = col.as_any().downcast_ref::<Time32SecondArray>().unwrap();
- c.value_as_time(row_index)
- .unwrap()
- .format(&self.time_format)
- .to_string()
- }
- DataType::Time32(TimeUnit::Millisecond) => {
- let c = col
- .as_any()
- .downcast_ref::<Time32MillisecondArray>()
- .unwrap();
- c.value_as_time(row_index)
- .unwrap()
- .format(&self.time_format)
- .to_string()
- }
- DataType::Time64(TimeUnit::Microsecond) => {
- let c = col
- .as_any()
- .downcast_ref::<Time64MicrosecondArray>()
- .unwrap();
- c.value_as_time(row_index)
- .unwrap()
- .format(&self.time_format)
- .to_string()
- }
- DataType::Time64(TimeUnit::Nanosecond) => {
- let c = col
- .as_any()
- .downcast_ref::<Time64NanosecondArray>()
- .unwrap();
- c.value_as_time(row_index)
- .unwrap()
- .format(&self.time_format)
- .to_string()
- }
- DataType::Timestamp(time_unit, _) => {
- use TimeUnit::*;
- let datetime = match time_unit {
- Second => col
- .as_any()
- .downcast_ref::<TimestampSecondArray>()
- .unwrap()
- .value_as_datetime(row_index)
- .unwrap(),
- Millisecond => col
- .as_any()
- .downcast_ref::<TimestampMillisecondArray>()
- .unwrap()
- .value_as_datetime(row_index)
- .unwrap(),
- Microsecond => col
- .as_any()
- .downcast_ref::<TimestampMicrosecondArray>()
- .unwrap()
- .value_as_datetime(row_index)
- .unwrap(),
- Nanosecond => col
- .as_any()
- .downcast_ref::<TimestampNanosecondArray>()
- .unwrap()
- .value_as_datetime(row_index)
- .unwrap(),
- };
- format!("{}", datetime.format(&self.timestamp_format))
- }
- DataType::Decimal(..) => make_string_from_decimal(col, row_index)?,
- t => {
- // List and Struct arrays not supported by the writer, any
- // other type needs to be implemented
- return Err(ArrowError::CsvError(format!(
- "CSV Writer does not support {:?} data type",
- t
- )));
- }
- };
- *item = string;
- }
- Ok(())
- }
-
- /// Write a vector of record batches to a writable object
- pub fn write(&mut self, batch: &RecordBatch) -> Result<()> {
- let num_columns = batch.num_columns();
- if self.beginning {
- if self.has_headers {
- let mut headers: Vec<String> = Vec::with_capacity(num_columns);
- batch
- .schema()
- .fields()
- .iter()
- .for_each(|field| headers.push(field.name().to_string()));
- self.writer.write_record(&headers[..])?;
- }
- self.beginning = false;
- }
-
- let columns: Vec<_> = batch
- .columns()
- .iter()
- .map(|array| match array.data_type() {
- DataType::Dictionary(_, value_type) => {
- crate::compute::kernels::cast::cast(array, &value_type)
- .expect("cannot cast dictionary to underlying values")
- }
- _ => array.clone(),
- })
- .collect();
-
- let mut buffer = vec!["".to_string(); batch.num_columns()];
-
- for row_index in 0..batch.num_rows() {
- self.convert(columns.as_slice(), row_index, &mut buffer)?;
- self.writer.write_record(&buffer)?;
- }
- self.writer.flush()?;
-
- Ok(())
- }
-}
-
-/// A CSV writer builder
-#[derive(Debug)]
-pub struct WriterBuilder {
- /// Optional column delimiter. Defaults to `b','`
- delimiter: Option<u8>,
- /// Whether to write column names as file headers. Defaults to `true`
- has_headers: bool,
- /// Optional date format for date arrays
- date_format: Option<String>,
- /// Optional datetime format for datetime arrays
- datetime_format: Option<String>,
- /// Optional timestamp format for timestamp arrays
- timestamp_format: Option<String>,
- /// Optional time format for time arrays
- time_format: Option<String>,
-}
-
-impl Default for WriterBuilder {
- fn default() -> Self {
- Self {
- has_headers: true,
- delimiter: None,
- date_format: Some(DEFAULT_DATE_FORMAT.to_string()),
- datetime_format: Some(DEFAULT_TIMESTAMP_FORMAT.to_string()),
- time_format: Some(DEFAULT_TIME_FORMAT.to_string()),
- timestamp_format: Some(DEFAULT_TIMESTAMP_FORMAT.to_string()),
- }
- }
-}
-
-impl WriterBuilder {
- /// Create a new builder for configuring CSV writing options.
- ///
- /// To convert a builder into a writer, call `WriterBuilder::build`
- ///
- /// # Example
- ///
- /// ```
- /// extern crate arrow;
- ///
- /// use arrow::csv;
- /// use std::fs::File;
- ///
- /// fn example() -> csv::Writer<File> {
- /// let file = File::create("target/out.csv").unwrap();
- ///
- /// // create a builder that doesn't write headers
- /// let builder = csv::WriterBuilder::new().has_headers(false);
- /// let writer = builder.build(file);
- ///
- /// writer
- /// }
- /// ```
- pub fn new() -> Self {
- Self::default()
- }
-
- /// Set whether to write headers
- pub fn has_headers(mut self, has_headers: bool) -> Self {
- self.has_headers = has_headers;
- self
- }
-
- /// Set the CSV file's column delimiter as a byte character
- pub fn with_delimiter(mut self, delimiter: u8) -> Self {
- self.delimiter = Some(delimiter);
- self
- }
-
- /// Set the CSV file's date format
- pub fn with_date_format(mut self, format: String) -> Self {
- self.date_format = Some(format);
- self
- }
-
- /// Set the CSV file's time format
- pub fn with_time_format(mut self, format: String) -> Self {
- self.time_format = Some(format);
- self
- }
-
- /// Set the CSV file's timestamp format
- pub fn with_timestamp_format(mut self, format: String) -> Self {
- self.timestamp_format = Some(format);
- self
- }
-
- /// Create a new `Writer`
- pub fn build<W: Write>(self, writer: W) -> Writer<W> {
- let delimiter = self.delimiter.unwrap_or(b',');
- let mut builder = csv_crate::WriterBuilder::new();
- let writer = builder.delimiter(delimiter).from_writer(writer);
- Writer {
- writer,
- delimiter,
- has_headers: self.has_headers,
- date_format: self
- .date_format
- .unwrap_or_else(|| DEFAULT_DATE_FORMAT.to_string()),
- datetime_format: self
- .datetime_format
- .unwrap_or_else(|| DEFAULT_TIMESTAMP_FORMAT.to_string()),
- time_format: self
- .time_format
- .unwrap_or_else(|| DEFAULT_TIME_FORMAT.to_string()),
- timestamp_format: self
- .timestamp_format
- .unwrap_or_else(|| DEFAULT_TIMESTAMP_FORMAT.to_string()),
- beginning: true,
- }
- }
-}
-
-#[cfg(test)]
-mod tests {
- use super::*;
-
- use crate::csv::Reader;
- use crate::datatypes::{Field, Schema};
- use crate::util::string_writer::StringWriter;
- use crate::util::test_util::get_temp_file;
- use std::fs::File;
- use std::io::{Cursor, Read};
- use std::sync::Arc;
-
- #[test]
- fn test_write_csv() {
- let schema = Schema::new(vec![
- Field::new("c1", DataType::Utf8, false),
- Field::new("c2", DataType::Float64, true),
- Field::new("c3", DataType::UInt32, false),
- Field::new("c4", DataType::Boolean, true),
- Field::new("c5", DataType::Timestamp(TimeUnit::Millisecond, None), true),
- Field::new("c6", DataType::Time32(TimeUnit::Second), false),
- Field::new(
- "c7",
- DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
- false,
- ),
- ]);
-
- let c1 = StringArray::from(vec![
- "Lorem ipsum dolor sit amet",
- "consectetur adipiscing elit",
- "sed do eiusmod tempor",
- ]);
- let c2 = PrimitiveArray::<Float64Type>::from(vec![
- Some(123.564532),
- None,
- Some(-556132.25),
- ]);
- let c3 = PrimitiveArray::<UInt32Type>::from(vec![3, 2, 1]);
- let c4 = BooleanArray::from(vec![Some(true), Some(false), None]);
- let c5 = TimestampMillisecondArray::from_opt_vec(
- vec![None, Some(1555584887378), Some(1555555555555)],
- None,
- );
- let c6 = Time32SecondArray::from(vec![1234, 24680, 85563]);
- let c7: DictionaryArray<Int32Type> =
- vec!["cupcakes", "cupcakes", "foo"].into_iter().collect();
-
- let batch = RecordBatch::try_new(
- Arc::new(schema),
- vec![
- Arc::new(c1),
- Arc::new(c2),
- Arc::new(c3),
- Arc::new(c4),
- Arc::new(c5),
- Arc::new(c6),
- Arc::new(c7),
- ],
- )
- .unwrap();
-
- let file = get_temp_file("columns.csv", &[]);
-
- let mut writer = Writer::new(file);
- let batches = vec![&batch, &batch];
- for batch in batches {
- writer.write(batch).unwrap();
- }
- // check that file was written successfully
- let mut file = File::open("target/debug/testdata/columns.csv").unwrap();
- let mut buffer: Vec<u8> = vec![];
- file.read_to_end(&mut buffer).unwrap();
-
- assert_eq!(
- r#"c1,c2,c3,c4,c5,c6,c7
-Lorem ipsum dolor sit amet,123.564532,3,true,,00:20:34,cupcakes
-consectetur adipiscing elit,,2,false,2019-04-18T10:54:47.378000000,06:51:20,cupcakes
-sed do eiusmod tempor,-556132.25,1,,2019-04-18T02:45:55.555000000,23:46:03,foo
-Lorem ipsum dolor sit amet,123.564532,3,true,,00:20:34,cupcakes
-consectetur adipiscing elit,,2,false,2019-04-18T10:54:47.378000000,06:51:20,cupcakes
-sed do eiusmod tempor,-556132.25,1,,2019-04-18T02:45:55.555000000,23:46:03,foo
-"#
- .to_string(),
- String::from_utf8(buffer).unwrap()
- );
- }
-
- #[test]
- fn test_write_csv_custom_options() {
- let schema = Schema::new(vec![
- Field::new("c1", DataType::Utf8, false),
- Field::new("c2", DataType::Float64, true),
- Field::new("c3", DataType::UInt32, false),
- Field::new("c4", DataType::Boolean, true),
- Field::new("c6", DataType::Time32(TimeUnit::Second), false),
- ]);
-
- let c1 = StringArray::from(vec![
- "Lorem ipsum dolor sit amet",
- "consectetur adipiscing elit",
- "sed do eiusmod tempor",
- ]);
- let c2 = PrimitiveArray::<Float64Type>::from(vec![
- Some(123.564532),
- None,
- Some(-556132.25),
- ]);
- let c3 = PrimitiveArray::<UInt32Type>::from(vec![3, 2, 1]);
- let c4 = BooleanArray::from(vec![Some(true), Some(false), None]);
- let c6 = Time32SecondArray::from(vec![1234, 24680, 85563]);
-
- let batch = RecordBatch::try_new(
- Arc::new(schema),
- vec![
- Arc::new(c1),
- Arc::new(c2),
- Arc::new(c3),
- Arc::new(c4),
- Arc::new(c6),
- ],
- )
- .unwrap();
-
- let file = get_temp_file("custom_options.csv", &[]);
-
- let builder = WriterBuilder::new()
- .has_headers(false)
- .with_delimiter(b'|')
- .with_time_format("%r".to_string());
- let mut writer = builder.build(file);
- let batches = vec![&batch];
- for batch in batches {
- writer.write(batch).unwrap();
- }
-
- // check that file was written successfully
- let mut file = File::open("target/debug/testdata/custom_options.csv").unwrap();
- let mut buffer: Vec<u8> = vec![];
- file.read_to_end(&mut buffer).unwrap();
-
- assert_eq!(
- "Lorem ipsum dolor sit amet|123.564532|3|true|12:20:34 AM\nconsectetur adipiscing elit||2|false|06:51:20 AM\nsed do eiusmod tempor|-556132.25|1||11:46:03 PM\n"
- .to_string(),
- String::from_utf8(buffer).unwrap()
- );
- }
-
- #[test]
- fn test_export_csv_string() {
- let schema = Schema::new(vec![
- Field::new("c1", DataType::Utf8, false),
- Field::new("c2", DataType::Float64, true),
- Field::new("c3", DataType::UInt32, false),
- Field::new("c4", DataType::Boolean, true),
- Field::new("c5", DataType::Timestamp(TimeUnit::Millisecond, None), true),
- Field::new("c6", DataType::Time32(TimeUnit::Second), false),
- Field::new("c7", DataType::Decimal(6, 2), false),
- ]);
-
- let c1 = StringArray::from(vec![
- "Lorem ipsum dolor sit amet",
- "consectetur adipiscing elit",
- "sed do eiusmod tempor",
- ]);
- let c2 = PrimitiveArray::<Float64Type>::from(vec![
- Some(123.564532),
- None,
- Some(-556132.25),
- ]);
- let c3 = PrimitiveArray::<UInt32Type>::from(vec![3, 2, 1]);
- let c4 = BooleanArray::from(vec![Some(true), Some(false), None]);
- let c5 = TimestampMillisecondArray::from_opt_vec(
- vec![None, Some(1555584887378), Some(1555555555555)],
- None,
- );
- let c6 = Time32SecondArray::from(vec![1234, 24680, 85563]);
- let mut c7_builder = DecimalBuilder::new(5, 6, 2);
- c7_builder.append_value(12345_i128).unwrap();
- c7_builder.append_value(-12345_i128).unwrap();
- c7_builder.append_null().unwrap();
- let c7 = c7_builder.finish();
-
- let batch = RecordBatch::try_new(
- Arc::new(schema),
- vec![
- Arc::new(c1),
- Arc::new(c2),
- Arc::new(c3),
- Arc::new(c4),
- Arc::new(c5),
- Arc::new(c6),
- Arc::new(c7),
- ],
- )
- .unwrap();
-
- let sw = StringWriter::new();
- let mut writer = Writer::new(sw);
- let batches = vec![&batch, &batch];
- for batch in batches {
- writer.write(batch).unwrap();
- }
-
- let left = "c1,c2,c3,c4,c5,c6,c7
-Lorem ipsum dolor sit amet,123.564532,3,true,,00:20:34,123.45
-consectetur adipiscing elit,,2,false,2019-04-18T10:54:47.378000000,06:51:20,-123.45
-sed do eiusmod tempor,-556132.25,1,,2019-04-18T02:45:55.555000000,23:46:03,
-Lorem ipsum dolor sit amet,123.564532,3,true,,00:20:34,123.45
-consectetur adipiscing elit,,2,false,2019-04-18T10:54:47.378000000,06:51:20,-123.45
-sed do eiusmod tempor,-556132.25,1,,2019-04-18T02:45:55.555000000,23:46:03,\n";
- let right = writer.writer.into_inner().map(|s| s.to_string());
- assert_eq!(Some(left.to_string()), right.ok());
- }
-
- #[test]
- fn test_conversion_consistency() {
- // test if we can serialize and deserialize whilst retaining the same type information/ precision
-
- let schema = Schema::new(vec![
- Field::new("c1", DataType::Date32, false),
- Field::new("c2", DataType::Date64, false),
- ]);
-
- let c1 = Date32Array::from(vec![3, 2, 1]);
- let c2 = Date64Array::from(vec![3, 2, 1]);
-
- let batch = RecordBatch::try_new(
- Arc::new(schema.clone()),
- vec![Arc::new(c1), Arc::new(c2)],
- )
- .unwrap();
-
- let builder = WriterBuilder::new().has_headers(false);
-
- let mut buf: Cursor<Vec<u8>> = Default::default();
- // drop the writer early to release the borrow.
- {
- let mut writer = builder.build(&mut buf);
- writer.write(&batch).unwrap();
- }
- buf.set_position(0);
-
- let mut reader = Reader::new(
- buf,
- Arc::new(schema),
- false,
- None,
- 3,
- // starting at row 2 and up to row 6.
- None,
- None,
- );
- let rb = reader.next().unwrap().unwrap();
- let c1 = rb.column(0).as_any().downcast_ref::<Date32Array>().unwrap();
- let c2 = rb.column(1).as_any().downcast_ref::<Date64Array>().unwrap();
-
- let actual = c1.into_iter().collect::<Vec<_>>();
- let expected = vec![Some(3), Some(2), Some(1)];
- assert_eq!(actual, expected);
- let actual = c2.into_iter().collect::<Vec<_>>();
- let expected = vec![Some(3), Some(2), Some(1)];
- assert_eq!(actual, expected);
- }
-}
diff --git a/arrow/src/datatypes/datatype.rs b/arrow/src/datatypes/datatype.rs
deleted file mode 100644
index 122cbdd..0000000
--- a/arrow/src/datatypes/datatype.rs
+++ /dev/null
@@ -1,477 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-use std::fmt;
-
-use serde_derive::{Deserialize, Serialize};
-use serde_json::{json, Value, Value::String as VString};
-
-use crate::error::{ArrowError, Result};
-
-use super::Field;
-
-/// The set of datatypes that are supported by this implementation of Apache Arrow.
-///
-/// The Arrow specification on data types includes some more types.
-/// See also [`Schema.fbs`](https://github.com/apache/arrow/blob/master/format/Schema.fbs)
-/// for Arrow's specification.
-///
-/// The variants of this enum include primitive fixed size types as well as parametric or
-/// nested types.
-/// Currently the Rust implementation supports the following nested types:
-/// - `List<T>`
-/// - `Struct<T, U, V, ...>`
-///
-/// Nested types can themselves be nested within other arrays.
-/// For more information on these types please see
-/// [the physical memory layout of Apache Arrow](https://arrow.apache.org/docs/format/Columnar.html#physical-memory-layout).
-#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Eq, Hash, PartialOrd, Ord)]
-pub enum DataType {
- /// Null type
- Null,
- /// A boolean datatype representing the values `true` and `false`.
- Boolean,
- /// A signed 8-bit integer.
- Int8,
- /// A signed 16-bit integer.
- Int16,
- /// A signed 32-bit integer.
- Int32,
- /// A signed 64-bit integer.
- Int64,
- /// An unsigned 8-bit integer.
- UInt8,
- /// An unsigned 16-bit integer.
- UInt16,
- /// An unsigned 32-bit integer.
- UInt32,
- /// An unsigned 64-bit integer.
- UInt64,
- /// A 16-bit floating point number.
- Float16,
- /// A 32-bit floating point number.
- Float32,
- /// A 64-bit floating point number.
- Float64,
- /// A timestamp with an optional timezone.
- ///
- /// Time is measured as a Unix epoch, counting the seconds from
- /// 00:00:00.000 on 1 January 1970, excluding leap seconds,
- /// as a 64-bit integer.
- ///
- /// The time zone is a string indicating the name of a time zone, one of:
- ///
- /// * As used in the Olson time zone database (the "tz database" or
- /// "tzdata"), such as "America/New_York"
- /// * An absolute time zone offset of the form +XX:XX or -XX:XX, such as +07:30
- Timestamp(TimeUnit, Option<String>),
- /// A 32-bit date representing the elapsed time since UNIX epoch (1970-01-01)
- /// in days (32 bits).
- Date32,
- /// A 64-bit date representing the elapsed time since UNIX epoch (1970-01-01)
- /// in milliseconds (64 bits). Values are evenly divisible by 86400000.
- Date64,
- /// A 32-bit time representing the elapsed time since midnight in the unit of `TimeUnit`.
- Time32(TimeUnit),
- /// A 64-bit time representing the elapsed time since midnight in the unit of `TimeUnit`.
- Time64(TimeUnit),
- /// Measure of elapsed time in either seconds, milliseconds, microseconds or nanoseconds.
- Duration(TimeUnit),
- /// 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).
- Interval(IntervalUnit),
- /// Opaque binary data of variable length.
- Binary,
- /// Opaque binary data of fixed size.
- /// Enum parameter specifies the number of bytes per value.
- FixedSizeBinary(i32),
- /// Opaque binary data of variable length and 64-bit offsets.
- LargeBinary,
- /// A variable-length string in Unicode with UTF-8 encoding.
- Utf8,
- /// A variable-length string in Unicode with UFT-8 encoding and 64-bit offsets.
- LargeUtf8,
- /// A list of some logical data type with variable length.
- List(Box<Field>),
- /// A list of some logical data type with fixed length.
- FixedSizeList(Box<Field>, i32),
- /// A list of some logical data type with variable length and 64-bit offsets.
- LargeList(Box<Field>),
- /// A nested datatype that contains a number of sub-fields.
- Struct(Vec<Field>),
- /// A nested datatype that can represent slots of differing types.
- Union(Vec<Field>),
- /// A dictionary encoded array (`key_type`, `value_type`), where
- /// each array element is an index of `key_type` into an
- /// associated dictionary of `value_type`.
- ///
- /// Dictionary arrays are used to store columns of `value_type`
- /// that contain many repeated values using less memory, but with
- /// a higher CPU overhead for some operations.
- ///
- /// This type mostly used to represent low cardinality string
- /// arrays or a limited set of primitive types as integers.
- Dictionary(Box<DataType>, Box<DataType>),
- /// Decimal value with precision and scale
- Decimal(usize, usize),
-}
-
-/// An absolute length of time in seconds, milliseconds, microseconds or nanoseconds.
-#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Eq, Hash, PartialOrd, Ord)]
-pub enum TimeUnit {
- /// Time in seconds.
- Second,
- /// Time in milliseconds.
- Millisecond,
- /// Time in microseconds.
- Microsecond,
- /// Time in nanoseconds.
- Nanosecond,
-}
-
-/// YEAR_MONTH or DAY_TIME interval in SQL style.
-#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Eq, Hash, PartialOrd, Ord)]
-pub enum IntervalUnit {
- /// Indicates the number of elapsed whole months, stored as 4-byte integers.
- YearMonth,
- /// Indicates the number of elapsed days and milliseconds,
- /// stored as 2 contiguous 32-bit integers (days, milliseconds) (8-bytes in total).
- DayTime,
-}
-
-impl fmt::Display for DataType {
- fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
- write!(f, "{:?}", self)
- }
-}
-
-impl DataType {
- /// Parse a data type from a JSON representation.
- pub(crate) fn from(json: &Value) -> Result<DataType> {
- let default_field = Field::new("", DataType::Boolean, true);
- match *json {
- Value::Object(ref map) => match map.get("name") {
- Some(s) if s == "null" => Ok(DataType::Null),
- Some(s) if s == "bool" => Ok(DataType::Boolean),
- Some(s) if s == "binary" => Ok(DataType::Binary),
- Some(s) if s == "largebinary" => Ok(DataType::LargeBinary),
- Some(s) if s == "utf8" => Ok(DataType::Utf8),
- Some(s) if s == "largeutf8" => Ok(DataType::LargeUtf8),
- Some(s) if s == "fixedsizebinary" => {
- // return a list with any type as its child isn't defined in the map
- if let Some(Value::Number(size)) = map.get("byteWidth") {
- Ok(DataType::FixedSizeBinary(size.as_i64().unwrap() as i32))
- } else {
- Err(ArrowError::ParseError(
- "Expecting a byteWidth for fixedsizebinary".to_string(),
- ))
- }
- }
- Some(s) if s == "decimal" => {
- // return a list with any type as its child isn't defined in the map
- let precision = match map.get("precision") {
- Some(p) => Ok(p.as_u64().unwrap() as usize),
- None => Err(ArrowError::ParseError(
- "Expecting a precision for decimal".to_string(),
- )),
- };
- let scale = match map.get("scale") {
- Some(s) => Ok(s.as_u64().unwrap() as usize),
- _ => Err(ArrowError::ParseError(
- "Expecting a scale for decimal".to_string(),
- )),
- };
-
- Ok(DataType::Decimal(precision?, scale?))
- }
- Some(s) if s == "floatingpoint" => match map.get("precision") {
- Some(p) if p == "HALF" => Ok(DataType::Float16),
- Some(p) if p == "SINGLE" => Ok(DataType::Float32),
- Some(p) if p == "DOUBLE" => Ok(DataType::Float64),
- _ => Err(ArrowError::ParseError(
- "floatingpoint precision missing or invalid".to_string(),
- )),
- },
- Some(s) if s == "timestamp" => {
- let unit = match map.get("unit") {
- Some(p) if p == "SECOND" => Ok(TimeUnit::Second),
- Some(p) if p == "MILLISECOND" => Ok(TimeUnit::Millisecond),
- Some(p) if p == "MICROSECOND" => Ok(TimeUnit::Microsecond),
- Some(p) if p == "NANOSECOND" => Ok(TimeUnit::Nanosecond),
- _ => Err(ArrowError::ParseError(
- "timestamp unit missing or invalid".to_string(),
- )),
- };
- let tz = match map.get("timezone") {
- None => Ok(None),
- Some(VString(tz)) => Ok(Some(tz.clone())),
- _ => Err(ArrowError::ParseError(
- "timezone must be a string".to_string(),
- )),
- };
- Ok(DataType::Timestamp(unit?, tz?))
- }
- Some(s) if s == "date" => match map.get("unit") {
- Some(p) if p == "DAY" => Ok(DataType::Date32),
- Some(p) if p == "MILLISECOND" => Ok(DataType::Date64),
- _ => Err(ArrowError::ParseError(
- "date unit missing or invalid".to_string(),
- )),
- },
- Some(s) if s == "time" => {
- let unit = match map.get("unit") {
- Some(p) if p == "SECOND" => Ok(TimeUnit::Second),
- Some(p) if p == "MILLISECOND" => Ok(TimeUnit::Millisecond),
- Some(p) if p == "MICROSECOND" => Ok(TimeUnit::Microsecond),
- Some(p) if p == "NANOSECOND" => Ok(TimeUnit::Nanosecond),
- _ => Err(ArrowError::ParseError(
- "time unit missing or invalid".to_string(),
- )),
- };
- match map.get("bitWidth") {
- Some(p) if p == 32 => Ok(DataType::Time32(unit?)),
- Some(p) if p == 64 => Ok(DataType::Time64(unit?)),
- _ => Err(ArrowError::ParseError(
- "time bitWidth missing or invalid".to_string(),
- )),
- }
- }
- Some(s) if s == "duration" => match map.get("unit") {
- Some(p) if p == "SECOND" => Ok(DataType::Duration(TimeUnit::Second)),
- Some(p) if p == "MILLISECOND" => {
- Ok(DataType::Duration(TimeUnit::Millisecond))
- }
- Some(p) if p == "MICROSECOND" => {
- Ok(DataType::Duration(TimeUnit::Microsecond))
- }
- Some(p) if p == "NANOSECOND" => {
- Ok(DataType::Duration(TimeUnit::Nanosecond))
- }
- _ => Err(ArrowError::ParseError(
- "time unit missing or invalid".to_string(),
- )),
- },
- Some(s) if s == "interval" => match map.get("unit") {
- Some(p) if p == "DAY_TIME" => {
- Ok(DataType::Interval(IntervalUnit::DayTime))
- }
- Some(p) if p == "YEAR_MONTH" => {
- Ok(DataType::Interval(IntervalUnit::YearMonth))
- }
- _ => Err(ArrowError::ParseError(
- "interval unit missing or invalid".to_string(),
- )),
- },
- Some(s) if s == "int" => match map.get("isSigned") {
- Some(&Value::Bool(true)) => match map.get("bitWidth") {
- Some(&Value::Number(ref n)) => match n.as_u64() {
- Some(8) => Ok(DataType::Int8),
- Some(16) => Ok(DataType::Int16),
- Some(32) => Ok(DataType::Int32),
- Some(64) => Ok(DataType::Int64),
- _ => Err(ArrowError::ParseError(
- "int bitWidth missing or invalid".to_string(),
- )),
- },
- _ => Err(ArrowError::ParseError(
- "int bitWidth missing or invalid".to_string(),
- )),
- },
- Some(&Value::Bool(false)) => match map.get("bitWidth") {
- Some(&Value::Number(ref n)) => match n.as_u64() {
- Some(8) => Ok(DataType::UInt8),
- Some(16) => Ok(DataType::UInt16),
- Some(32) => Ok(DataType::UInt32),
- Some(64) => Ok(DataType::UInt64),
- _ => Err(ArrowError::ParseError(
- "int bitWidth missing or invalid".to_string(),
- )),
- },
- _ => Err(ArrowError::ParseError(
- "int bitWidth missing or invalid".to_string(),
- )),
- },
- _ => Err(ArrowError::ParseError(
- "int signed missing or invalid".to_string(),
- )),
- },
- Some(s) if s == "list" => {
- // return a list with any type as its child isn't defined in the map
- Ok(DataType::List(Box::new(default_field)))
- }
- Some(s) if s == "largelist" => {
- // return a largelist with any type as its child isn't defined in the map
- Ok(DataType::LargeList(Box::new(default_field)))
- }
- Some(s) if s == "fixedsizelist" => {
- // return a list with any type as its child isn't defined in the map
- if let Some(Value::Number(size)) = map.get("listSize") {
- Ok(DataType::FixedSizeList(
- Box::new(default_field),
- size.as_i64().unwrap() as i32,
- ))
- } else {
- Err(ArrowError::ParseError(
- "Expecting a listSize for fixedsizelist".to_string(),
- ))
- }
- }
- Some(s) if s == "struct" => {
- // return an empty `struct` type as its children aren't defined in the map
- Ok(DataType::Struct(vec![]))
- }
- Some(other) => Err(ArrowError::ParseError(format!(
- "invalid or unsupported type name: {} in {:?}",
- other, json
- ))),
- None => Err(ArrowError::ParseError("type name missing".to_string())),
- },
- _ => Err(ArrowError::ParseError(
- "invalid json value type".to_string(),
- )),
- }
- }
-
- /// Generate a JSON representation of the data type.
- pub fn to_json(&self) -> Value {
- match self {
- DataType::Null => json!({"name": "null"}),
- DataType::Boolean => json!({"name": "bool"}),
- DataType::Int8 => json!({"name": "int", "bitWidth": 8, "isSigned": true}),
- DataType::Int16 => json!({"name": "int", "bitWidth": 16, "isSigned": true}),
- DataType::Int32 => json!({"name": "int", "bitWidth": 32, "isSigned": true}),
- DataType::Int64 => json!({"name": "int", "bitWidth": 64, "isSigned": true}),
- DataType::UInt8 => json!({"name": "int", "bitWidth": 8, "isSigned": false}),
- DataType::UInt16 => json!({"name": "int", "bitWidth": 16, "isSigned": false}),
- DataType::UInt32 => json!({"name": "int", "bitWidth": 32, "isSigned": false}),
- DataType::UInt64 => json!({"name": "int", "bitWidth": 64, "isSigned": false}),
- DataType::Float16 => json!({"name": "floatingpoint", "precision": "HALF"}),
- DataType::Float32 => json!({"name": "floatingpoint", "precision": "SINGLE"}),
- DataType::Float64 => json!({"name": "floatingpoint", "precision": "DOUBLE"}),
- DataType::Utf8 => json!({"name": "utf8"}),
- DataType::LargeUtf8 => json!({"name": "largeutf8"}),
- DataType::Binary => json!({"name": "binary"}),
- DataType::LargeBinary => json!({"name": "largebinary"}),
- DataType::FixedSizeBinary(byte_width) => {
- json!({"name": "fixedsizebinary", "byteWidth": byte_width})
- }
- DataType::Struct(_) => json!({"name": "struct"}),
- DataType::Union(_) => json!({"name": "union"}),
- DataType::List(_) => json!({ "name": "list"}),
- DataType::LargeList(_) => json!({ "name": "largelist"}),
- DataType::FixedSizeList(_, length) => {
- json!({"name":"fixedsizelist", "listSize": length})
- }
- DataType::Time32(unit) => {
- json!({"name": "time", "bitWidth": 32, "unit": match unit {
- TimeUnit::Second => "SECOND",
- TimeUnit::Millisecond => "MILLISECOND",
- TimeUnit::Microsecond => "MICROSECOND",
- TimeUnit::Nanosecond => "NANOSECOND",
- }})
- }
- DataType::Time64(unit) => {
- json!({"name": "time", "bitWidth": 64, "unit": match unit {
- TimeUnit::Second => "SECOND",
- TimeUnit::Millisecond => "MILLISECOND",
- TimeUnit::Microsecond => "MICROSECOND",
- TimeUnit::Nanosecond => "NANOSECOND",
- }})
- }
- DataType::Date32 => {
- json!({"name": "date", "unit": "DAY"})
- }
- DataType::Date64 => {
- json!({"name": "date", "unit": "MILLISECOND"})
- }
- DataType::Timestamp(unit, None) => {
- json!({"name": "timestamp", "unit": match unit {
- TimeUnit::Second => "SECOND",
- TimeUnit::Millisecond => "MILLISECOND",
- TimeUnit::Microsecond => "MICROSECOND",
- TimeUnit::Nanosecond => "NANOSECOND",
- }})
- }
- DataType::Timestamp(unit, Some(tz)) => {
- json!({"name": "timestamp", "unit": match unit {
- TimeUnit::Second => "SECOND",
- TimeUnit::Millisecond => "MILLISECOND",
- TimeUnit::Microsecond => "MICROSECOND",
- TimeUnit::Nanosecond => "NANOSECOND",
- }, "timezone": tz})
- }
- DataType::Interval(unit) => json!({"name": "interval", "unit": match unit {
- IntervalUnit::YearMonth => "YEAR_MONTH",
- IntervalUnit::DayTime => "DAY_TIME",
- }}),
- DataType::Duration(unit) => json!({"name": "duration", "unit": match unit {
- TimeUnit::Second => "SECOND",
- TimeUnit::Millisecond => "MILLISECOND",
- TimeUnit::Microsecond => "MICROSECOND",
- TimeUnit::Nanosecond => "NANOSECOND",
- }}),
- DataType::Dictionary(_, _) => json!({ "name": "dictionary"}),
- DataType::Decimal(precision, scale) => {
- json!({"name": "decimal", "precision": precision, "scale": scale})
- }
- }
- }
-
- /// Returns true if this type is numeric: (UInt*, Unit*, or Float*).
- pub fn is_numeric(t: &DataType) -> bool {
- use DataType::*;
- matches!(
- t,
- UInt8
- | UInt16
- | UInt32
- | UInt64
- | Int8
- | Int16
- | Int32
- | Int64
- | Float32
- | Float64
- )
- }
-
- /// Compares the datatype with another, ignoring nested field names
- /// and metadata.
- pub(crate) fn equals_datatype(&self, other: &DataType) -> bool {
- match (&self, other) {
- (DataType::List(a), DataType::List(b))
- | (DataType::LargeList(a), DataType::LargeList(b)) => {
- a.is_nullable() == b.is_nullable()
- && a.data_type().equals_datatype(b.data_type())
- }
- (DataType::FixedSizeList(a, a_size), DataType::FixedSizeList(b, b_size)) => {
- a_size == b_size
- && a.is_nullable() == b.is_nullable()
- && a.data_type().equals_datatype(b.data_type())
- }
- (DataType::Struct(a), DataType::Struct(b)) => {
- a.len() == b.len()
- && a.iter().zip(b).all(|(a, b)| {
- a.is_nullable() == b.is_nullable()
- && a.data_type().equals_datatype(b.data_type())
- })
- }
- _ => self == other,
- }
- }
-}
diff --git a/arrow/src/datatypes/field.rs b/arrow/src/datatypes/field.rs
deleted file mode 100644
index a471f12..0000000
--- a/arrow/src/datatypes/field.rs
+++ /dev/null
@@ -1,541 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-use std::collections::BTreeMap;
-
-use serde_derive::{Deserialize, Serialize};
-use serde_json::{json, Value};
-
-use crate::error::{ArrowError, Result};
-
-use super::DataType;
-
-/// Contains the meta-data for a single relative type.
-///
-/// The `Schema` object is an ordered collection of `Field` objects.
-#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Eq, Hash, PartialOrd, Ord)]
-pub struct Field {
- name: String,
- data_type: DataType,
- nullable: bool,
- dict_id: i64,
- dict_is_ordered: bool,
- /// A map of key-value pairs containing additional custom meta data.
- #[serde(skip_serializing_if = "Option::is_none")]
- metadata: Option<BTreeMap<String, String>>,
-}
-
-impl Field {
- /// Creates a new field
- pub fn new(name: &str, data_type: DataType, nullable: bool) -> Self {
- Field {
- name: name.to_string(),
- data_type,
- nullable,
- dict_id: 0,
- dict_is_ordered: false,
- metadata: None,
- }
- }
-
- /// Creates a new field
- pub fn new_dict(
- name: &str,
- data_type: DataType,
- nullable: bool,
- dict_id: i64,
- dict_is_ordered: bool,
- ) -> Self {
- Field {
- name: name.to_string(),
- data_type,
- nullable,
- dict_id,
- dict_is_ordered,
- metadata: None,
- }
- }
-
- /// Sets the `Field`'s optional custom metadata.
- /// The metadata is set as `None` for empty map.
- #[inline]
- pub fn set_metadata(&mut self, metadata: Option<BTreeMap<String, String>>) {
- // To make serde happy, convert Some(empty_map) to None.
- self.metadata = None;
- if let Some(v) = metadata {
- if !v.is_empty() {
- self.metadata = Some(v);
- }
- }
- }
-
- /// Returns the immutable reference to the `Field`'s optional custom metadata.
- #[inline]
- pub const fn metadata(&self) -> &Option<BTreeMap<String, String>> {
- &self.metadata
- }
-
- /// Returns an immutable reference to the `Field`'s name.
- #[inline]
- pub const fn name(&self) -> &String {
- &self.name
- }
-
- /// Returns an immutable reference to the `Field`'s data-type.
- #[inline]
- pub const fn data_type(&self) -> &DataType {
- &self.data_type
- }
-
- /// Indicates whether this `Field` supports null values.
- #[inline]
- pub const fn is_nullable(&self) -> bool {
- self.nullable
- }
-
- /// Returns the dictionary ID, if this is a dictionary type.
- #[inline]
- pub const fn dict_id(&self) -> Option<i64> {
- match self.data_type {
- DataType::Dictionary(_, _) => Some(self.dict_id),
- _ => None,
- }
- }
-
- /// Returns whether this `Field`'s dictionary is ordered, if this is a dictionary type.
- #[inline]
- pub const fn dict_is_ordered(&self) -> Option<bool> {
- match self.data_type {
- DataType::Dictionary(_, _) => Some(self.dict_is_ordered),
- _ => None,
- }
- }
-
- /// Parse a `Field` definition from a JSON representation.
- pub fn from(json: &Value) -> Result<Self> {
- match *json {
- Value::Object(ref map) => {
- let name = match map.get("name") {
- Some(&Value::String(ref name)) => name.to_string(),
- _ => {
- return Err(ArrowError::ParseError(
- "Field missing 'name' attribute".to_string(),
- ));
- }
- };
- let nullable = match map.get("nullable") {
- Some(&Value::Bool(b)) => b,
- _ => {
- return Err(ArrowError::ParseError(
- "Field missing 'nullable' attribute".to_string(),
- ));
- }
- };
- let data_type = match map.get("type") {
- Some(t) => DataType::from(t)?,
- _ => {
- return Err(ArrowError::ParseError(
- "Field missing 'type' attribute".to_string(),
- ));
- }
- };
-
- // Referenced example file: testing/data/arrow-ipc-stream/integration/1.0.0-littleendian/generated_custom_metadata.json.gz
- let metadata = match map.get("metadata") {
- Some(&Value::Array(ref values)) => {
- let mut res: BTreeMap<String, String> = BTreeMap::new();
- for value in values {
- match value.as_object() {
- Some(map) => {
- if map.len() != 2 {
- return Err(ArrowError::ParseError(
- "Field 'metadata' must have exact two entries for each key-value map".to_string(),
- ));
- }
- if let (Some(k), Some(v)) =
- (map.get("key"), map.get("value"))
- {
- if let (Some(k_str), Some(v_str)) =
- (k.as_str(), v.as_str())
- {
- res.insert(
- k_str.to_string().clone(),
- v_str.to_string().clone(),
- );
- } else {
- return Err(ArrowError::ParseError("Field 'metadata' must have map value of string type".to_string()));
- }
- } else {
- return Err(ArrowError::ParseError("Field 'metadata' lacks map keys named \"key\" or \"value\"".to_string()));
- }
- }
- _ => {
- return Err(ArrowError::ParseError(
- "Field 'metadata' contains non-object key-value pair".to_string(),
- ));
- }
- }
- }
- Some(res)
- }
- // We also support map format, because Schema's metadata supports this.
- // See https://github.com/apache/arrow/pull/5907
- Some(&Value::Object(ref values)) => {
- let mut res: BTreeMap<String, String> = BTreeMap::new();
- for (k, v) in values {
- if let Some(str_value) = v.as_str() {
- res.insert(k.clone(), str_value.to_string().clone());
- } else {
- return Err(ArrowError::ParseError(
- format!("Field 'metadata' contains non-string value for key {}", k),
- ));
- }
- }
- Some(res)
- }
- Some(_) => {
- return Err(ArrowError::ParseError(
- "Field `metadata` is not json array".to_string(),
- ));
- }
- _ => None,
- };
-
- // if data_type is a struct or list, get its children
- let data_type = match data_type {
- DataType::List(_)
- | DataType::LargeList(_)
- | DataType::FixedSizeList(_, _) => match map.get("children") {
- Some(Value::Array(values)) => {
- if values.len() != 1 {
- return Err(ArrowError::ParseError(
- "Field 'children' must have one element for a list data type".to_string(),
- ));
- }
- match data_type {
- DataType::List(_) => {
- DataType::List(Box::new(Self::from(&values[0])?))
- }
- DataType::LargeList(_) => {
- DataType::LargeList(Box::new(Self::from(&values[0])?))
- }
- DataType::FixedSizeList(_, int) => DataType::FixedSizeList(
- Box::new(Self::from(&values[0])?),
- int,
- ),
- _ => unreachable!(
- "Data type should be a list, largelist or fixedsizelist"
- ),
- }
- }
- Some(_) => {
- return Err(ArrowError::ParseError(
- "Field 'children' must be an array".to_string(),
- ))
- }
- None => {
- return Err(ArrowError::ParseError(
- "Field missing 'children' attribute".to_string(),
- ));
- }
- },
- DataType::Struct(mut fields) => match map.get("children") {
- Some(Value::Array(values)) => {
- let struct_fields: Result<Vec<Field>> =
- values.iter().map(|v| Field::from(v)).collect();
- fields.append(&mut struct_fields?);
- DataType::Struct(fields)
- }
- Some(_) => {
- return Err(ArrowError::ParseError(
- "Field 'children' must be an array".to_string(),
- ))
- }
- None => {
- return Err(ArrowError::ParseError(
- "Field missing 'children' attribute".to_string(),
- ));
- }
- },
- _ => data_type,
- };
-
- let mut dict_id = 0;
- let mut dict_is_ordered = false;
-
- let data_type = match map.get("dictionary") {
- Some(dictionary) => {
- let index_type = match dictionary.get("indexType") {
- Some(t) => DataType::from(t)?,
- _ => {
- return Err(ArrowError::ParseError(
- "Field missing 'indexType' attribute".to_string(),
- ));
- }
- };
- dict_id = match dictionary.get("id") {
- Some(Value::Number(n)) => n.as_i64().unwrap(),
- _ => {
- return Err(ArrowError::ParseError(
- "Field missing 'id' attribute".to_string(),
- ));
- }
- };
- dict_is_ordered = match dictionary.get("isOrdered") {
- Some(&Value::Bool(n)) => n,
- _ => {
- return Err(ArrowError::ParseError(
- "Field missing 'isOrdered' attribute".to_string(),
- ));
- }
- };
- DataType::Dictionary(Box::new(index_type), Box::new(data_type))
- }
- _ => data_type,
- };
- Ok(Field {
- name,
- data_type,
- nullable,
- dict_id,
- dict_is_ordered,
- metadata,
- })
- }
- _ => Err(ArrowError::ParseError(
- "Invalid json value type for field".to_string(),
- )),
- }
- }
-
- /// Generate a JSON representation of the `Field`.
- pub fn to_json(&self) -> Value {
- let children: Vec<Value> = match self.data_type() {
- DataType::Struct(fields) => fields.iter().map(|f| f.to_json()).collect(),
- DataType::List(field) => vec![field.to_json()],
- DataType::LargeList(field) => vec![field.to_json()],
- DataType::FixedSizeList(field, _) => vec![field.to_json()],
- _ => vec![],
- };
- match self.data_type() {
- DataType::Dictionary(ref index_type, ref value_type) => json!({
- "name": self.name,
- "nullable": self.nullable,
- "type": value_type.to_json(),
- "children": children,
- "dictionary": {
- "id": self.dict_id,
- "indexType": index_type.to_json(),
- "isOrdered": self.dict_is_ordered
- }
- }),
- _ => json!({
- "name": self.name,
- "nullable": self.nullable,
- "type": self.data_type.to_json(),
- "children": children
- }),
- }
- }
-
- /// Merge field into self if it is compatible. Struct will be merged recursively.
- /// NOTE: `self` may be updated to unexpected state in case of merge failure.
- ///
- /// Example:
- ///
- /// ```
- /// use arrow::datatypes::*;
- ///
- /// let mut field = Field::new("c1", DataType::Int64, false);
- /// assert!(field.try_merge(&Field::new("c1", DataType::Int64, true)).is_ok());
- /// assert!(field.is_nullable());
- /// ```
- pub fn try_merge(&mut self, from: &Field) -> Result<()> {
- // merge metadata
- match (self.metadata(), from.metadata()) {
- (Some(self_metadata), Some(from_metadata)) => {
- let mut merged = self_metadata.clone();
- for (key, from_value) in from_metadata {
- if let Some(self_value) = self_metadata.get(key) {
- if self_value != from_value {
- return Err(ArrowError::SchemaError(format!(
- "Fail to merge field due to conflicting metadata data value for key {}", key),
- ));
- }
- } else {
- merged.insert(key.clone(), from_value.clone());
- }
- }
- self.set_metadata(Some(merged));
- }
- (None, Some(from_metadata)) => {
- self.set_metadata(Some(from_metadata.clone()));
- }
- _ => {}
- }
- if from.dict_id != self.dict_id {
- return Err(ArrowError::SchemaError(
- "Fail to merge schema Field due to conflicting dict_id".to_string(),
- ));
- }
- if from.dict_is_ordered != self.dict_is_ordered {
- return Err(ArrowError::SchemaError(
- "Fail to merge schema Field due to conflicting dict_is_ordered"
- .to_string(),
- ));
- }
- match &mut self.data_type {
- DataType::Struct(nested_fields) => match &from.data_type {
- DataType::Struct(from_nested_fields) => {
- for from_field in from_nested_fields {
- let mut is_new_field = true;
- for self_field in nested_fields.iter_mut() {
- if self_field.name != from_field.name {
- continue;
- }
- is_new_field = false;
- self_field.try_merge(&from_field)?;
- }
- if is_new_field {
- nested_fields.push(from_field.clone());
- }
- }
- }
- _ => {
- return Err(ArrowError::SchemaError(
- "Fail to merge schema Field due to conflicting datatype"
- .to_string(),
- ));
- }
- },
- DataType::Union(nested_fields) => match &from.data_type {
- DataType::Union(from_nested_fields) => {
- for from_field in from_nested_fields {
- let mut is_new_field = true;
- for self_field in nested_fields.iter_mut() {
- if from_field == self_field {
- is_new_field = false;
- break;
- }
- }
- if is_new_field {
- nested_fields.push(from_field.clone());
- }
- }
- }
- _ => {
- return Err(ArrowError::SchemaError(
- "Fail to merge schema Field due to conflicting datatype"
- .to_string(),
- ));
- }
- },
- DataType::Null
- | DataType::Boolean
- | DataType::Int8
- | DataType::Int16
- | DataType::Int32
- | DataType::Int64
- | DataType::UInt8
- | DataType::UInt16
- | DataType::UInt32
- | DataType::UInt64
- | DataType::Float16
- | DataType::Float32
- | DataType::Float64
- | DataType::Timestamp(_, _)
- | DataType::Date32
- | DataType::Date64
- | DataType::Time32(_)
- | DataType::Time64(_)
- | DataType::Duration(_)
- | DataType::Binary
- | DataType::LargeBinary
- | DataType::Interval(_)
- | DataType::LargeList(_)
- | DataType::List(_)
- | DataType::Dictionary(_, _)
- | DataType::FixedSizeList(_, _)
- | DataType::FixedSizeBinary(_)
- | DataType::Utf8
- | DataType::LargeUtf8
- | DataType::Decimal(_, _) => {
- if self.data_type != from.data_type {
- return Err(ArrowError::SchemaError(
- "Fail to merge schema Field due to conflicting datatype"
- .to_string(),
- ));
- }
- }
- }
- if from.nullable {
- self.nullable = from.nullable;
- }
-
- Ok(())
- }
-
- /// Check to see if `self` is a superset of `other` field. Superset is defined as:
- ///
- /// * if nullability doesn't match, self needs to be nullable
- /// * self.metadata is a superset of other.metadata
- /// * all other fields are equal
- pub fn contains(&self, other: &Field) -> bool {
- if self.name != other.name
- || self.data_type != other.data_type
- || self.dict_id != other.dict_id
- || self.dict_is_ordered != other.dict_is_ordered
- {
- return false;
- }
-
- if self.nullable != other.nullable && !self.nullable {
- return false;
- }
-
- // make sure self.metadata is a superset of other.metadata
- match (&self.metadata, &other.metadata) {
- (None, Some(_)) => {
- return false;
- }
- (Some(self_meta), Some(other_meta)) => {
- for (k, v) in other_meta.iter() {
- match self_meta.get(k) {
- Some(s) => {
- if s != v {
- return false;
- }
- }
- None => {
- return false;
- }
- }
- }
- }
- _ => {}
- }
-
- true
- }
-}
-
-// TODO: improve display with crate https://crates.io/crates/derive_more ?
-impl std::fmt::Display for Field {
- fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
- write!(f, "{:?}", self)
- }
-}
diff --git a/arrow/src/datatypes/mod.rs b/arrow/src/datatypes/mod.rs
deleted file mode 100644
index 175b50b..0000000
--- a/arrow/src/datatypes/mod.rs
+++ /dev/null
@@ -1,1241 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-//! Defines the logical data types of Arrow arrays.
-//!
-//! The most important things you might be looking for are:
-//! * [`Schema`](crate::datatypes::Schema) to describe a schema.
-//! * [`Field`](crate::datatypes::Field) to describe one field within a schema.
-//! * [`DataType`](crate::datatypes::DataType) to describe the type of a field.
-
-use std::sync::Arc;
-
-mod native;
-pub use native::*;
-mod field;
-pub use field::*;
-mod schema;
-pub use schema::*;
-mod numeric;
-pub use numeric::*;
-mod types;
-pub use types::*;
-mod datatype;
-pub use datatype::*;
-
-/// A reference-counted reference to a [`Schema`](crate::datatypes::Schema).
-pub type SchemaRef = Arc<Schema>;
-
-#[cfg(test)]
-mod tests {
- use super::*;
- use crate::error::Result;
- use serde_json::Value::{Bool, Number as VNumber};
- use serde_json::{Number, Value};
- use std::{
- collections::{BTreeMap, HashMap},
- f32::NAN,
- };
-
- #[test]
- fn test_list_datatype_equality() {
- // tests that list type equality is checked while ignoring list names
- let list_a = DataType::List(Box::new(Field::new("item", DataType::Int32, true)));
- let list_b = DataType::List(Box::new(Field::new("array", DataType::Int32, true)));
- let list_c = DataType::List(Box::new(Field::new("item", DataType::Int32, false)));
- let list_d = DataType::List(Box::new(Field::new("item", DataType::UInt32, true)));
- assert!(list_a.equals_datatype(&list_b));
- assert!(!list_a.equals_datatype(&list_c));
- assert!(!list_b.equals_datatype(&list_c));
- assert!(!list_a.equals_datatype(&list_d));
-
- let list_e =
- DataType::FixedSizeList(Box::new(Field::new("item", list_a, false)), 3);
- let list_f =
- DataType::FixedSizeList(Box::new(Field::new("array", list_b, false)), 3);
- let list_g = DataType::FixedSizeList(
- Box::new(Field::new("item", DataType::FixedSizeBinary(3), true)),
- 3,
- );
- assert!(list_e.equals_datatype(&list_f));
- assert!(!list_e.equals_datatype(&list_g));
- assert!(!list_f.equals_datatype(&list_g));
-
- let list_h = DataType::Struct(vec![Field::new("f1", list_e, true)]);
- let list_i = DataType::Struct(vec![Field::new("f1", list_f.clone(), true)]);
- let list_j = DataType::Struct(vec![Field::new("f1", list_f.clone(), false)]);
- let list_k = DataType::Struct(vec![
- Field::new("f1", list_f.clone(), false),
- Field::new("f2", list_g.clone(), false),
- Field::new("f3", DataType::Utf8, true),
- ]);
- let list_l = DataType::Struct(vec![
- Field::new("ff1", list_f.clone(), false),
- Field::new("ff2", list_g.clone(), false),
- Field::new("ff3", DataType::LargeUtf8, true),
- ]);
- let list_m = DataType::Struct(vec![
- Field::new("ff1", list_f, false),
- Field::new("ff2", list_g, false),
- Field::new("ff3", DataType::Utf8, true),
- ]);
- assert!(list_h.equals_datatype(&list_i));
- assert!(!list_h.equals_datatype(&list_j));
- assert!(!list_k.equals_datatype(&list_l));
- assert!(list_k.equals_datatype(&list_m));
- }
-
- #[test]
- fn create_struct_type() {
- let _person = DataType::Struct(vec![
- Field::new("first_name", DataType::Utf8, false),
- Field::new("last_name", DataType::Utf8, false),
- Field::new(
- "address",
- DataType::Struct(vec![
- Field::new("street", DataType::Utf8, false),
- Field::new("zip", DataType::UInt16, false),
- ]),
- false,
- ),
- ]);
- }
-
- #[test]
- fn serde_struct_type() {
- let kv_array = [("k".to_string(), "v".to_string())];
- let field_metadata: BTreeMap<String, String> = kv_array.iter().cloned().collect();
-
- // Non-empty map: should be converted as JSON obj { ... }
- let mut first_name = Field::new("first_name", DataType::Utf8, false);
- first_name.set_metadata(Some(field_metadata));
-
- // Empty map: should be omitted.
- let mut last_name = Field::new("last_name", DataType::Utf8, false);
- last_name.set_metadata(Some(BTreeMap::default()));
-
- let person = DataType::Struct(vec![
- first_name,
- last_name,
- Field::new(
- "address",
- DataType::Struct(vec![
- Field::new("street", DataType::Utf8, false),
- Field::new("zip", DataType::UInt16, false),
- ]),
- false,
- ),
- ]);
-
- let serialized = serde_json::to_string(&person).unwrap();
-
- // NOTE that this is testing the default (derived) serialization format, not the
- // JSON format specified in metadata.md
-
- assert_eq!(
- "{\"Struct\":[\
- {\"name\":\"first_name\",\"data_type\":\"Utf8\",\"nullable\":false,\"dict_id\":0,\"dict_is_ordered\":false,\"metadata\":{\"k\":\"v\"}},\
- {\"name\":\"last_name\",\"data_type\":\"Utf8\",\"nullable\":false,\"dict_id\":0,\"dict_is_ordered\":false},\
- {\"name\":\"address\",\"data_type\":{\"Struct\":\
- [{\"name\":\"street\",\"data_type\":\"Utf8\",\"nullable\":false,\"dict_id\":0,\"dict_is_ordered\":false},\
- {\"name\":\"zip\",\"data_type\":\"UInt16\",\"nullable\":false,\"dict_id\":0,\"dict_is_ordered\":false}\
- ]},\"nullable\":false,\"dict_id\":0,\"dict_is_ordered\":false}]}",
- serialized
- );
-
- let deserialized = serde_json::from_str(&serialized).unwrap();
-
- assert_eq!(person, deserialized);
- }
-
- #[test]
- fn struct_field_to_json() {
- let f = Field::new(
- "address",
- DataType::Struct(vec![
- Field::new("street", DataType::Utf8, false),
- Field::new("zip", DataType::UInt16, false),
- ]),
- false,
- );
- let value: Value = serde_json::from_str(
- r#"{
- "name": "address",
- "nullable": false,
- "type": {
- "name": "struct"
- },
- "children": [
- {
- "name": "street",
- "nullable": false,
- "type": {
- "name": "utf8"
- },
- "children": []
- },
- {
- "name": "zip",
- "nullable": false,
- "type": {
- "name": "int",
- "bitWidth": 16,
- "isSigned": false
- },
- "children": []
- }
- ]
- }"#,
- )
- .unwrap();
- assert_eq!(value, f.to_json());
- }
-
- #[test]
- fn primitive_field_to_json() {
- let f = Field::new("first_name", DataType::Utf8, false);
- let value: Value = serde_json::from_str(
- r#"{
- "name": "first_name",
- "nullable": false,
- "type": {
- "name": "utf8"
- },
- "children": []
- }"#,
- )
- .unwrap();
- assert_eq!(value, f.to_json());
- }
- #[test]
- fn parse_struct_from_json() {
- let json = r#"
- {
- "name": "address",
- "type": {
- "name": "struct"
- },
- "nullable": false,
- "children": [
- {
- "name": "street",
- "type": {
- "name": "utf8"
- },
- "nullable": false,
- "children": []
- },
- {
- "name": "zip",
- "type": {
- "name": "int",
- "isSigned": false,
- "bitWidth": 16
- },
- "nullable": false,
- "children": []
- }
- ]
- }
- "#;
- let value: Value = serde_json::from_str(json).unwrap();
- let dt = Field::from(&value).unwrap();
-
- let expected = Field::new(
- "address",
- DataType::Struct(vec![
- Field::new("street", DataType::Utf8, false),
- Field::new("zip", DataType::UInt16, false),
- ]),
- false,
- );
-
- assert_eq!(expected, dt);
- }
-
- #[test]
- fn parse_utf8_from_json() {
- let json = "{\"name\":\"utf8\"}";
- let value: Value = serde_json::from_str(json).unwrap();
- let dt = DataType::from(&value).unwrap();
- assert_eq!(DataType::Utf8, dt);
- }
-
- #[test]
- fn parse_int32_from_json() {
- let json = "{\"name\": \"int\", \"isSigned\": true, \"bitWidth\": 32}";
- let value: Value = serde_json::from_str(json).unwrap();
- let dt = DataType::from(&value).unwrap();
- assert_eq!(DataType::Int32, dt);
- }
-
- #[test]
- fn schema_json() {
- // Add some custom metadata
- let metadata: HashMap<String, String> =
- [("Key".to_string(), "Value".to_string())]
- .iter()
- .cloned()
- .collect();
-
- let schema = Schema::new_with_metadata(
- vec![
- Field::new("c1", DataType::Utf8, false),
- Field::new("c2", DataType::Binary, false),
- Field::new("c3", DataType::FixedSizeBinary(3), false),
- Field::new("c4", DataType::Boolean, false),
- Field::new("c5", DataType::Date32, false),
- Field::new("c6", DataType::Date64, false),
- Field::new("c7", DataType::Time32(TimeUnit::Second), false),
- Field::new("c8", DataType::Time32(TimeUnit::Millisecond), false),
- Field::new("c9", DataType::Time32(TimeUnit::Microsecond), false),
- Field::new("c10", DataType::Time32(TimeUnit::Nanosecond), false),
- Field::new("c11", DataType::Time64(TimeUnit::Second), false),
- Field::new("c12", DataType::Time64(TimeUnit::Millisecond), false),
- Field::new("c13", DataType::Time64(TimeUnit::Microsecond), false),
- Field::new("c14", DataType::Time64(TimeUnit::Nanosecond), false),
- Field::new("c15", DataType::Timestamp(TimeUnit::Second, None), false),
- Field::new(
- "c16",
- DataType::Timestamp(TimeUnit::Millisecond, Some("UTC".to_string())),
- false,
- ),
- Field::new(
- "c17",
- DataType::Timestamp(
- TimeUnit::Microsecond,
- Some("Africa/Johannesburg".to_string()),
- ),
- false,
- ),
- Field::new(
- "c18",
- DataType::Timestamp(TimeUnit::Nanosecond, None),
- false,
- ),
- Field::new("c19", DataType::Interval(IntervalUnit::DayTime), false),
- Field::new("c20", DataType::Interval(IntervalUnit::YearMonth), false),
- Field::new(
- "c21",
- DataType::List(Box::new(Field::new("item", DataType::Boolean, true))),
- false,
- ),
- Field::new(
- "c22",
- DataType::FixedSizeList(
- Box::new(Field::new("bools", DataType::Boolean, false)),
- 5,
- ),
- false,
- ),
- Field::new(
- "c23",
- DataType::List(Box::new(Field::new(
- "inner_list",
- DataType::List(Box::new(Field::new(
- "struct",
- DataType::Struct(vec![]),
- true,
- ))),
- false,
- ))),
- true,
- ),
- Field::new(
- "c24",
- DataType::Struct(vec![
- Field::new("a", DataType::Utf8, false),
- Field::new("b", DataType::UInt16, false),
- ]),
- false,
- ),
- Field::new("c25", DataType::Interval(IntervalUnit::YearMonth), true),
- Field::new("c26", DataType::Interval(IntervalUnit::DayTime), true),
- Field::new("c27", DataType::Duration(TimeUnit::Second), false),
- Field::new("c28", DataType::Duration(TimeUnit::Millisecond), false),
- Field::new("c29", DataType::Duration(TimeUnit::Microsecond), false),
- Field::new("c30", DataType::Duration(TimeUnit::Nanosecond), false),
- Field::new_dict(
- "c31",
- DataType::Dictionary(
- Box::new(DataType::Int32),
- Box::new(DataType::Utf8),
- ),
- true,
- 123,
- true,
- ),
- Field::new("c32", DataType::LargeBinary, true),
- Field::new("c33", DataType::LargeUtf8, true),
- Field::new(
- "c34",
- DataType::LargeList(Box::new(Field::new(
- "inner_large_list",
- DataType::LargeList(Box::new(Field::new(
- "struct",
- DataType::Struct(vec![]),
- false,
- ))),
- true,
- ))),
- true,
- ),
- ],
- metadata,
- );
-
- let expected = schema.to_json();
- let json = r#"{
- "fields": [
- {
- "name": "c1",
- "nullable": false,
- "type": {
- "name": "utf8"
- },
- "children": []
- },
- {
- "name": "c2",
- "nullable": false,
- "type": {
- "name": "binary"
- },
- "children": []
- },
- {
- "name": "c3",
- "nullable": false,
- "type": {
- "name": "fixedsizebinary",
- "byteWidth": 3
- },
- "children": []
- },
- {
- "name": "c4",
- "nullable": false,
- "type": {
- "name": "bool"
- },
- "children": []
- },
- {
- "name": "c5",
- "nullable": false,
- "type": {
- "name": "date",
- "unit": "DAY"
- },
- "children": []
- },
- {
- "name": "c6",
- "nullable": false,
- "type": {
- "name": "date",
- "unit": "MILLISECOND"
- },
- "children": []
- },
- {
- "name": "c7",
- "nullable": false,
- "type": {
- "name": "time",
- "bitWidth": 32,
- "unit": "SECOND"
- },
- "children": []
- },
- {
- "name": "c8",
- "nullable": false,
- "type": {
- "name": "time",
- "bitWidth": 32,
- "unit": "MILLISECOND"
- },
- "children": []
- },
- {
- "name": "c9",
- "nullable": false,
- "type": {
- "name": "time",
- "bitWidth": 32,
- "unit": "MICROSECOND"
- },
- "children": []
- },
- {
- "name": "c10",
- "nullable": false,
- "type": {
- "name": "time",
- "bitWidth": 32,
- "unit": "NANOSECOND"
- },
- "children": []
- },
- {
- "name": "c11",
- "nullable": false,
- "type": {
- "name": "time",
- "bitWidth": 64,
- "unit": "SECOND"
- },
- "children": []
- },
- {
- "name": "c12",
- "nullable": false,
- "type": {
- "name": "time",
- "bitWidth": 64,
- "unit": "MILLISECOND"
- },
- "children": []
- },
- {
- "name": "c13",
- "nullable": false,
- "type": {
- "name": "time",
- "bitWidth": 64,
- "unit": "MICROSECOND"
- },
- "children": []
- },
- {
- "name": "c14",
- "nullable": false,
- "type": {
- "name": "time",
- "bitWidth": 64,
- "unit": "NANOSECOND"
- },
- "children": []
- },
- {
- "name": "c15",
- "nullable": false,
- "type": {
- "name": "timestamp",
- "unit": "SECOND"
- },
- "children": []
- },
- {
- "name": "c16",
- "nullable": false,
- "type": {
- "name": "timestamp",
- "unit": "MILLISECOND",
- "timezone": "UTC"
- },
- "children": []
- },
- {
- "name": "c17",
- "nullable": false,
- "type": {
- "name": "timestamp",
- "unit": "MICROSECOND",
- "timezone": "Africa/Johannesburg"
- },
- "children": []
- },
- {
- "name": "c18",
- "nullable": false,
- "type": {
- "name": "timestamp",
- "unit": "NANOSECOND"
- },
- "children": []
- },
- {
- "name": "c19",
- "nullable": false,
- "type": {
- "name": "interval",
- "unit": "DAY_TIME"
- },
- "children": []
- },
- {
- "name": "c20",
- "nullable": false,
- "type": {
- "name": "interval",
- "unit": "YEAR_MONTH"
- },
- "children": []
- },
- {
- "name": "c21",
- "nullable": false,
- "type": {
- "name": "list"
- },
- "children": [
- {
- "name": "item",
- "nullable": true,
- "type": {
- "name": "bool"
- },
- "children": []
- }
- ]
- },
- {
- "name": "c22",
- "nullable": false,
- "type": {
- "name": "fixedsizelist",
- "listSize": 5
- },
- "children": [
- {
- "name": "bools",
- "nullable": false,
- "type": {
- "name": "bool"
- },
- "children": []
- }
- ]
- },
- {
- "name": "c23",
- "nullable": true,
- "type": {
- "name": "list"
- },
- "children": [
- {
- "name": "inner_list",
- "nullable": false,
- "type": {
- "name": "list"
- },
- "children": [
- {
- "name": "struct",
- "nullable": true,
- "type": {
- "name": "struct"
- },
- "children": []
- }
- ]
- }
- ]
- },
- {
- "name": "c24",
- "nullable": false,
- "type": {
- "name": "struct"
- },
- "children": [
- {
- "name": "a",
- "nullable": false,
- "type": {
- "name": "utf8"
- },
- "children": []
- },
- {
- "name": "b",
- "nullable": false,
- "type": {
- "name": "int",
- "bitWidth": 16,
- "isSigned": false
- },
- "children": []
- }
- ]
- },
- {
- "name": "c25",
- "nullable": true,
- "type": {
- "name": "interval",
- "unit": "YEAR_MONTH"
- },
- "children": []
- },
- {
- "name": "c26",
- "nullable": true,
- "type": {
- "name": "interval",
- "unit": "DAY_TIME"
- },
- "children": []
- },
- {
- "name": "c27",
- "nullable": false,
- "type": {
- "name": "duration",
- "unit": "SECOND"
- },
- "children": []
- },
- {
- "name": "c28",
- "nullable": false,
- "type": {
- "name": "duration",
- "unit": "MILLISECOND"
- },
- "children": []
- },
- {
- "name": "c29",
- "nullable": false,
- "type": {
- "name": "duration",
- "unit": "MICROSECOND"
- },
- "children": []
- },
- {
- "name": "c30",
- "nullable": false,
- "type": {
- "name": "duration",
- "unit": "NANOSECOND"
- },
- "children": []
- },
- {
- "name": "c31",
- "nullable": true,
- "children": [],
- "type": {
- "name": "utf8"
- },
- "dictionary": {
- "id": 123,
- "indexType": {
- "name": "int",
- "bitWidth": 32,
- "isSigned": true
- },
- "isOrdered": true
- }
- },
- {
- "name": "c32",
- "nullable": true,
- "type": {
- "name": "largebinary"
- },
- "children": []
- },
- {
- "name": "c33",
- "nullable": true,
- "type": {
- "name": "largeutf8"
- },
- "children": []
- },
- {
- "name": "c34",
- "nullable": true,
- "type": {
- "name": "largelist"
- },
- "children": [
- {
- "name": "inner_large_list",
- "nullable": true,
- "type": {
- "name": "largelist"
- },
- "children": [
- {
- "name": "struct",
- "nullable": false,
- "type": {
- "name": "struct"
- },
- "children": []
- }
- ]
- }
- ]
- }
- ],
- "metadata" : {
- "Key": "Value"
- }
- }"#;
- let value: Value = serde_json::from_str(&json).unwrap();
- assert_eq!(expected, value);
-
- // convert back to a schema
- let value: Value = serde_json::from_str(&json).unwrap();
- let schema2 = Schema::from(&value).unwrap();
-
- assert_eq!(schema, schema2);
-
- // Check that empty metadata produces empty value in JSON and can be parsed
- let json = r#"{
- "fields": [
- {
- "name": "c1",
- "nullable": false,
- "type": {
- "name": "utf8"
- },
- "children": []
- }
- ],
- "metadata": {}
- }"#;
- let value: Value = serde_json::from_str(&json).unwrap();
- let schema = Schema::from(&value).unwrap();
- assert!(schema.metadata.is_empty());
-
- // Check that metadata field is not required in the JSON.
- let json = r#"{
- "fields": [
- {
- "name": "c1",
- "nullable": false,
- "type": {
- "name": "utf8"
- },
- "children": []
- }
- ]
- }"#;
- let value: Value = serde_json::from_str(&json).unwrap();
- let schema = Schema::from(&value).unwrap();
- assert!(schema.metadata.is_empty());
- }
-
- #[test]
- fn create_schema_string() {
- let schema = person_schema();
- assert_eq!(schema.to_string(),
- "Field { name: \"first_name\", data_type: Utf8, nullable: false, dict_id: 0, dict_is_ordered: false, metadata: Some({\"k\": \"v\"}) }, \
- Field { name: \"last_name\", data_type: Utf8, nullable: false, dict_id: 0, dict_is_ordered: false, metadata: None }, \
- Field { name: \"address\", data_type: Struct([\
- Field { name: \"street\", data_type: Utf8, nullable: false, dict_id: 0, dict_is_ordered: false, metadata: None }, \
- Field { name: \"zip\", data_type: UInt16, nullable: false, dict_id: 0, dict_is_ordered: false, metadata: None }\
- ]), nullable: false, dict_id: 0, dict_is_ordered: false, metadata: None }, \
- Field { name: \"interests\", data_type: Dictionary(Int32, Utf8), nullable: true, dict_id: 123, dict_is_ordered: true, metadata: None }")
- }
-
- #[test]
- fn schema_field_accessors() {
- let schema = person_schema();
-
- // test schema accessors
- assert_eq!(schema.fields().len(), 4);
-
- // test field accessors
- let first_name = &schema.fields()[0];
- assert_eq!(first_name.name(), "first_name");
- assert_eq!(first_name.data_type(), &DataType::Utf8);
- assert_eq!(first_name.is_nullable(), false);
- assert_eq!(first_name.dict_id(), None);
- assert_eq!(first_name.dict_is_ordered(), None);
-
- let metadata = first_name.metadata();
- assert!(metadata.is_some());
- let md = metadata.as_ref().unwrap();
- assert_eq!(md.len(), 1);
- let key = md.get("k");
- assert!(key.is_some());
- assert_eq!(key.unwrap(), "v");
-
- let interests = &schema.fields()[3];
- assert_eq!(interests.name(), "interests");
- assert_eq!(
- interests.data_type(),
- &DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8))
- );
- assert_eq!(interests.dict_id(), Some(123));
- assert_eq!(interests.dict_is_ordered(), Some(true));
- }
-
- #[test]
- #[should_panic(
- expected = "Unable to get field named \\\"nickname\\\". Valid fields: [\\\"first_name\\\", \\\"last_name\\\", \\\"address\\\", \\\"interests\\\"]"
- )]
- fn schema_index_of() {
- let schema = person_schema();
- assert_eq!(schema.index_of("first_name").unwrap(), 0);
- assert_eq!(schema.index_of("last_name").unwrap(), 1);
- schema.index_of("nickname").unwrap();
- }
-
- #[test]
- #[should_panic(
- expected = "Unable to get field named \\\"nickname\\\". Valid fields: [\\\"first_name\\\", \\\"last_name\\\", \\\"address\\\", \\\"interests\\\"]"
- )]
- fn schema_field_with_name() {
- let schema = person_schema();
- assert_eq!(
- schema.field_with_name("first_name").unwrap().name(),
- "first_name"
- );
- assert_eq!(
- schema.field_with_name("last_name").unwrap().name(),
- "last_name"
- );
- schema.field_with_name("nickname").unwrap();
- }
-
- #[test]
- fn schema_field_with_dict_id() {
- let schema = person_schema();
-
- let fields_dict_123: Vec<_> = schema
- .fields_with_dict_id(123)
- .iter()
- .map(|f| f.name())
- .collect();
- assert_eq!(fields_dict_123, vec!["interests"]);
-
- assert!(schema.fields_with_dict_id(456).is_empty());
- }
-
- #[test]
- fn schema_equality() {
- let schema1 = Schema::new(vec![
- Field::new("c1", DataType::Utf8, false),
- Field::new("c2", DataType::Float64, true),
- Field::new("c3", DataType::LargeBinary, true),
- ]);
- let schema2 = Schema::new(vec![
- Field::new("c1", DataType::Utf8, false),
- Field::new("c2", DataType::Float64, true),
- Field::new("c3", DataType::LargeBinary, true),
- ]);
-
- assert_eq!(schema1, schema2);
-
- let schema3 = Schema::new(vec![
- Field::new("c1", DataType::Utf8, false),
- Field::new("c2", DataType::Float32, true),
- ]);
- let schema4 = Schema::new(vec![
- Field::new("C1", DataType::Utf8, false),
- Field::new("C2", DataType::Float64, true),
- ]);
-
- assert!(schema1 != schema3);
- assert!(schema1 != schema4);
- assert!(schema2 != schema3);
- assert!(schema2 != schema4);
- assert!(schema3 != schema4);
-
- let mut f = Field::new("c1", DataType::Utf8, false);
- f.set_metadata(Some(
- [("foo".to_string(), "bar".to_string())]
- .iter()
- .cloned()
- .collect(),
- ));
- let schema5 = Schema::new(vec![
- f,
- Field::new("c2", DataType::Float64, true),
- Field::new("c3", DataType::LargeBinary, true),
- ]);
- assert!(schema1 != schema5);
- }
-
- #[test]
- fn test_arrow_native_type_to_json() {
- assert_eq!(Some(Bool(true)), true.into_json_value());
- assert_eq!(Some(VNumber(Number::from(1))), 1i8.into_json_value());
- assert_eq!(Some(VNumber(Number::from(1))), 1i16.into_json_value());
- assert_eq!(Some(VNumber(Number::from(1))), 1i32.into_json_value());
- assert_eq!(Some(VNumber(Number::from(1))), 1i64.into_json_value());
- assert_eq!(Some(VNumber(Number::from(1))), 1u8.into_json_value());
- assert_eq!(Some(VNumber(Number::from(1))), 1u16.into_json_value());
- assert_eq!(Some(VNumber(Number::from(1))), 1u32.into_json_value());
- assert_eq!(Some(VNumber(Number::from(1))), 1u64.into_json_value());
- assert_eq!(
- Some(VNumber(Number::from_f64(0.01f64).unwrap())),
- 0.01.into_json_value()
- );
- assert_eq!(
- Some(VNumber(Number::from_f64(0.01f64).unwrap())),
- 0.01f64.into_json_value()
- );
- assert_eq!(None, NAN.into_json_value());
- }
-
- fn person_schema() -> Schema {
- let kv_array = [("k".to_string(), "v".to_string())];
- let field_metadata: BTreeMap<String, String> = kv_array.iter().cloned().collect();
- let mut first_name = Field::new("first_name", DataType::Utf8, false);
- first_name.set_metadata(Some(field_metadata));
-
- Schema::new(vec![
- first_name,
- Field::new("last_name", DataType::Utf8, false),
- Field::new(
- "address",
- DataType::Struct(vec![
- Field::new("street", DataType::Utf8, false),
- Field::new("zip", DataType::UInt16, false),
- ]),
- false,
- ),
- Field::new_dict(
- "interests",
- DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
- true,
- 123,
- true,
- ),
- ])
- }
-
- #[test]
- fn test_try_merge_field_with_metadata() {
- // 1. Different values for the same key should cause error.
- let metadata1: BTreeMap<String, String> =
- [("foo".to_string(), "bar".to_string())]
- .iter()
- .cloned()
- .collect();
- let mut f1 = Field::new("first_name", DataType::Utf8, false);
- f1.set_metadata(Some(metadata1));
-
- let metadata2: BTreeMap<String, String> =
- [("foo".to_string(), "baz".to_string())]
- .iter()
- .cloned()
- .collect();
- let mut f2 = Field::new("first_name", DataType::Utf8, false);
- f2.set_metadata(Some(metadata2));
-
- assert!(
- Schema::try_merge(vec![Schema::new(vec![f1]), Schema::new(vec![f2])])
- .is_err()
- );
-
- // 2. None + Some
- let mut f1 = Field::new("first_name", DataType::Utf8, false);
- let metadata2: BTreeMap<String, String> =
- [("missing".to_string(), "value".to_string())]
- .iter()
- .cloned()
- .collect();
- let mut f2 = Field::new("first_name", DataType::Utf8, false);
- f2.set_metadata(Some(metadata2));
-
- assert!(f1.try_merge(&f2).is_ok());
- assert!(f1.metadata().is_some());
- assert_eq!(
- f1.metadata().as_ref().unwrap(),
- f2.metadata().as_ref().unwrap()
- );
-
- // 3. Some + Some
- let mut f1 = Field::new("first_name", DataType::Utf8, false);
- f1.set_metadata(Some(
- [("foo".to_string(), "bar".to_string())]
- .iter()
- .cloned()
- .collect(),
- ));
- let mut f2 = Field::new("first_name", DataType::Utf8, false);
- f2.set_metadata(Some(
- [("foo2".to_string(), "bar2".to_string())]
- .iter()
- .cloned()
- .collect(),
- ));
-
- assert!(f1.try_merge(&f2).is_ok());
- assert!(f1.metadata().is_some());
- assert_eq!(
- f1.metadata().clone().unwrap(),
- [
- ("foo".to_string(), "bar".to_string()),
- ("foo2".to_string(), "bar2".to_string())
- ]
- .iter()
- .cloned()
- .collect()
- );
-
- // 4. Some + None.
- let mut f1 = Field::new("first_name", DataType::Utf8, false);
- f1.set_metadata(Some(
- [("foo".to_string(), "bar".to_string())]
- .iter()
- .cloned()
- .collect(),
- ));
- let f2 = Field::new("first_name", DataType::Utf8, false);
- assert!(f1.try_merge(&f2).is_ok());
- assert!(f1.metadata().is_some());
- assert_eq!(
- f1.metadata().clone().unwrap(),
- [("foo".to_string(), "bar".to_string())]
- .iter()
- .cloned()
- .collect()
- );
-
- // 5. None + None.
- let mut f1 = Field::new("first_name", DataType::Utf8, false);
- let f2 = Field::new("first_name", DataType::Utf8, false);
- assert!(f1.try_merge(&f2).is_ok());
- assert!(f1.metadata().is_none());
- }
-
- #[test]
- fn test_schema_merge() -> Result<()> {
- let merged = Schema::try_merge(vec![
- Schema::new(vec![
- Field::new("first_name", DataType::Utf8, false),
- Field::new("last_name", DataType::Utf8, false),
- Field::new(
- "address",
- DataType::Struct(vec![Field::new("zip", DataType::UInt16, false)]),
- false,
- ),
- ]),
- Schema::new_with_metadata(
- vec![
- // nullable merge
- Field::new("last_name", DataType::Utf8, true),
- Field::new(
- "address",
- DataType::Struct(vec![
- // add new nested field
- Field::new("street", DataType::Utf8, false),
- // nullable merge on nested field
- Field::new("zip", DataType::UInt16, true),
- ]),
- false,
- ),
- // new field
- Field::new("number", DataType::Utf8, true),
- ],
- [("foo".to_string(), "bar".to_string())]
- .iter()
- .cloned()
- .collect::<HashMap<String, String>>(),
- ),
- ])?;
-
- assert_eq!(
- merged,
- Schema::new_with_metadata(
- vec![
- Field::new("first_name", DataType::Utf8, false),
- Field::new("last_name", DataType::Utf8, true),
- Field::new(
- "address",
- DataType::Struct(vec![
- Field::new("zip", DataType::UInt16, true),
- Field::new("street", DataType::Utf8, false),
- ]),
- false,
- ),
- Field::new("number", DataType::Utf8, true),
- ],
- [("foo".to_string(), "bar".to_string())]
- .iter()
- .cloned()
- .collect::<HashMap<String, String>>()
- )
- );
-
- // support merge union fields
- assert_eq!(
- Schema::try_merge(vec![
- Schema::new(vec![Field::new(
- "c1",
- DataType::Union(vec![
- Field::new("c11", DataType::Utf8, true),
- Field::new("c12", DataType::Utf8, true),
- ]),
- false
- ),]),
- Schema::new(vec![Field::new(
- "c1",
- DataType::Union(vec![
- Field::new("c12", DataType::Utf8, true),
- Field::new("c13", DataType::Time64(TimeUnit::Second), true),
- ]),
- false
- ),])
- ])?,
- Schema::new(vec![Field::new(
- "c1",
- DataType::Union(vec![
- Field::new("c11", DataType::Utf8, true),
- Field::new("c12", DataType::Utf8, true),
- Field::new("c13", DataType::Time64(TimeUnit::Second), true),
- ]),
- false
- ),]),
- );
-
- // incompatible field should throw error
- assert!(Schema::try_merge(vec![
- Schema::new(vec![
- Field::new("first_name", DataType::Utf8, false),
- Field::new("last_name", DataType::Utf8, false),
- ]),
- Schema::new(vec![Field::new("last_name", DataType::Int64, false),])
- ])
- .is_err());
-
- // incompatible metadata should throw error
- assert!(Schema::try_merge(vec![
- Schema::new_with_metadata(
- vec![Field::new("first_name", DataType::Utf8, false)],
- [("foo".to_string(), "bar".to_string()),]
- .iter()
- .cloned()
- .collect::<HashMap<String, String>>()
- ),
- Schema::new_with_metadata(
- vec![Field::new("last_name", DataType::Utf8, false)],
- [("foo".to_string(), "baz".to_string()),]
- .iter()
- .cloned()
- .collect::<HashMap<String, String>>()
- )
- ])
- .is_err());
-
- Ok(())
- }
-}
diff --git a/arrow/src/datatypes/native.rs b/arrow/src/datatypes/native.rs
deleted file mode 100644
index 6e8cf89..0000000
--- a/arrow/src/datatypes/native.rs
+++ /dev/null
@@ -1,333 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-use serde_json::{Number, Value};
-
-use super::DataType;
-
-/// Trait declaring any type that is serializable to JSON. This includes all primitive types (bool, i32, etc.).
-pub trait JsonSerializable: 'static {
- fn into_json_value(self) -> Option<Value>;
-}
-
-/// Trait expressing a Rust type that has the same in-memory representation
-/// as Arrow. This includes `i16`, `f32`, but excludes `bool` (which in arrow is represented in bits).
-/// In little endian machines, types that implement [`ArrowNativeType`] can be memcopied to arrow buffers
-/// as is.
-pub trait ArrowNativeType:
- std::fmt::Debug
- + Send
- + Sync
- + Copy
- + PartialOrd
- + std::str::FromStr
- + Default
- + JsonSerializable
-{
- /// Convert native type from usize.
- #[inline]
- fn from_usize(_: usize) -> Option<Self> {
- None
- }
-
- /// Convert native type to usize.
- #[inline]
- fn to_usize(&self) -> Option<usize> {
- None
- }
-
- /// Convert native type to isize.
- #[inline]
- fn to_isize(&self) -> Option<isize> {
- None
- }
-
- /// Convert native type from i32.
- #[inline]
- fn from_i32(_: i32) -> Option<Self> {
- None
- }
-
- /// Convert native type from i64.
- #[inline]
- fn from_i64(_: i64) -> Option<Self> {
- None
- }
-}
-
-/// Trait bridging the dynamic-typed nature of Arrow (via [`DataType`]) with the
-/// static-typed nature of rust types ([`ArrowNativeType`]) for all types that implement [`ArrowNativeType`].
-pub trait ArrowPrimitiveType: 'static {
- /// Corresponding Rust native type for the primitive type.
- type Native: ArrowNativeType;
-
- /// the corresponding Arrow data type of this primitive type.
- const DATA_TYPE: DataType;
-
- /// Returns the byte width of this primitive type.
- fn get_byte_width() -> usize {
- std::mem::size_of::<Self::Native>()
- }
-
- /// Returns a default value of this primitive type.
- ///
- /// This is useful for aggregate array ops like `sum()`, `mean()`.
- fn default_value() -> Self::Native {
- Default::default()
- }
-}
-
-impl JsonSerializable for bool {
- fn into_json_value(self) -> Option<Value> {
- Some(self.into())
- }
-}
-
-impl JsonSerializable for i8 {
- fn into_json_value(self) -> Option<Value> {
- Some(self.into())
- }
-}
-
-impl ArrowNativeType for i8 {
- #[inline]
- fn from_usize(v: usize) -> Option<Self> {
- num::FromPrimitive::from_usize(v)
- }
-
- #[inline]
- fn to_usize(&self) -> Option<usize> {
- num::ToPrimitive::to_usize(self)
- }
-
- #[inline]
- fn to_isize(&self) -> Option<isize> {
- num::ToPrimitive::to_isize(self)
- }
-}
-
-impl JsonSerializable for i16 {
- fn into_json_value(self) -> Option<Value> {
- Some(self.into())
- }
-}
-
-impl ArrowNativeType for i16 {
- #[inline]
- fn from_usize(v: usize) -> Option<Self> {
- num::FromPrimitive::from_usize(v)
- }
-
- #[inline]
- fn to_usize(&self) -> Option<usize> {
- num::ToPrimitive::to_usize(self)
- }
-
- #[inline]
- fn to_isize(&self) -> Option<isize> {
- num::ToPrimitive::to_isize(self)
- }
-}
-
-impl JsonSerializable for i32 {
- fn into_json_value(self) -> Option<Value> {
- Some(self.into())
- }
-}
-
-impl ArrowNativeType for i32 {
- #[inline]
- fn from_usize(v: usize) -> Option<Self> {
- num::FromPrimitive::from_usize(v)
- }
-
- #[inline]
- fn to_usize(&self) -> Option<usize> {
- num::ToPrimitive::to_usize(self)
- }
-
- #[inline]
- fn to_isize(&self) -> Option<isize> {
- num::ToPrimitive::to_isize(self)
- }
-
- /// Convert native type from i32.
- #[inline]
- fn from_i32(val: i32) -> Option<Self> {
- Some(val)
- }
-}
-
-impl JsonSerializable for i64 {
- fn into_json_value(self) -> Option<Value> {
- Some(Value::Number(Number::from(self)))
- }
-}
-
-impl ArrowNativeType for i64 {
- #[inline]
- fn from_usize(v: usize) -> Option<Self> {
- num::FromPrimitive::from_usize(v)
- }
-
- #[inline]
- fn to_usize(&self) -> Option<usize> {
- num::ToPrimitive::to_usize(self)
- }
-
- #[inline]
- fn to_isize(&self) -> Option<isize> {
- num::ToPrimitive::to_isize(self)
- }
-
- /// Convert native type from i64.
- #[inline]
- fn from_i64(val: i64) -> Option<Self> {
- Some(val)
- }
-}
-
-impl JsonSerializable for u8 {
- fn into_json_value(self) -> Option<Value> {
- Some(self.into())
- }
-}
-
-impl ArrowNativeType for u8 {
- #[inline]
- fn from_usize(v: usize) -> Option<Self> {
- num::FromPrimitive::from_usize(v)
- }
-
- #[inline]
- fn to_usize(&self) -> Option<usize> {
- num::ToPrimitive::to_usize(self)
- }
-
- #[inline]
- fn to_isize(&self) -> Option<isize> {
- num::ToPrimitive::to_isize(self)
- }
-}
-
-impl JsonSerializable for u16 {
- fn into_json_value(self) -> Option<Value> {
- Some(self.into())
- }
-}
-
-impl ArrowNativeType for u16 {
- #[inline]
- fn from_usize(v: usize) -> Option<Self> {
- num::FromPrimitive::from_usize(v)
- }
-
- #[inline]
- fn to_usize(&self) -> Option<usize> {
- num::ToPrimitive::to_usize(self)
- }
-
- #[inline]
- fn to_isize(&self) -> Option<isize> {
- num::ToPrimitive::to_isize(self)
- }
-}
-
-impl JsonSerializable for u32 {
- fn into_json_value(self) -> Option<Value> {
- Some(self.into())
- }
-}
-
-impl ArrowNativeType for u32 {
- #[inline]
- fn from_usize(v: usize) -> Option<Self> {
- num::FromPrimitive::from_usize(v)
- }
-
- #[inline]
- fn to_usize(&self) -> Option<usize> {
- num::ToPrimitive::to_usize(self)
- }
-
- #[inline]
- fn to_isize(&self) -> Option<isize> {
- num::ToPrimitive::to_isize(self)
- }
-}
-
-impl JsonSerializable for u64 {
- fn into_json_value(self) -> Option<Value> {
- Some(self.into())
- }
-}
-
-impl ArrowNativeType for u64 {
- #[inline]
- fn from_usize(v: usize) -> Option<Self> {
- num::FromPrimitive::from_usize(v)
- }
-
- #[inline]
- fn to_usize(&self) -> Option<usize> {
- num::ToPrimitive::to_usize(self)
- }
-
- #[inline]
- fn to_isize(&self) -> Option<isize> {
- num::ToPrimitive::to_isize(self)
- }
-}
-
-impl JsonSerializable for f32 {
- fn into_json_value(self) -> Option<Value> {
- Number::from_f64(f64::round(self as f64 * 1000.0) / 1000.0).map(Value::Number)
- }
-}
-
-impl JsonSerializable for f64 {
- fn into_json_value(self) -> Option<Value> {
- Number::from_f64(self).map(Value::Number)
- }
-}
-
-impl ArrowNativeType for f32 {}
-impl ArrowNativeType for f64 {}
-
-/// Allows conversion from supported Arrow types to a byte slice.
-pub trait ToByteSlice {
- /// Converts this instance into a byte slice
- fn to_byte_slice(&self) -> &[u8];
-}
-
-impl<T: ArrowNativeType> ToByteSlice for [T] {
- #[inline]
- fn to_byte_slice(&self) -> &[u8] {
- let raw_ptr = self.as_ptr() as *const T as *const u8;
- unsafe {
- std::slice::from_raw_parts(raw_ptr, self.len() * std::mem::size_of::<T>())
- }
- }
-}
-
-impl<T: ArrowNativeType> ToByteSlice for T {
- #[inline]
- fn to_byte_slice(&self) -> &[u8] {
- let raw_ptr = self as *const T as *const u8;
- unsafe { std::slice::from_raw_parts(raw_ptr, std::mem::size_of::<T>()) }
- }
-}
diff --git a/arrow/src/datatypes/numeric.rs b/arrow/src/datatypes/numeric.rs
deleted file mode 100644
index 39c6732..0000000
--- a/arrow/src/datatypes/numeric.rs
+++ /dev/null
@@ -1,536 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-use super::*;
-#[cfg(feature = "simd")]
-use packed_simd::*;
-#[cfg(feature = "simd")]
-use std::ops::{
- Add, BitAnd, BitAndAssign, BitOr, BitOrAssign, Div, Mul, Neg, Not, Rem, Sub,
-};
-
-/// A subtype of primitive type that represents numeric values.
-///
-/// SIMD operations are defined in this trait if available on the target system.
-#[cfg(feature = "simd")]
-pub trait ArrowNumericType: ArrowPrimitiveType
-where
- Self::Simd: Add<Output = Self::Simd>
- + Sub<Output = Self::Simd>
- + Mul<Output = Self::Simd>
- + Div<Output = Self::Simd>
- + Rem<Output = Self::Simd>
- + Copy,
- Self::SimdMask: BitAnd<Output = Self::SimdMask>
- + BitOr<Output = Self::SimdMask>
- + BitAndAssign
- + BitOrAssign
- + Not<Output = Self::SimdMask>
- + Copy,
-{
- /// Defines the SIMD type that should be used for this numeric type
- type Simd;
-
- /// Defines the SIMD Mask type that should be used for this numeric type
- type SimdMask;
-
- /// The number of SIMD lanes available
- fn lanes() -> usize;
-
- /// Initializes a SIMD register to a constant value
- fn init(value: Self::Native) -> Self::Simd;
-
- /// Loads a slice into a SIMD register
- fn load(slice: &[Self::Native]) -> Self::Simd;
-
- /// Creates a new SIMD mask for this SIMD type filling it with `value`
- fn mask_init(value: bool) -> Self::SimdMask;
-
- /// Creates a new SIMD mask for this SIMD type from the lower-most bits of the given `mask`.
- /// The number of bits used corresponds to the number of lanes of this type
- fn mask_from_u64(mask: u64) -> Self::SimdMask;
-
- /// Creates a bitmask from the given SIMD mask.
- /// Each bit corresponds to one vector lane, starting with the least-significant bit.
- fn mask_to_u64(mask: &Self::SimdMask) -> u64;
-
- /// Gets the value of a single lane in a SIMD mask
- fn mask_get(mask: &Self::SimdMask, idx: usize) -> bool;
-
- /// Sets the value of a single lane of a SIMD mask
- fn mask_set(mask: Self::SimdMask, idx: usize, value: bool) -> Self::SimdMask;
-
- /// Selects elements of `a` and `b` using `mask`
- fn mask_select(mask: Self::SimdMask, a: Self::Simd, b: Self::Simd) -> Self::Simd;
-
- /// Returns `true` if any of the lanes in the mask are `true`
- fn mask_any(mask: Self::SimdMask) -> bool;
-
- /// Performs a SIMD binary operation
- fn bin_op<F: Fn(Self::Simd, Self::Simd) -> Self::Simd>(
- left: Self::Simd,
- right: Self::Simd,
- op: F,
- ) -> Self::Simd;
-
- /// SIMD version of equal
- fn eq(left: Self::Simd, right: Self::Simd) -> Self::SimdMask;
-
- /// SIMD version of not equal
- fn ne(left: Self::Simd, right: Self::Simd) -> Self::SimdMask;
-
- /// SIMD version of less than
- fn lt(left: Self::Simd, right: Self::Simd) -> Self::SimdMask;
-
- /// SIMD version of less than or equal to
- fn le(left: Self::Simd, right: Self::Simd) -> Self::SimdMask;
-
- /// SIMD version of greater than
- fn gt(left: Self::Simd, right: Self::Simd) -> Self::SimdMask;
-
- /// SIMD version of greater than or equal to
- fn ge(left: Self::Simd, right: Self::Simd) -> Self::SimdMask;
-
- /// Writes a SIMD result back to a slice
- fn write(simd_result: Self::Simd, slice: &mut [Self::Native]);
-
- fn unary_op<F: Fn(Self::Simd) -> Self::Simd>(a: Self::Simd, op: F) -> Self::Simd;
-}
-
-#[cfg(not(feature = "simd"))]
-pub trait ArrowNumericType: ArrowPrimitiveType {}
-
-macro_rules! make_numeric_type {
- ($impl_ty:ty, $native_ty:ty, $simd_ty:ident, $simd_mask_ty:ident) => {
- #[cfg(feature = "simd")]
- impl ArrowNumericType for $impl_ty {
- type Simd = $simd_ty;
-
- type SimdMask = $simd_mask_ty;
-
- #[inline]
- fn lanes() -> usize {
- Self::Simd::lanes()
- }
-
- #[inline]
- fn init(value: Self::Native) -> Self::Simd {
- Self::Simd::splat(value)
- }
-
- #[inline]
- fn load(slice: &[Self::Native]) -> Self::Simd {
- unsafe { Self::Simd::from_slice_unaligned_unchecked(slice) }
- }
-
- #[inline]
- fn mask_init(value: bool) -> Self::SimdMask {
- Self::SimdMask::splat(value)
- }
-
- #[inline]
- fn mask_from_u64(mask: u64) -> Self::SimdMask {
- // this match will get removed by the compiler since the number of lanes is known at
- // compile-time for each concrete numeric type
- match Self::lanes() {
- 8 => {
- // the bit position in each lane indicates the index of that lane
- let vecidx = i64x8::new(1, 2, 4, 8, 16, 32, 64, 128);
-
- // broadcast the lowermost 8 bits of mask to each lane
- let vecmask = i64x8::splat((mask & 0xFF) as i64);
- // compute whether the bit corresponding to each lanes index is set
- let vecmask = (vecidx & vecmask).eq(vecidx);
-
- // transmute is necessary because the different match arms return different
- // mask types, at runtime only one of those expressions will exist per type,
- // with the type being equal to `SimdMask`.
- unsafe { std::mem::transmute(vecmask) }
- }
- 16 => {
- // same general logic as for 8 lanes, extended to 16 bits
- let vecidx = i32x16::new(
- 1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096,
- 8192, 16384, 32768,
- );
-
- let vecmask = i32x16::splat((mask & 0xFFFF) as i32);
- let vecmask = (vecidx & vecmask).eq(vecidx);
-
- unsafe { std::mem::transmute(vecmask) }
- }
- 32 => {
- // compute two separate m32x16 vector masks from from the lower-most 32 bits of `mask`
- // and then combine them into one m16x32 vector mask by writing and reading a temporary
- let tmp = &mut [0_i16; 32];
-
- let vecidx = i32x16::new(
- 1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096,
- 8192, 16384, 32768,
- );
-
- let vecmask = i32x16::splat((mask & 0xFFFF) as i32);
- let vecmask = (vecidx & vecmask).eq(vecidx);
-
- i16x16::from_cast(vecmask)
- .write_to_slice_unaligned(&mut tmp[0..16]);
-
- let vecmask = i32x16::splat(((mask >> 16) & 0xFFFF) as i32);
- let vecmask = (vecidx & vecmask).eq(vecidx);
-
- i16x16::from_cast(vecmask)
- .write_to_slice_unaligned(&mut tmp[16..32]);
-
- unsafe { std::mem::transmute(i16x32::from_slice_unaligned(tmp)) }
- }
- 64 => {
- // compute four m32x16 vector masks from from all 64 bits of `mask`
- // and convert them into one m8x64 vector mask by writing and reading a temporary
- let tmp = &mut [0_i8; 64];
-
- let vecidx = i32x16::new(
- 1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096,
- 8192, 16384, 32768,
- );
-
- let vecmask = i32x16::splat((mask & 0xFFFF) as i32);
- let vecmask = (vecidx & vecmask).eq(vecidx);
-
- i8x16::from_cast(vecmask)
- .write_to_slice_unaligned(&mut tmp[0..16]);
-
- let vecmask = i32x16::splat(((mask >> 16) & 0xFFFF) as i32);
- let vecmask = (vecidx & vecmask).eq(vecidx);
-
- i8x16::from_cast(vecmask)
- .write_to_slice_unaligned(&mut tmp[16..32]);
-
- let vecmask = i32x16::splat(((mask >> 32) & 0xFFFF) as i32);
- let vecmask = (vecidx & vecmask).eq(vecidx);
-
- i8x16::from_cast(vecmask)
- .write_to_slice_unaligned(&mut tmp[32..48]);
-
- let vecmask = i32x16::splat(((mask >> 48) & 0xFFFF) as i32);
- let vecmask = (vecidx & vecmask).eq(vecidx);
-
- i8x16::from_cast(vecmask)
- .write_to_slice_unaligned(&mut tmp[48..64]);
-
- unsafe { std::mem::transmute(i8x64::from_slice_unaligned(tmp)) }
- }
- _ => panic!("Invalid number of vector lanes"),
- }
- }
-
- #[inline]
- fn mask_to_u64(mask: &Self::SimdMask) -> u64 {
- mask.bitmask() as u64
- }
-
- #[inline]
- fn mask_get(mask: &Self::SimdMask, idx: usize) -> bool {
- unsafe { mask.extract_unchecked(idx) }
- }
-
- #[inline]
- fn mask_set(mask: Self::SimdMask, idx: usize, value: bool) -> Self::SimdMask {
- unsafe { mask.replace_unchecked(idx, value) }
- }
-
- /// Selects elements of `a` and `b` using `mask`
- #[inline]
- fn mask_select(
- mask: Self::SimdMask,
- a: Self::Simd,
- b: Self::Simd,
- ) -> Self::Simd {
- mask.select(a, b)
- }
-
- #[inline]
- fn mask_any(mask: Self::SimdMask) -> bool {
- mask.any()
- }
-
- #[inline]
- fn bin_op<F: Fn(Self::Simd, Self::Simd) -> Self::Simd>(
- left: Self::Simd,
- right: Self::Simd,
- op: F,
- ) -> Self::Simd {
- op(left, right)
- }
-
- #[inline]
- fn eq(left: Self::Simd, right: Self::Simd) -> Self::SimdMask {
- left.eq(right)
- }
-
- #[inline]
- fn ne(left: Self::Simd, right: Self::Simd) -> Self::SimdMask {
- left.ne(right)
- }
-
- #[inline]
- fn lt(left: Self::Simd, right: Self::Simd) -> Self::SimdMask {
- left.lt(right)
- }
-
- #[inline]
- fn le(left: Self::Simd, right: Self::Simd) -> Self::SimdMask {
- left.le(right)
- }
-
- #[inline]
- fn gt(left: Self::Simd, right: Self::Simd) -> Self::SimdMask {
- left.gt(right)
- }
-
- #[inline]
- fn ge(left: Self::Simd, right: Self::Simd) -> Self::SimdMask {
- left.ge(right)
- }
-
- #[inline]
- fn write(simd_result: Self::Simd, slice: &mut [Self::Native]) {
- unsafe { simd_result.write_to_slice_unaligned_unchecked(slice) };
- }
-
- #[inline]
- fn unary_op<F: Fn(Self::Simd) -> Self::Simd>(
- a: Self::Simd,
- op: F,
- ) -> Self::Simd {
- op(a)
- }
- }
-
- #[cfg(not(feature = "simd"))]
- impl ArrowNumericType for $impl_ty {}
- };
-}
-
-make_numeric_type!(Int8Type, i8, i8x64, m8x64);
-make_numeric_type!(Int16Type, i16, i16x32, m16x32);
-make_numeric_type!(Int32Type, i32, i32x16, m32x16);
-make_numeric_type!(Int64Type, i64, i64x8, m64x8);
-make_numeric_type!(UInt8Type, u8, u8x64, m8x64);
-make_numeric_type!(UInt16Type, u16, u16x32, m16x32);
-make_numeric_type!(UInt32Type, u32, u32x16, m32x16);
-make_numeric_type!(UInt64Type, u64, u64x8, m64x8);
-make_numeric_type!(Float32Type, f32, f32x16, m32x16);
-make_numeric_type!(Float64Type, f64, f64x8, m64x8);
-
-make_numeric_type!(TimestampSecondType, i64, i64x8, m64x8);
-make_numeric_type!(TimestampMillisecondType, i64, i64x8, m64x8);
-make_numeric_type!(TimestampMicrosecondType, i64, i64x8, m64x8);
-make_numeric_type!(TimestampNanosecondType, i64, i64x8, m64x8);
-make_numeric_type!(Date32Type, i32, i32x16, m32x16);
-make_numeric_type!(Date64Type, i64, i64x8, m64x8);
-make_numeric_type!(Time32SecondType, i32, i32x16, m32x16);
-make_numeric_type!(Time32MillisecondType, i32, i32x16, m32x16);
-make_numeric_type!(Time64MicrosecondType, i64, i64x8, m64x8);
-make_numeric_type!(Time64NanosecondType, i64, i64x8, m64x8);
-make_numeric_type!(IntervalYearMonthType, i32, i32x16, m32x16);
-make_numeric_type!(IntervalDayTimeType, i64, i64x8, m64x8);
-make_numeric_type!(DurationSecondType, i64, i64x8, m64x8);
-make_numeric_type!(DurationMillisecondType, i64, i64x8, m64x8);
-make_numeric_type!(DurationMicrosecondType, i64, i64x8, m64x8);
-make_numeric_type!(DurationNanosecondType, i64, i64x8, m64x8);
-
-/// A subtype of primitive type that represents signed numeric values.
-///
-/// SIMD operations are defined in this trait if available on the target system.
-#[cfg(feature = "simd")]
-pub trait ArrowSignedNumericType: ArrowNumericType
-where
- Self::SignedSimd: Neg<Output = Self::SignedSimd>,
-{
- /// Defines the SIMD type that should be used for this numeric type
- type SignedSimd;
-
- /// Loads a slice of signed numeric type into a SIMD register
- fn load_signed(slice: &[Self::Native]) -> Self::SignedSimd;
-
- /// Performs a SIMD unary operation on signed numeric type
- fn signed_unary_op<F: Fn(Self::SignedSimd) -> Self::SignedSimd>(
- a: Self::SignedSimd,
- op: F,
- ) -> Self::SignedSimd;
-
- /// Writes a signed SIMD result back to a slice
- fn write_signed(simd_result: Self::SignedSimd, slice: &mut [Self::Native]);
-}
-
-#[cfg(not(feature = "simd"))]
-pub trait ArrowSignedNumericType: ArrowNumericType
-where
- Self::Native: std::ops::Neg<Output = Self::Native>,
-{
-}
-
-macro_rules! make_signed_numeric_type {
- ($impl_ty:ty, $simd_ty:ident) => {
- #[cfg(feature = "simd")]
- impl ArrowSignedNumericType for $impl_ty {
- type SignedSimd = $simd_ty;
-
- #[inline]
- fn load_signed(slice: &[Self::Native]) -> Self::SignedSimd {
- unsafe { Self::SignedSimd::from_slice_unaligned_unchecked(slice) }
- }
-
- #[inline]
- fn signed_unary_op<F: Fn(Self::SignedSimd) -> Self::SignedSimd>(
- a: Self::SignedSimd,
- op: F,
- ) -> Self::SignedSimd {
- op(a)
- }
-
- #[inline]
- fn write_signed(simd_result: Self::SignedSimd, slice: &mut [Self::Native]) {
- unsafe { simd_result.write_to_slice_unaligned_unchecked(slice) };
- }
- }
-
- #[cfg(not(feature = "simd"))]
- impl ArrowSignedNumericType for $impl_ty {}
- };
-}
-
-make_signed_numeric_type!(Int8Type, i8x64);
-make_signed_numeric_type!(Int16Type, i16x32);
-make_signed_numeric_type!(Int32Type, i32x16);
-make_signed_numeric_type!(Int64Type, i64x8);
-make_signed_numeric_type!(Float32Type, f32x16);
-make_signed_numeric_type!(Float64Type, f64x8);
-
-#[cfg(feature = "simd")]
-pub trait ArrowFloatNumericType: ArrowNumericType {
- fn pow(base: Self::Simd, raise: Self::Simd) -> Self::Simd;
-}
-
-#[cfg(not(feature = "simd"))]
-pub trait ArrowFloatNumericType: ArrowNumericType {}
-
-macro_rules! make_float_numeric_type {
- ($impl_ty:ty, $simd_ty:ident) => {
- #[cfg(feature = "simd")]
- impl ArrowFloatNumericType for $impl_ty {
- #[inline]
- fn pow(base: Self::Simd, raise: Self::Simd) -> Self::Simd {
- base.powf(raise)
- }
- }
-
- #[cfg(not(feature = "simd"))]
- impl ArrowFloatNumericType for $impl_ty {}
- };
-}
-
-make_float_numeric_type!(Float32Type, f32x16);
-make_float_numeric_type!(Float64Type, f64x8);
-
-#[cfg(all(test, simd_x86))]
-mod tests {
- use crate::datatypes::{
- ArrowNumericType, Float32Type, Float64Type, Int32Type, Int64Type, Int8Type,
- UInt16Type,
- };
- use packed_simd::*;
- use FromCast;
-
- /// calculate the expected mask by iterating over all bits
- macro_rules! expected_mask {
- ($T:ty, $MASK:expr) => {{
- let mask = $MASK;
- // simd width of all types is currently 64 bytes -> 512 bits
- let lanes = 64 / std::mem::size_of::<$T>();
- // translate each set bit into a value of all ones (-1) of the correct type
- (0..lanes)
- .map(|i| (if (mask & (1 << i)) != 0 { -1 } else { 0 }))
- .collect::<Vec<$T>>()
- }};
- }
-
- #[test]
- fn test_mask_f64() {
- let mask = 0b10101010;
- let actual = Float64Type::mask_from_u64(mask);
- let expected = expected_mask!(i64, mask);
- let expected = m64x8::from_cast(i64x8::from_slice_unaligned(expected.as_slice()));
-
- assert_eq!(expected, actual);
- }
-
- #[test]
- fn test_mask_u64() {
- let mask = 0b01010101;
- let actual = Int64Type::mask_from_u64(mask);
- let expected = expected_mask!(i64, mask);
- let expected = m64x8::from_cast(i64x8::from_slice_unaligned(expected.as_slice()));
-
- assert_eq!(expected, actual);
- }
-
- #[test]
- fn test_mask_f32() {
- let mask = 0b10101010_10101010;
- let actual = Float32Type::mask_from_u64(mask);
- let expected = expected_mask!(i32, mask);
- let expected =
- m32x16::from_cast(i32x16::from_slice_unaligned(expected.as_slice()));
-
- assert_eq!(expected, actual);
- }
-
- #[test]
- fn test_mask_i32() {
- let mask = 0b01010101_01010101;
- let actual = Int32Type::mask_from_u64(mask);
- let expected = expected_mask!(i32, mask);
- let expected =
- m32x16::from_cast(i32x16::from_slice_unaligned(expected.as_slice()));
-
- assert_eq!(expected, actual);
- }
-
- #[test]
- fn test_mask_u16() {
- let mask = 0b01010101_01010101_10101010_10101010;
- let actual = UInt16Type::mask_from_u64(mask);
- let expected = expected_mask!(i16, mask);
- dbg!(&expected);
- let expected =
- m16x32::from_cast(i16x32::from_slice_unaligned(expected.as_slice()));
-
- assert_eq!(expected, actual);
- }
-
- #[test]
- fn test_mask_i8() {
- let mask =
- 0b01010101_01010101_10101010_10101010_01010101_01010101_10101010_10101010;
- let actual = Int8Type::mask_from_u64(mask);
- let expected = expected_mask!(i8, mask);
- let expected = m8x64::from_cast(i8x64::from_slice_unaligned(expected.as_slice()));
-
- assert_eq!(expected, actual);
- }
-}
diff --git a/arrow/src/datatypes/schema.rs b/arrow/src/datatypes/schema.rs
deleted file mode 100644
index cfc0744..0000000
--- a/arrow/src/datatypes/schema.rs
+++ /dev/null
@@ -1,370 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-use std::collections::HashMap;
-use std::default::Default;
-use std::fmt;
-
-use serde_derive::{Deserialize, Serialize};
-use serde_json::{json, Value};
-
-use crate::error::{ArrowError, Result};
-
-use super::Field;
-
-/// Describes the meta-data of an ordered sequence of relative types.
-///
-/// Note that this information is only part of the meta-data and not part of the physical
-/// memory layout.
-#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Eq)]
-pub struct Schema {
- pub(crate) fields: Vec<Field>,
- /// A map of key-value pairs containing additional meta data.
- #[serde(skip_serializing_if = "HashMap::is_empty")]
- #[serde(default)]
- pub(crate) metadata: HashMap<String, String>,
-}
-
-impl Schema {
- /// Creates an empty `Schema`
- pub fn empty() -> Self {
- Self {
- fields: vec![],
- metadata: HashMap::new(),
- }
- }
-
- /// Creates a new `Schema` from a sequence of `Field` values.
- ///
- /// # Example
- ///
- /// ```
- /// # extern crate arrow;
- /// # use arrow::datatypes::{Field, DataType, Schema};
- /// let field_a = Field::new("a", DataType::Int64, false);
- /// let field_b = Field::new("b", DataType::Boolean, false);
- ///
- /// let schema = Schema::new(vec![field_a, field_b]);
- /// ```
- pub fn new(fields: Vec<Field>) -> Self {
- Self::new_with_metadata(fields, HashMap::new())
- }
-
- /// Creates a new `Schema` from a sequence of `Field` values
- /// and adds additional metadata in form of key value pairs.
- ///
- /// # Example
- ///
- /// ```
- /// # extern crate arrow;
- /// # use arrow::datatypes::{Field, DataType, Schema};
- /// # use std::collections::HashMap;
- /// let field_a = Field::new("a", DataType::Int64, false);
- /// let field_b = Field::new("b", DataType::Boolean, false);
- ///
- /// let mut metadata: HashMap<String, String> = HashMap::new();
- /// metadata.insert("row_count".to_string(), "100".to_string());
- ///
- /// let schema = Schema::new_with_metadata(vec![field_a, field_b], metadata);
- /// ```
- #[inline]
- pub const fn new_with_metadata(
- fields: Vec<Field>,
- metadata: HashMap<String, String>,
- ) -> Self {
- Self { fields, metadata }
- }
-
- /// Merge schema into self if it is compatible. Struct fields will be merged recursively.
- ///
- /// Example:
- ///
- /// ```
- /// use arrow::datatypes::*;
- ///
- /// let merged = Schema::try_merge(vec![
- /// Schema::new(vec![
- /// Field::new("c1", DataType::Int64, false),
- /// Field::new("c2", DataType::Utf8, false),
- /// ]),
- /// Schema::new(vec![
- /// Field::new("c1", DataType::Int64, true),
- /// Field::new("c2", DataType::Utf8, false),
- /// Field::new("c3", DataType::Utf8, false),
- /// ]),
- /// ]).unwrap();
- ///
- /// assert_eq!(
- /// merged,
- /// Schema::new(vec![
- /// Field::new("c1", DataType::Int64, true),
- /// Field::new("c2", DataType::Utf8, false),
- /// Field::new("c3", DataType::Utf8, false),
- /// ]),
- /// );
- /// ```
- pub fn try_merge(schemas: impl IntoIterator<Item = Self>) -> Result<Self> {
- schemas
- .into_iter()
- .try_fold(Self::empty(), |mut merged, schema| {
- let Schema { metadata, fields } = schema;
- for (key, value) in metadata.into_iter() {
- // merge metadata
- if let Some(old_val) = merged.metadata.get(&key) {
- if old_val != &value {
- return Err(ArrowError::SchemaError(
- "Fail to merge schema due to conflicting metadata."
- .to_string(),
- ));
- }
- }
- merged.metadata.insert(key, value);
- }
- // merge fields
- for field in fields.into_iter() {
- let mut new_field = true;
- for merged_field in &mut merged.fields {
- if field.name() != merged_field.name() {
- continue;
- }
- new_field = false;
- merged_field.try_merge(&field)?
- }
- // found a new field, add to field list
- if new_field {
- merged.fields.push(field);
- }
- }
- Ok(merged)
- })
- }
-
- /// Returns an immutable reference of the vector of `Field` instances.
- #[inline]
- pub const fn fields(&self) -> &Vec<Field> {
- &self.fields
- }
-
- /// Returns an immutable reference of a specific `Field` instance selected using an
- /// offset within the internal `fields` vector.
- pub fn field(&self, i: usize) -> &Field {
- &self.fields[i]
- }
-
- /// Returns an immutable reference of a specific `Field` instance selected by name.
- pub fn field_with_name(&self, name: &str) -> Result<&Field> {
- Ok(&self.fields[self.index_of(name)?])
- }
-
- /// Returns a vector of immutable references to all `Field` instances selected by
- /// the dictionary ID they use.
- pub fn fields_with_dict_id(&self, dict_id: i64) -> Vec<&Field> {
- self.fields
- .iter()
- .filter(|f| f.dict_id() == Some(dict_id))
- .collect()
- }
-
- /// Find the index of the column with the given name.
- pub fn index_of(&self, name: &str) -> Result<usize> {
- for i in 0..self.fields.len() {
- if self.fields[i].name() == name {
- return Ok(i);
- }
- }
- let valid_fields: Vec<String> =
- self.fields.iter().map(|f| f.name().clone()).collect();
- Err(ArrowError::InvalidArgumentError(format!(
- "Unable to get field named \"{}\". Valid fields: {:?}",
- name, valid_fields
- )))
- }
-
- /// Returns an immutable reference to the Map of custom metadata key-value pairs.
- #[inline]
- pub const fn metadata(&self) -> &HashMap<String, String> {
- &self.metadata
- }
-
- /// Look up a column by name and return a immutable reference to the column along with
- /// its index.
- pub fn column_with_name(&self, name: &str) -> Option<(usize, &Field)> {
- self.fields
- .iter()
- .enumerate()
- .find(|&(_, c)| c.name() == name)
- }
-
- /// Generate a JSON representation of the `Schema`.
- pub fn to_json(&self) -> Value {
- json!({
- "fields": self.fields.iter().map(|field| field.to_json()).collect::<Vec<Value>>(),
- "metadata": serde_json::to_value(&self.metadata).unwrap()
- })
- }
-
- /// Parse a `Schema` definition from a JSON representation.
- pub fn from(json: &Value) -> Result<Self> {
- match *json {
- Value::Object(ref schema) => {
- let fields = if let Some(Value::Array(fields)) = schema.get("fields") {
- fields
- .iter()
- .map(|f| Field::from(f))
- .collect::<Result<_>>()?
- } else {
- return Err(ArrowError::ParseError(
- "Schema fields should be an array".to_string(),
- ));
- };
-
- let metadata = if let Some(value) = schema.get("metadata") {
- Self::from_metadata(value)?
- } else {
- HashMap::default()
- };
-
- Ok(Self { fields, metadata })
- }
- _ => Err(ArrowError::ParseError(
- "Invalid json value type for schema".to_string(),
- )),
- }
- }
-
- /// Parse a `metadata` definition from a JSON representation.
- /// The JSON can either be an Object or an Array of Objects.
- fn from_metadata(json: &Value) -> Result<HashMap<String, String>> {
- match json {
- Value::Array(_) => {
- let mut hashmap = HashMap::new();
- let values: Vec<MetadataKeyValue> = serde_json::from_value(json.clone())
- .map_err(|_| {
- ArrowError::JsonError(
- "Unable to parse object into key-value pair".to_string(),
- )
- })?;
- for meta in values {
- hashmap.insert(meta.key.clone(), meta.value);
- }
- Ok(hashmap)
- }
- Value::Object(md) => md
- .iter()
- .map(|(k, v)| {
- if let Value::String(v) = v {
- Ok((k.to_string(), v.to_string()))
- } else {
- Err(ArrowError::ParseError(
- "metadata `value` field must be a string".to_string(),
- ))
- }
- })
- .collect::<Result<_>>(),
- _ => Err(ArrowError::ParseError(
- "`metadata` field must be an object".to_string(),
- )),
- }
- }
-
- /// Check to see if `self` is a superset of `other` schema. Here are the comparision rules:
- ///
- /// * `self` and `other` should contain the same number of fields
- /// * for every field `f` in `other`, the field in `self` with corresponding index should be a
- /// superset of `f`.
- /// * self.metadata is a superset of other.metadata
- ///
- /// In other words, any record conforms to `other` should also conform to `self`.
- pub fn contains(&self, other: &Schema) -> bool {
- if self.fields.len() != other.fields.len() {
- return false;
- }
-
- for (i, field) in other.fields.iter().enumerate() {
- if !self.fields[i].contains(field) {
- return false;
- }
- }
-
- // make sure self.metadata is a superset of other.metadata
- for (k, v) in &other.metadata {
- match self.metadata.get(k) {
- Some(s) => {
- if s != v {
- return false;
- }
- }
- None => {
- return false;
- }
- }
- }
-
- true
- }
-}
-
-impl fmt::Display for Schema {
- fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
- f.write_str(
- &self
- .fields
- .iter()
- .map(|c| c.to_string())
- .collect::<Vec<String>>()
- .join(", "),
- )
- }
-}
-
-#[derive(Deserialize)]
-struct MetadataKeyValue {
- key: String,
- value: String,
-}
-
-#[cfg(test)]
-mod tests {
- use crate::datatypes::DataType;
-
- use super::*;
-
- #[test]
- fn test_ser_de_metadata() {
- // ser/de with empty metadata
- let mut schema = Schema::new(vec![
- Field::new("name", DataType::Utf8, false),
- Field::new("address", DataType::Utf8, false),
- Field::new("priority", DataType::UInt8, false),
- ]);
-
- let json = serde_json::to_string(&schema).unwrap();
- let de_schema = serde_json::from_str(&json).unwrap();
-
- assert_eq!(schema, de_schema);
-
- // ser/de with non-empty metadata
- schema.metadata = [("key".to_owned(), "val".to_owned())]
- .iter()
- .cloned()
- .collect();
- let json = serde_json::to_string(&schema).unwrap();
- let de_schema = serde_json::from_str(&json).unwrap();
-
- assert_eq!(schema, de_schema);
- }
-}
diff --git a/arrow/src/datatypes/types.rs b/arrow/src/datatypes/types.rs
deleted file mode 100644
index 30c9aae..0000000
--- a/arrow/src/datatypes/types.rs
+++ /dev/null
@@ -1,185 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-use super::{ArrowPrimitiveType, DataType, IntervalUnit, TimeUnit};
-
-// BooleanType is special: its bit-width is not the size of the primitive type, and its `index`
-// operation assumes bit-packing.
-#[derive(Debug)]
-pub struct BooleanType {}
-
-impl BooleanType {
- pub const DATA_TYPE: DataType = DataType::Boolean;
-}
-
-macro_rules! make_type {
- ($name:ident, $native_ty:ty, $data_ty:expr) => {
- #[derive(Debug)]
- pub struct $name {}
-
- impl ArrowPrimitiveType for $name {
- type Native = $native_ty;
- const DATA_TYPE: DataType = $data_ty;
- }
- };
-}
-
-make_type!(Int8Type, i8, DataType::Int8);
-make_type!(Int16Type, i16, DataType::Int16);
-make_type!(Int32Type, i32, DataType::Int32);
-make_type!(Int64Type, i64, DataType::Int64);
-make_type!(UInt8Type, u8, DataType::UInt8);
-make_type!(UInt16Type, u16, DataType::UInt16);
-make_type!(UInt32Type, u32, DataType::UInt32);
-make_type!(UInt64Type, u64, DataType::UInt64);
-make_type!(Float32Type, f32, DataType::Float32);
-make_type!(Float64Type, f64, DataType::Float64);
-make_type!(
- TimestampSecondType,
- i64,
- DataType::Timestamp(TimeUnit::Second, None)
-);
-make_type!(
- TimestampMillisecondType,
- i64,
- DataType::Timestamp(TimeUnit::Millisecond, None)
-);
-make_type!(
- TimestampMicrosecondType,
- i64,
- DataType::Timestamp(TimeUnit::Microsecond, None)
-);
-make_type!(
- TimestampNanosecondType,
- i64,
- DataType::Timestamp(TimeUnit::Nanosecond, None)
-);
-make_type!(Date32Type, i32, DataType::Date32);
-make_type!(Date64Type, i64, DataType::Date64);
-make_type!(Time32SecondType, i32, DataType::Time32(TimeUnit::Second));
-make_type!(
- Time32MillisecondType,
- i32,
- DataType::Time32(TimeUnit::Millisecond)
-);
-make_type!(
- Time64MicrosecondType,
- i64,
- DataType::Time64(TimeUnit::Microsecond)
-);
-make_type!(
- Time64NanosecondType,
- i64,
- DataType::Time64(TimeUnit::Nanosecond)
-);
-make_type!(
- IntervalYearMonthType,
- i32,
- DataType::Interval(IntervalUnit::YearMonth)
-);
-make_type!(
- IntervalDayTimeType,
- i64,
- DataType::Interval(IntervalUnit::DayTime)
-);
-make_type!(
- DurationSecondType,
- i64,
- DataType::Duration(TimeUnit::Second)
-);
-make_type!(
- DurationMillisecondType,
- i64,
- DataType::Duration(TimeUnit::Millisecond)
-);
-make_type!(
- DurationMicrosecondType,
- i64,
- DataType::Duration(TimeUnit::Microsecond)
-);
-make_type!(
- DurationNanosecondType,
- i64,
- DataType::Duration(TimeUnit::Nanosecond)
-);
-
-/// A subtype of primitive type that represents legal dictionary keys.
-/// See <https://arrow.apache.org/docs/format/Columnar.html>
-pub trait ArrowDictionaryKeyType: ArrowPrimitiveType {}
-
-impl ArrowDictionaryKeyType for Int8Type {}
-
-impl ArrowDictionaryKeyType for Int16Type {}
-
-impl ArrowDictionaryKeyType for Int32Type {}
-
-impl ArrowDictionaryKeyType for Int64Type {}
-
-impl ArrowDictionaryKeyType for UInt8Type {}
-
-impl ArrowDictionaryKeyType for UInt16Type {}
-
-impl ArrowDictionaryKeyType for UInt32Type {}
-
-impl ArrowDictionaryKeyType for UInt64Type {}
-
-/// A subtype of primitive type that represents temporal values.
-pub trait ArrowTemporalType: ArrowPrimitiveType {}
-
-impl ArrowTemporalType for TimestampSecondType {}
-impl ArrowTemporalType for TimestampMillisecondType {}
-impl ArrowTemporalType for TimestampMicrosecondType {}
-impl ArrowTemporalType for TimestampNanosecondType {}
-impl ArrowTemporalType for Date32Type {}
-impl ArrowTemporalType for Date64Type {}
-impl ArrowTemporalType for Time32SecondType {}
-impl ArrowTemporalType for Time32MillisecondType {}
-impl ArrowTemporalType for Time64MicrosecondType {}
-impl ArrowTemporalType for Time64NanosecondType {}
-// impl ArrowTemporalType for IntervalYearMonthType {}
-// impl ArrowTemporalType for IntervalDayTimeType {}
-impl ArrowTemporalType for DurationSecondType {}
-impl ArrowTemporalType for DurationMillisecondType {}
-impl ArrowTemporalType for DurationMicrosecondType {}
-impl ArrowTemporalType for DurationNanosecondType {}
-
-/// A timestamp type allows us to create array builders that take a timestamp.
-pub trait ArrowTimestampType: ArrowTemporalType {
- /// Returns the `TimeUnit` of this timestamp.
- fn get_time_unit() -> TimeUnit;
-}
-
-impl ArrowTimestampType for TimestampSecondType {
- fn get_time_unit() -> TimeUnit {
- TimeUnit::Second
- }
-}
-impl ArrowTimestampType for TimestampMillisecondType {
- fn get_time_unit() -> TimeUnit {
- TimeUnit::Millisecond
- }
-}
-impl ArrowTimestampType for TimestampMicrosecondType {
- fn get_time_unit() -> TimeUnit {
- TimeUnit::Microsecond
- }
-}
-impl ArrowTimestampType for TimestampNanosecondType {
- fn get_time_unit() -> TimeUnit {
- TimeUnit::Nanosecond
- }
-}
diff --git a/arrow/src/error.rs b/arrow/src/error.rs
deleted file mode 100644
index 86896c0..0000000
--- a/arrow/src/error.rs
+++ /dev/null
@@ -1,134 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-//! Defines `ArrowError` for representing failures in various Arrow operations.
-use std::fmt::{Debug, Display, Formatter};
-use std::io::Write;
-
-use std::error::Error;
-
-/// Many different operations in the `arrow` crate return this error type.
-#[derive(Debug)]
-pub enum ArrowError {
- /// Returned when functionality is not yet available.
- NotYetImplemented(String),
- ExternalError(Box<dyn Error + Send + Sync>),
- CastError(String),
- MemoryError(String),
- ParseError(String),
- SchemaError(String),
- ComputeError(String),
- DivideByZero,
- CsvError(String),
- JsonError(String),
- IoError(String),
- InvalidArgumentError(String),
- ParquetError(String),
- /// Error during import or export to/from the C Data Interface
- CDataInterface(String),
- DictionaryKeyOverflowError,
-}
-
-impl ArrowError {
- /// Wraps an external error in an `ArrowError`.
- pub fn from_external_error(
- error: Box<dyn ::std::error::Error + Send + Sync>,
- ) -> Self {
- Self::ExternalError(error)
- }
-}
-
-impl From<::std::io::Error> for ArrowError {
- fn from(error: std::io::Error) -> Self {
- ArrowError::IoError(error.to_string())
- }
-}
-
-#[cfg(feature = "csv")]
-impl From<csv_crate::Error> for ArrowError {
- fn from(error: csv_crate::Error) -> Self {
- match error.kind() {
- csv_crate::ErrorKind::Io(error) => ArrowError::CsvError(error.to_string()),
- csv_crate::ErrorKind::Utf8 { pos: _, err } => ArrowError::CsvError(format!(
- "Encountered UTF-8 error while reading CSV file: {}",
- err.to_string()
- )),
- csv_crate::ErrorKind::UnequalLengths {
- expected_len, len, ..
- } => ArrowError::CsvError(format!(
- "Encountered unequal lengths between records on CSV file. Expected {} \
- records, found {} records",
- len, expected_len
- )),
- _ => ArrowError::CsvError("Error reading CSV file".to_string()),
- }
- }
-}
-
-impl From<::std::string::FromUtf8Error> for ArrowError {
- fn from(error: std::string::FromUtf8Error) -> Self {
- ArrowError::ParseError(error.to_string())
- }
-}
-
-impl From<serde_json::Error> for ArrowError {
- fn from(error: serde_json::Error) -> Self {
- ArrowError::JsonError(error.to_string())
- }
-}
-
-impl<W: Write> From<::std::io::IntoInnerError<W>> for ArrowError {
- fn from(error: std::io::IntoInnerError<W>) -> Self {
- ArrowError::IoError(error.to_string())
- }
-}
-
-impl Display for ArrowError {
- fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
- match self {
- ArrowError::NotYetImplemented(source) => {
- write!(f, "Not yet implemented: {}", &source)
- }
- ArrowError::ExternalError(source) => write!(f, "External error: {}", &source),
- ArrowError::CastError(desc) => write!(f, "Cast error: {}", desc),
- ArrowError::MemoryError(desc) => write!(f, "Memory error: {}", desc),
- ArrowError::ParseError(desc) => write!(f, "Parser error: {}", desc),
- ArrowError::SchemaError(desc) => write!(f, "Schema error: {}", desc),
- ArrowError::ComputeError(desc) => write!(f, "Compute error: {}", desc),
- ArrowError::DivideByZero => write!(f, "Divide by zero error"),
- ArrowError::CsvError(desc) => write!(f, "Csv error: {}", desc),
- ArrowError::JsonError(desc) => write!(f, "Json error: {}", desc),
- ArrowError::IoError(desc) => write!(f, "Io error: {}", desc),
- ArrowError::InvalidArgumentError(desc) => {
- write!(f, "Invalid argument error: {}", desc)
- }
- ArrowError::ParquetError(desc) => {
- write!(f, "Parquet argument error: {}", desc)
- }
- ArrowError::CDataInterface(desc) => {
- write!(f, "C Data interface error: {}", desc)
- }
- ArrowError::DictionaryKeyOverflowError => {
- write!(f, "Dictionary key bigger than the key type")
- }
- }
- }
-}
-
-impl Error for ArrowError {}
-
-pub type Result<T> = std::result::Result<T, ArrowError>;
diff --git a/arrow/src/ffi.rs b/arrow/src/ffi.rs
deleted file mode 100644
index 42dc440..0000000
--- a/arrow/src/ffi.rs
+++ /dev/null
@@ -1,1082 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-//! Contains declarations to bind to the [C Data Interface](https://arrow.apache.org/docs/format/CDataInterface.html).
-//!
-//! Generally, this module is divided in two main interfaces:
-//! One interface maps C ABI to native Rust types, i.e. convert c-pointers, c_char, to native rust.
-//! This is handled by [FFI_ArrowSchema] and [FFI_ArrowArray].
-//!
-//! The second interface maps native Rust types to the Rust-specific implementation of Arrow such as `format` to `Datatype`,
-//! `Buffer`, etc. This is handled by `ArrowArray`.
-//!
-//! ```rust
-//! # use std::sync::Arc;
-//! # use arrow::array::{Int32Array, Array, ArrayData, make_array_from_raw};
-//! # use arrow::error::{Result, ArrowError};
-//! # use arrow::compute::kernels::arithmetic;
-//! # use std::convert::TryFrom;
-//! # fn main() -> Result<()> {
-//! // create an array natively
-//! let array = Int32Array::from(vec![Some(1), None, Some(3)]);
-//!
-//! // export it
-//! let (array_ptr, schema_ptr) = array.to_raw()?;
-//!
-//! // consumed and used by something else...
-//!
-//! // import it
-//! let array = unsafe { make_array_from_raw(array_ptr, schema_ptr)? };
-//!
-//! // perform some operation
-//! let array = array.as_any().downcast_ref::<Int32Array>().ok_or(
-//! ArrowError::ParseError("Expects an int32".to_string()),
-//! )?;
-//! let array = arithmetic::add(&array, &array)?;
-//!
-//! // verify
-//! assert_eq!(array, Int32Array::from(vec![Some(2), None, Some(6)]));
-//!
-//! // (drop/release)
-//! Ok(())
-//! }
-//! ```
-
-/*
-# Design:
-
-Main assumptions:
-* A memory region is deallocated according it its own release mechanism.
-* Rust shares memory regions between arrays.
-* A memory region should be deallocated when no-one is using it.
-
-The design of this module is as follows:
-
-`ArrowArray` contains two `Arc`s, one per ABI-compatible `struct`, each containing data
-according to the C Data Interface. These Arcs are used for ref counting of the structs
-within Rust and lifetime management.
-
-Each ABI-compatible `struct` knowns how to `drop` itself, calling `release`.
-
-To import an array, unsafely create an `ArrowArray` from two pointers using [ArrowArray::try_from_raw].
-To export an array, create an `ArrowArray` using [ArrowArray::try_new].
-*/
-
-use std::{
- ffi::CStr,
- ffi::CString,
- iter,
- mem::size_of,
- ptr::{self, NonNull},
- sync::Arc,
-};
-
-use crate::array::ArrayData;
-use crate::buffer::Buffer;
-use crate::datatypes::{DataType, Field, TimeUnit};
-use crate::error::{ArrowError, Result};
-use crate::util::bit_util;
-
-#[allow(dead_code)]
-struct SchemaPrivateData {
- field: Field,
- children_ptr: Box<[*mut FFI_ArrowSchema]>,
-}
-
-/// ABI-compatible struct for `ArrowSchema` from C Data Interface
-/// See <https://arrow.apache.org/docs/format/CDataInterface.html#structure-definitions>
-/// This was created by bindgen
-#[repr(C)]
-#[derive(Debug)]
-pub struct FFI_ArrowSchema {
- format: *const ::std::os::raw::c_char,
- name: *const ::std::os::raw::c_char,
- metadata: *const ::std::os::raw::c_char,
- flags: i64,
- n_children: i64,
- children: *mut *mut FFI_ArrowSchema,
- dictionary: *mut FFI_ArrowSchema,
- release: ::std::option::Option<unsafe extern "C" fn(arg1: *mut FFI_ArrowSchema)>,
- private_data: *mut ::std::os::raw::c_void,
-}
-
-// callback used to drop [FFI_ArrowSchema] when it is exported.
-unsafe extern "C" fn release_schema(schema: *mut FFI_ArrowSchema) {
- if schema.is_null() {
- return;
- }
- let schema = &mut *schema;
-
- // take ownership back to release it.
- CString::from_raw(schema.format as *mut std::os::raw::c_char);
- CString::from_raw(schema.name as *mut std::os::raw::c_char);
- let private = Box::from_raw(schema.private_data as *mut SchemaPrivateData);
- for child in private.children_ptr.iter() {
- let _ = Box::from_raw(*child);
- }
-
- schema.release = None;
-}
-
-impl FFI_ArrowSchema {
- /// create a new [`Ffi_ArrowSchema`]. This fails if the fields' [`DataType`] is not supported.
- fn try_new(field: Field) -> Result<FFI_ArrowSchema> {
- let format = to_format(field.data_type())?;
- let name = field.name().clone();
-
- // allocate (and hold) the children
- let children_vec = match field.data_type() {
- DataType::List(field) => {
- vec![Box::new(FFI_ArrowSchema::try_new(field.as_ref().clone())?)]
- }
- DataType::LargeList(field) => {
- vec![Box::new(FFI_ArrowSchema::try_new(field.as_ref().clone())?)]
- }
- DataType::Struct(fields) => fields
- .iter()
- .map(|field| Ok(Box::new(FFI_ArrowSchema::try_new(field.clone())?)))
- .collect::<Result<Vec<_>>>()?,
- _ => vec![],
- };
- // note: this cannot be done along with the above because the above is fallible and this op leaks.
- let children_ptr = children_vec
- .into_iter()
- .map(Box::into_raw)
- .collect::<Box<_>>();
- let n_children = children_ptr.len() as i64;
-
- let flags = field.is_nullable() as i64 * 2;
-
- let mut private = Box::new(SchemaPrivateData {
- field,
- children_ptr,
- });
-
- // <https://arrow.apache.org/docs/format/CDataInterface.html#c.ArrowSchema>
- Ok(FFI_ArrowSchema {
- format: CString::new(format).unwrap().into_raw(),
- name: CString::new(name).unwrap().into_raw(),
- metadata: std::ptr::null_mut(),
- flags,
- n_children,
- children: private.children_ptr.as_mut_ptr(),
- dictionary: std::ptr::null_mut(),
- release: Some(release_schema),
- private_data: Box::into_raw(private) as *mut ::std::os::raw::c_void,
- })
- }
-
- /// create an empty [FFI_ArrowSchema]
- fn empty() -> Self {
- Self {
- format: std::ptr::null_mut(),
- name: std::ptr::null_mut(),
- metadata: std::ptr::null_mut(),
- flags: 0,
- n_children: 0,
- children: ptr::null_mut(),
- dictionary: std::ptr::null_mut(),
- release: None,
- private_data: std::ptr::null_mut(),
- }
- }
-
- /// returns the format of this schema.
- pub fn format(&self) -> &str {
- assert!(!self.format.is_null());
- // safe because the lifetime of `self.format` equals `self`
- unsafe { CStr::from_ptr(self.format) }
- .to_str()
- .expect("The external API has a non-utf8 as format")
- }
-
- /// returns the name of this schema.
- pub fn name(&self) -> &str {
- assert!(!self.name.is_null());
- // safe because the lifetime of `self.name` equals `self`
- unsafe { CStr::from_ptr(self.name) }.to_str().unwrap()
- }
-
- pub fn child(&self, index: usize) -> &Self {
- assert!(index < self.n_children as usize);
- assert!(!self.name.is_null());
- unsafe { self.children.add(index).as_ref().unwrap().as_ref().unwrap() }
- }
-
- pub fn nullable(&self) -> bool {
- (self.flags / 2) & 1 == 1
- }
-}
-
-impl Drop for FFI_ArrowSchema {
- fn drop(&mut self) {
- match self.release {
- None => (),
- Some(release) => unsafe { release(self) },
- };
- }
-}
-
-/// See https://arrow.apache.org/docs/format/CDataInterface.html#data-type-description-format-strings
-fn to_field(schema: &FFI_ArrowSchema) -> Result<Field> {
- let data_type = match schema.format() {
- "n" => DataType::Null,
- "b" => DataType::Boolean,
- "c" => DataType::Int8,
- "C" => DataType::UInt8,
- "s" => DataType::Int16,
- "S" => DataType::UInt16,
- "i" => DataType::Int32,
- "I" => DataType::UInt32,
- "l" => DataType::Int64,
- "L" => DataType::UInt64,
- "e" => DataType::Float16,
- "f" => DataType::Float32,
- "g" => DataType::Float64,
- "z" => DataType::Binary,
- "Z" => DataType::LargeBinary,
- "u" => DataType::Utf8,
- "U" => DataType::LargeUtf8,
- "tdD" => DataType::Date32,
- "tdm" => DataType::Date64,
- "tts" => DataType::Time32(TimeUnit::Second),
- "ttm" => DataType::Time32(TimeUnit::Millisecond),
- "ttu" => DataType::Time64(TimeUnit::Microsecond),
- "ttn" => DataType::Time64(TimeUnit::Nanosecond),
- "+l" => {
- let child = schema.child(0);
- DataType::List(Box::new(to_field(child)?))
- }
- "+L" => {
- let child = schema.child(0);
- DataType::LargeList(Box::new(to_field(child)?))
- }
- "+s" => {
- let children = (0..schema.n_children as usize)
- .map(|x| to_field(schema.child(x)))
- .collect::<Result<Vec<_>>>()?;
- DataType::Struct(children)
- }
- other => {
- return Err(ArrowError::CDataInterface(format!(
- "The datatype \"{:?}\" is still not supported in Rust implementation",
- other
- )))
- }
- };
- Ok(Field::new(schema.name(), data_type, schema.nullable()))
-}
-
-/// See https://arrow.apache.org/docs/format/CDataInterface.html#data-type-description-format-strings
-fn to_format(data_type: &DataType) -> Result<String> {
- Ok(match data_type {
- DataType::Null => "n",
- DataType::Boolean => "b",
- DataType::Int8 => "c",
- DataType::UInt8 => "C",
- DataType::Int16 => "s",
- DataType::UInt16 => "S",
- DataType::Int32 => "i",
- DataType::UInt32 => "I",
- DataType::Int64 => "l",
- DataType::UInt64 => "L",
- DataType::Float16 => "e",
- DataType::Float32 => "f",
- DataType::Float64 => "g",
- DataType::Binary => "z",
- DataType::LargeBinary => "Z",
- DataType::Utf8 => "u",
- DataType::LargeUtf8 => "U",
- DataType::Date32 => "tdD",
- DataType::Date64 => "tdm",
- DataType::Time32(TimeUnit::Second) => "tts",
- DataType::Time32(TimeUnit::Millisecond) => "ttm",
- DataType::Time64(TimeUnit::Microsecond) => "ttu",
- DataType::Time64(TimeUnit::Nanosecond) => "ttn",
- DataType::List(_) => "+l",
- DataType::LargeList(_) => "+L",
- DataType::Struct(_) => "+s",
- z => {
- return Err(ArrowError::CDataInterface(format!(
- "The datatype \"{:?}\" is still not supported in Rust implementation",
- z
- )))
- }
- }
- .to_string())
-}
-
-// returns the number of bits that buffer `i` (in the C data interface) is expected to have.
-// This is set by the Arrow specification
-fn bit_width(data_type: &DataType, i: usize) -> Result<usize> {
- Ok(match (data_type, i) {
- // the null buffer is bit sized
- (_, 0) => 1,
- // primitive types first buffer's size is given by the native types
- (DataType::Boolean, 1) => 1,
- (DataType::UInt8, 1) => size_of::<u8>() * 8,
- (DataType::UInt16, 1) => size_of::<u16>() * 8,
- (DataType::UInt32, 1) => size_of::<u32>() * 8,
- (DataType::UInt64, 1) => size_of::<u64>() * 8,
- (DataType::Int8, 1) => size_of::<i8>() * 8,
- (DataType::Int16, 1) => size_of::<i16>() * 8,
- (DataType::Int32, 1) | (DataType::Date32, 1) | (DataType::Time32(_), 1) => size_of::<i32>() * 8,
- (DataType::Int64, 1) | (DataType::Date64, 1) | (DataType::Time64(_), 1) => size_of::<i64>() * 8,
- (DataType::Float32, 1) => size_of::<f32>() * 8,
- (DataType::Float64, 1) => size_of::<f64>() * 8,
- // primitive types have a single buffer
- (DataType::Boolean, _) |
- (DataType::UInt8, _) |
- (DataType::UInt16, _) |
- (DataType::UInt32, _) |
- (DataType::UInt64, _) |
- (DataType::Int8, _) |
- (DataType::Int16, _) |
- (DataType::Int32, _) | (DataType::Date32, _) | (DataType::Time32(_), _) |
- (DataType::Int64, _) | (DataType::Date64, _) | (DataType::Time64(_), _) |
- (DataType::Float32, _) |
- (DataType::Float64, _) => {
- return Err(ArrowError::CDataInterface(format!(
- "The datatype \"{:?}\" expects 2 buffers, but requested {}. Please verify that the C data interface is correctly implemented.",
- data_type, i
- )))
- }
- // Variable-sized binaries: have two buffers.
- // "small": first buffer is i32, second is in bytes
- (DataType::Utf8, 1) | (DataType::Binary, 1) | (DataType::List(_), 1) => size_of::<i32>() * 8,
- (DataType::Utf8, 2) | (DataType::Binary, 2) | (DataType::List(_), 2) => size_of::<u8>() * 8,
- (DataType::Utf8, _) | (DataType::Binary, _) | (DataType::List(_), _)=> {
- return Err(ArrowError::CDataInterface(format!(
- "The datatype \"{:?}\" expects 3 buffers, but requested {}. Please verify that the C data interface is correctly implemented.",
- data_type, i
- )))
- }
- // Variable-sized binaries: have two buffers.
- // LargeUtf8: first buffer is i64, second is in bytes
- (DataType::LargeUtf8, 1) | (DataType::LargeBinary, 1) | (DataType::LargeList(_), 1) => size_of::<i64>() * 8,
- (DataType::LargeUtf8, 2) | (DataType::LargeBinary, 2) | (DataType::LargeList(_), 2)=> size_of::<u8>() * 8,
- (DataType::LargeUtf8, _) | (DataType::LargeBinary, _) | (DataType::LargeList(_), _)=> {
- return Err(ArrowError::CDataInterface(format!(
- "The datatype \"{:?}\" expects 3 buffers, but requested {}. Please verify that the C data interface is correctly implemented.",
- data_type, i
- )))
- }
- _ => {
- return Err(ArrowError::CDataInterface(format!(
- "The datatype \"{:?}\" is still not supported in Rust implementation",
- data_type
- )))
- }
- })
-}
-
-/// ABI-compatible struct for ArrowArray from C Data Interface
-/// See <https://arrow.apache.org/docs/format/CDataInterface.html#structure-definitions>
-/// This was created by bindgen
-#[repr(C)]
-#[derive(Debug)]
-pub struct FFI_ArrowArray {
- pub(crate) length: i64,
- pub(crate) null_count: i64,
- pub(crate) offset: i64,
- pub(crate) n_buffers: i64,
- pub(crate) n_children: i64,
- pub(crate) buffers: *mut *const ::std::os::raw::c_void,
- children: *mut *mut FFI_ArrowArray,
- dictionary: *mut FFI_ArrowArray,
- release: ::std::option::Option<unsafe extern "C" fn(arg1: *mut FFI_ArrowArray)>,
- // When exported, this MUST contain everything that is owned by this array.
- // for example, any buffer pointed to in `buffers` must be here, as well as the `buffers` pointer
- // itself.
- // In other words, everything in [FFI_ArrowArray] must be owned by `private_data` and can assume
- // that they do not outlive `private_data`.
- private_data: *mut ::std::os::raw::c_void,
-}
-
-impl Drop for FFI_ArrowArray {
- fn drop(&mut self) {
- match self.release {
- None => (),
- Some(release) => unsafe { release(self) },
- };
- }
-}
-
-// callback used to drop [FFI_ArrowArray] when it is exported
-unsafe extern "C" fn release_array(array: *mut FFI_ArrowArray) {
- if array.is_null() {
- return;
- }
- let array = &mut *array;
-
- // take ownership of `private_data`, therefore dropping it`
- let private = Box::from_raw(array.private_data as *mut PrivateData);
- for child in private.children.iter() {
- let _ = Box::from_raw(*child);
- }
-
- array.release = None;
-}
-
-struct PrivateData {
- buffers: Vec<Option<Buffer>>,
- buffers_ptr: Box<[*const std::os::raw::c_void]>,
- children: Box<[*mut FFI_ArrowArray]>,
-}
-
-impl FFI_ArrowArray {
- /// creates a new `FFI_ArrowArray` from existing data.
- /// # Safety
- /// This method releases `buffers`. Consumers of this struct *must* call `release` before
- /// releasing this struct, or contents in `buffers` leak.
- fn new(data: &ArrayData) -> Self {
- // * insert the null buffer at the start
- // * make all others `Option<Buffer>`.
- let buffers = iter::once(data.null_buffer().cloned())
- .chain(data.buffers().iter().map(|b| Some(b.clone())))
- .collect::<Vec<_>>();
- let n_buffers = buffers.len() as i64;
-
- let buffers_ptr = buffers
- .iter()
- .map(|maybe_buffer| match maybe_buffer {
- // note that `raw_data` takes into account the buffer's offset
- Some(b) => b.as_ptr() as *const std::os::raw::c_void,
- None => std::ptr::null(),
- })
- .collect::<Box<[_]>>();
-
- let children = data
- .child_data()
- .iter()
- .map(|child| Box::into_raw(Box::new(FFI_ArrowArray::new(child))))
- .collect::<Box<_>>();
- let n_children = children.len() as i64;
-
- // create the private data owning everything.
- // any other data must be added here, e.g. via a struct, to track lifetime.
- let mut private_data = Box::new(PrivateData {
- buffers,
- buffers_ptr,
- children,
- });
-
- Self {
- length: data.len() as i64,
- null_count: data.null_count() as i64,
- offset: data.offset() as i64,
- n_buffers,
- n_children,
- buffers: private_data.buffers_ptr.as_mut_ptr(),
- children: private_data.children.as_mut_ptr(),
- dictionary: std::ptr::null_mut(),
- release: Some(release_array),
- private_data: Box::into_raw(private_data) as *mut ::std::os::raw::c_void,
- }
- }
-
- // create an empty `FFI_ArrowArray`, which can be used to import data into
- fn empty() -> Self {
- Self {
- length: 0,
- null_count: 0,
- offset: 0,
- n_buffers: 0,
- n_children: 0,
- buffers: std::ptr::null_mut(),
- children: std::ptr::null_mut(),
- dictionary: std::ptr::null_mut(),
- release: None,
- private_data: std::ptr::null_mut(),
- }
- }
-
- /// the length of the array
- pub fn len(&self) -> usize {
- self.length as usize
- }
-
- /// whether the array is empty
- pub fn is_empty(&self) -> bool {
- self.length == 0
- }
-
- /// the offset of the array
- pub fn offset(&self) -> usize {
- self.offset as usize
- }
-
- /// the null count of the array
- pub fn null_count(&self) -> usize {
- self.null_count as usize
- }
-}
-
-/// returns a new buffer corresponding to the index `i` of the FFI array. It may not exist (null pointer).
-/// `bits` is the number of bits that the native type of this buffer has.
-/// The size of the buffer will be `ceil(self.length * bits, 8)`.
-/// # Panic
-/// This function panics if `i` is larger or equal to `n_buffers`.
-/// # Safety
-/// This function assumes that `ceil(self.length * bits, 8)` is the size of the buffer
-unsafe fn create_buffer(
- owner: Arc<FFI_ArrowArray>,
- array: &FFI_ArrowArray,
- index: usize,
- len: usize,
-) -> Option<Buffer> {
- if array.buffers.is_null() {
- return None;
- }
- let buffers = array.buffers as *mut *const u8;
-
- assert!(index < array.n_buffers as usize);
- let ptr = *buffers.add(index);
-
- NonNull::new(ptr as *mut u8).map(|ptr| Buffer::from_unowned(ptr, len, owner))
-}
-
-fn create_child(
- owner: Arc<FFI_ArrowArray>,
- array: &FFI_ArrowArray,
- schema: &FFI_ArrowSchema,
- index: usize,
-) -> ArrowArrayChild<'static> {
- assert!(index < array.n_children as usize);
- assert!(!array.children.is_null());
- assert!(!array.children.is_null());
- unsafe {
- let arr_ptr = *array.children.add(index);
- let schema_ptr = *schema.children.add(index);
- assert!(!arr_ptr.is_null());
- assert!(!schema_ptr.is_null());
- let arr_ptr = &*arr_ptr;
- let schema_ptr = &*schema_ptr;
- ArrowArrayChild::from_raw(arr_ptr, schema_ptr, owner)
- }
-}
-
-pub trait ArrowArrayRef {
- fn to_data(&self) -> Result<ArrayData> {
- let data_type = self.data_type()?;
- let len = self.array().len();
- let offset = self.array().offset();
- let null_count = self.array().null_count();
- let buffers = self.buffers()?;
- let null_bit_buffer = self.null_bit_buffer();
-
- let child_data = (0..self.array().n_children as usize)
- .map(|i| {
- let child = self.child(i);
- child.to_data()
- })
- .map(|d| d.unwrap())
- .collect();
-
- Ok(ArrayData::new(
- data_type,
- len,
- Some(null_count),
- null_bit_buffer,
- offset,
- buffers,
- child_data,
- ))
- }
-
- /// returns all buffers, as organized by Rust (i.e. null buffer is skipped)
- fn buffers(&self) -> Result<Vec<Buffer>> {
- (0..self.array().n_buffers - 1)
- .map(|index| {
- // + 1: skip null buffer
- let index = (index + 1) as usize;
-
- let len = self.buffer_len(index)?;
-
- unsafe { create_buffer(self.owner().clone(), self.array(), index, len) }
- .ok_or_else(|| {
- ArrowError::CDataInterface(format!(
- "The external buffer at position {} is null.",
- index - 1
- ))
- })
- })
- .collect()
- }
-
- /// Returns the length, in bytes, of the buffer `i` (indexed according to the C data interface)
- // Rust implementation uses fixed-sized buffers, which require knowledge of their `len`.
- // for variable-sized buffers, such as the second buffer of a stringArray, we need
- // to fetch offset buffer's len to build the second buffer.
- fn buffer_len(&self, i: usize) -> Result<usize> {
- // Inner type is not important for buffer length.
- let data_type = &self.data_type()?;
-
- Ok(match (data_type, i) {
- (DataType::Utf8, 1)
- | (DataType::LargeUtf8, 1)
- | (DataType::Binary, 1)
- | (DataType::LargeBinary, 1)
- | (DataType::List(_), 1)
- | (DataType::LargeList(_), 1) => {
- // the len of the offset buffer (buffer 1) equals length + 1
- let bits = bit_width(data_type, i)?;
- debug_assert_eq!(bits % 8, 0);
- (self.array().length as usize + 1) * (bits / 8)
- }
- (DataType::Utf8, 2) | (DataType::Binary, 2) | (DataType::List(_), 2) => {
- // the len of the data buffer (buffer 2) equals the last value of the offset buffer (buffer 1)
- let len = self.buffer_len(1)?;
- // first buffer is the null buffer => add(1)
- // we assume that pointer is aligned for `i32`, as Utf8 uses `i32` offsets.
- #[allow(clippy::cast_ptr_alignment)]
- let offset_buffer = unsafe {
- *(self.array().buffers as *mut *const u8).add(1) as *const i32
- };
- // get last offset
- (unsafe { *offset_buffer.add(len / size_of::<i32>() - 1) }) as usize
- }
- (DataType::LargeUtf8, 2)
- | (DataType::LargeBinary, 2)
- | (DataType::LargeList(_), 2) => {
- // the len of the data buffer (buffer 2) equals the last value of the offset buffer (buffer 1)
- let len = self.buffer_len(1)?;
- // first buffer is the null buffer => add(1)
- // we assume that pointer is aligned for `i64`, as Large uses `i64` offsets.
- #[allow(clippy::cast_ptr_alignment)]
- let offset_buffer = unsafe {
- *(self.array().buffers as *mut *const u8).add(1) as *const i64
- };
- // get last offset
- (unsafe { *offset_buffer.add(len / size_of::<i64>() - 1) }) as usize
- }
- // buffer len of primitive types
- _ => {
- let bits = bit_width(data_type, i)?;
- bit_util::ceil(self.array().length as usize * bits, 8)
- }
- })
- }
-
- /// returns the null bit buffer.
- /// Rust implementation uses a buffer that is not part of the array of buffers.
- /// The C Data interface's null buffer is part of the array of buffers.
- fn null_bit_buffer(&self) -> Option<Buffer> {
- // similar to `self.buffer_len(0)`, but without `Result`.
- let buffer_len = bit_util::ceil(self.array().length as usize, 8);
-
- unsafe { create_buffer(self.owner().clone(), self.array(), 0, buffer_len) }
- }
-
- fn child(&self, index: usize) -> ArrowArrayChild {
- create_child(self.owner().clone(), self.array(), self.schema(), index)
- }
-
- fn owner(&self) -> &Arc<FFI_ArrowArray>;
- fn array(&self) -> &FFI_ArrowArray;
- fn schema(&self) -> &FFI_ArrowSchema;
- fn data_type(&self) -> Result<DataType>;
-}
-
-/// Struct used to move an Array from and to the C Data Interface.
-/// Its main responsibility is to expose functionality that requires
-/// both [FFI_ArrowArray] and [FFI_ArrowSchema].
-///
-/// This struct has two main paths:
-///
-/// ## Import from the C Data Interface
-/// * [ArrowArray::empty] to allocate memory to be filled by an external call
-/// * [ArrowArray::try_from_raw] to consume two non-null allocated pointers
-/// ## Export to the C Data Interface
-/// * [ArrowArray::try_new] to create a new [ArrowArray] from Rust-specific information
-/// * [ArrowArray::into_raw] to expose two pointers for [FFI_ArrowArray] and [FFI_ArrowSchema].
-///
-/// # Safety
-/// Whoever creates this struct is responsible for releasing their resources. Specifically,
-/// consumers *must* call [ArrowArray::into_raw] and take ownership of the individual pointers,
-/// calling [FFI_ArrowArray::release] and [FFI_ArrowSchema::release] accordingly.
-///
-/// Furthermore, this struct assumes that the incoming data agrees with the C data interface.
-#[derive(Debug)]
-pub struct ArrowArray {
- array: Arc<FFI_ArrowArray>,
- schema: Arc<FFI_ArrowSchema>,
-}
-
-#[derive(Debug)]
-pub struct ArrowArrayChild<'a> {
- array: &'a FFI_ArrowArray,
- schema: &'a FFI_ArrowSchema,
- owner: Arc<FFI_ArrowArray>,
-}
-
-impl ArrowArrayRef for ArrowArray {
- /// the data_type as declared in the schema
- fn data_type(&self) -> Result<DataType> {
- to_field(&self.schema).map(|x| x.data_type().clone())
- }
-
- fn array(&self) -> &FFI_ArrowArray {
- self.array.as_ref()
- }
-
- fn schema(&self) -> &FFI_ArrowSchema {
- self.schema.as_ref()
- }
-
- fn owner(&self) -> &Arc<FFI_ArrowArray> {
- &self.array
- }
-}
-
-impl<'a> ArrowArrayRef for ArrowArrayChild<'a> {
- /// the data_type as declared in the schema
- fn data_type(&self) -> Result<DataType> {
- to_field(self.schema).map(|x| x.data_type().clone())
- }
-
- fn array(&self) -> &FFI_ArrowArray {
- self.array
- }
-
- fn schema(&self) -> &FFI_ArrowSchema {
- self.schema
- }
-
- fn owner(&self) -> &Arc<FFI_ArrowArray> {
- &self.owner
- }
-}
-
-impl ArrowArray {
- /// creates a new `ArrowArray`. This is used to export to the C Data Interface.
- /// # Safety
- /// See safety of [ArrowArray]
- #[allow(clippy::too_many_arguments)]
- pub unsafe fn try_new(data: ArrayData) -> Result<Self> {
- let field = Field::new("", data.data_type().clone(), data.null_count() != 0);
- let array = Arc::new(FFI_ArrowArray::new(&data));
- let schema = Arc::new(FFI_ArrowSchema::try_new(field)?);
-
- Ok(ArrowArray { array, schema })
- }
-
- /// creates a new [ArrowArray] from two pointers. Used to import from the C Data Interface.
- /// # Safety
- /// See safety of [ArrowArray]
- /// # Error
- /// Errors if any of the pointers is null
- pub unsafe fn try_from_raw(
- array: *const FFI_ArrowArray,
- schema: *const FFI_ArrowSchema,
- ) -> Result<Self> {
- if array.is_null() || schema.is_null() {
- return Err(ArrowError::MemoryError(
- "At least one of the pointers passed to `try_from_raw` is null"
- .to_string(),
- ));
- };
- Ok(Self {
- array: Arc::from_raw(array as *mut FFI_ArrowArray),
- schema: Arc::from_raw(schema as *mut FFI_ArrowSchema),
- })
- }
-
- /// creates a new empty [ArrowArray]. Used to import from the C Data Interface.
- /// # Safety
- /// See safety of [ArrowArray]
- pub unsafe fn empty() -> Self {
- let schema = Arc::new(FFI_ArrowSchema::empty());
- let array = Arc::new(FFI_ArrowArray::empty());
- ArrowArray { array, schema }
- }
-
- /// exports [ArrowArray] to the C Data Interface
- pub fn into_raw(this: ArrowArray) -> (*const FFI_ArrowArray, *const FFI_ArrowSchema) {
- (Arc::into_raw(this.array), Arc::into_raw(this.schema))
- }
-}
-
-impl<'a> ArrowArrayChild<'a> {
- fn from_raw(
- array: &'a FFI_ArrowArray,
- schema: &'a FFI_ArrowSchema,
- owner: Arc<FFI_ArrowArray>,
- ) -> Self {
- Self {
- array,
- schema,
- owner,
- }
- }
-}
-
-#[cfg(test)]
-mod tests {
- use super::*;
- use crate::array::{
- make_array, Array, ArrayData, BinaryOffsetSizeTrait, BooleanArray,
- GenericBinaryArray, GenericListArray, GenericStringArray, Int32Array,
- OffsetSizeTrait, StringOffsetSizeTrait, Time32MillisecondArray,
- };
- use crate::compute::kernels;
- use crate::datatypes::Field;
- use std::convert::TryFrom;
- use std::iter::FromIterator;
-
- #[test]
- fn test_round_trip() -> Result<()> {
- // create an array natively
- let array = Int32Array::from(vec![1, 2, 3]);
-
- // export it
- let array = ArrowArray::try_from(array.data().clone())?;
-
- // (simulate consumer) import it
- let data = ArrayData::try_from(array)?;
- let array = make_array(data);
-
- // perform some operation
- let array = array.as_any().downcast_ref::<Int32Array>().unwrap();
- let array = kernels::arithmetic::add(&array, &array).unwrap();
-
- // verify
- assert_eq!(array, Int32Array::from(vec![2, 4, 6]));
-
- // (drop/release)
- Ok(())
- }
- // case with nulls is tested in the docs, through the example on this module.
-
- fn test_generic_string<Offset: StringOffsetSizeTrait>() -> Result<()> {
- // create an array natively
- let array =
- GenericStringArray::<Offset>::from(vec![Some("a"), None, Some("aaa")]);
-
- // export it
- let array = ArrowArray::try_from(array.data().clone())?;
-
- // (simulate consumer) import it
- let data = ArrayData::try_from(array)?;
- let array = make_array(data);
-
- // perform some operation
- let array = kernels::concat::concat(&[array.as_ref(), array.as_ref()]).unwrap();
- let array = array
- .as_any()
- .downcast_ref::<GenericStringArray<Offset>>()
- .unwrap();
-
- // verify
- let expected = GenericStringArray::<Offset>::from(vec![
- Some("a"),
- None,
- Some("aaa"),
- Some("a"),
- None,
- Some("aaa"),
- ]);
- assert_eq!(array, &expected);
-
- // (drop/release)
- Ok(())
- }
-
- #[test]
- fn test_string() -> Result<()> {
- test_generic_string::<i32>()
- }
-
- #[test]
- fn test_large_string() -> Result<()> {
- test_generic_string::<i64>()
- }
-
- fn test_generic_list<Offset: OffsetSizeTrait>() -> Result<()> {
- // Construct a value array
- let value_data = ArrayData::builder(DataType::Int32)
- .len(8)
- .add_buffer(Buffer::from_slice_ref(&[0, 1, 2, 3, 4, 5, 6, 7]))
- .build();
-
- // Construct a buffer for value offsets, for the nested array:
- // [[0, 1, 2], [3, 4, 5], [6, 7]]
- let value_offsets = Buffer::from_iter(
- [0usize, 3, 6, 8]
- .iter()
- .map(|i| Offset::from_usize(*i).unwrap()),
- );
-
- // Construct a list array from the above two
- let list_data_type = match std::mem::size_of::<Offset>() {
- 4 => DataType::List(Box::new(Field::new("item", DataType::Int32, false))),
- _ => {
- DataType::LargeList(Box::new(Field::new("item", DataType::Int32, false)))
- }
- };
-
- let list_data = ArrayData::builder(list_data_type)
- .len(3)
- .add_buffer(value_offsets)
- .add_child_data(value_data)
- .build();
-
- // create an array natively
- let array = GenericListArray::<Offset>::from(list_data.clone());
-
- // export it
- let array = ArrowArray::try_from(array.data().clone())?;
-
- // (simulate consumer) import it
- let data = ArrayData::try_from(array)?;
- let array = make_array(data);
-
- // downcast
- let array = array
- .as_any()
- .downcast_ref::<GenericListArray<Offset>>()
- .unwrap();
-
- dbg!(&array);
-
- // verify
- let expected = GenericListArray::<Offset>::from(list_data);
- assert_eq!(&array.value(0), &expected.value(0));
- assert_eq!(&array.value(1), &expected.value(1));
- assert_eq!(&array.value(2), &expected.value(2));
-
- // (drop/release)
- Ok(())
- }
-
- #[test]
- fn test_list() -> Result<()> {
- test_generic_list::<i32>()
- }
-
- #[test]
- fn test_large_list() -> Result<()> {
- test_generic_list::<i64>()
- }
-
- fn test_generic_binary<Offset: BinaryOffsetSizeTrait>() -> Result<()> {
- // create an array natively
- let array: Vec<Option<&[u8]>> = vec![Some(b"a"), None, Some(b"aaa")];
- let array = GenericBinaryArray::<Offset>::from(array);
-
- // export it
- let array = ArrowArray::try_from(array.data().clone())?;
-
- // (simulate consumer) import it
- let data = ArrayData::try_from(array)?;
- let array = make_array(data);
-
- // perform some operation
- let array = kernels::concat::concat(&[array.as_ref(), array.as_ref()]).unwrap();
- let array = array
- .as_any()
- .downcast_ref::<GenericBinaryArray<Offset>>()
- .unwrap();
-
- // verify
- let expected: Vec<Option<&[u8]>> = vec![
- Some(b"a"),
- None,
- Some(b"aaa"),
- Some(b"a"),
- None,
- Some(b"aaa"),
- ];
- let expected = GenericBinaryArray::<Offset>::from(expected);
- assert_eq!(array, &expected);
-
- // (drop/release)
- Ok(())
- }
-
- #[test]
- fn test_binary() -> Result<()> {
- test_generic_binary::<i32>()
- }
-
- #[test]
- fn test_large_binary() -> Result<()> {
- test_generic_binary::<i64>()
- }
-
- #[test]
- fn test_bool() -> Result<()> {
- // create an array natively
- let array = BooleanArray::from(vec![None, Some(true), Some(false)]);
-
- // export it
- let array = ArrowArray::try_from(array.data().clone())?;
-
- // (simulate consumer) import it
- let data = ArrayData::try_from(array)?;
- let array = make_array(data);
-
- // perform some operation
- let array = array.as_any().downcast_ref::<BooleanArray>().unwrap();
- let array = kernels::boolean::not(&array)?;
-
- // verify
- assert_eq!(
- array,
- BooleanArray::from(vec![None, Some(false), Some(true)])
- );
-
- // (drop/release)
- Ok(())
- }
-
- #[test]
- fn test_time32() -> Result<()> {
- // create an array natively
- let array = Time32MillisecondArray::from(vec![None, Some(1), Some(2)]);
-
- // export it
- let array = ArrowArray::try_from(array.data().clone())?;
-
- // (simulate consumer) import it
- let data = ArrayData::try_from(array)?;
- let array = make_array(data);
-
- // perform some operation
- let array = kernels::concat::concat(&[array.as_ref(), array.as_ref()]).unwrap();
- let array = array
- .as_any()
- .downcast_ref::<Time32MillisecondArray>()
- .unwrap();
-
- // verify
- assert_eq!(
- array,
- &Time32MillisecondArray::from(vec![
- None,
- Some(1),
- Some(2),
- None,
- Some(1),
- Some(2)
- ])
- );
-
- // (drop/release)
- Ok(())
- }
-}
diff --git a/arrow/src/ipc/convert.rs b/arrow/src/ipc/convert.rs
deleted file mode 100644
index 59d4d0b..0000000
--- a/arrow/src/ipc/convert.rs
+++ /dev/null
@@ -1,871 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-//! Utilities for converting between IPC types and native Arrow types
-
-use crate::datatypes::{DataType, Field, IntervalUnit, Schema, TimeUnit};
-use crate::error::{ArrowError, Result};
-use crate::ipc;
-
-use flatbuffers::{
- FlatBufferBuilder, ForwardsUOffset, UnionWIPOffset, Vector, WIPOffset,
-};
-use std::collections::{BTreeMap, HashMap};
-
-use DataType::*;
-
-/// Serialize a schema in IPC format
-pub fn schema_to_fb(schema: &Schema) -> FlatBufferBuilder {
- let mut fbb = FlatBufferBuilder::new();
-
- let root = schema_to_fb_offset(&mut fbb, schema);
-
- fbb.finish(root, None);
-
- fbb
-}
-
-pub fn schema_to_fb_offset<'a>(
- fbb: &mut FlatBufferBuilder<'a>,
- schema: &Schema,
-) -> WIPOffset<ipc::Schema<'a>> {
- let mut fields = vec![];
- for field in schema.fields() {
- let fb_field = build_field(fbb, field);
- fields.push(fb_field);
- }
-
- let mut custom_metadata = vec![];
- for (k, v) in schema.metadata() {
- let fb_key_name = fbb.create_string(k.as_str());
- let fb_val_name = fbb.create_string(v.as_str());
-
- let mut kv_builder = ipc::KeyValueBuilder::new(fbb);
- kv_builder.add_key(fb_key_name);
- kv_builder.add_value(fb_val_name);
- custom_metadata.push(kv_builder.finish());
- }
-
- let fb_field_list = fbb.create_vector(&fields);
- let fb_metadata_list = fbb.create_vector(&custom_metadata);
-
- let mut builder = ipc::SchemaBuilder::new(fbb);
- builder.add_fields(fb_field_list);
- builder.add_custom_metadata(fb_metadata_list);
- builder.finish()
-}
-
-/// Convert an IPC Field to Arrow Field
-impl<'a> From<ipc::Field<'a>> for Field {
- fn from(field: ipc::Field) -> Field {
- let mut arrow_field = if let Some(dictionary) = field.dictionary() {
- Field::new_dict(
- field.name().unwrap(),
- get_data_type(field, true),
- field.nullable(),
- dictionary.id(),
- dictionary.isOrdered(),
- )
- } else {
- Field::new(
- field.name().unwrap(),
- get_data_type(field, true),
- field.nullable(),
- )
- };
-
- let mut metadata = None;
- if let Some(list) = field.custom_metadata() {
- let mut metadata_map = BTreeMap::default();
- for kv in list {
- if let (Some(k), Some(v)) = (kv.key(), kv.value()) {
- metadata_map.insert(k.to_string(), v.to_string());
- }
- }
- metadata = Some(metadata_map);
- }
-
- arrow_field.set_metadata(metadata);
- arrow_field
- }
-}
-
-/// Deserialize a Schema table from IPC format to Schema data type
-pub fn fb_to_schema(fb: ipc::Schema) -> Schema {
- let mut fields: Vec<Field> = vec![];
- let c_fields = fb.fields().unwrap();
- let len = c_fields.len();
- for i in 0..len {
- let c_field: ipc::Field = c_fields.get(i);
- match c_field.type_type() {
- ipc::Type::Decimal if fb.endianness() == ipc::Endianness::Big => {
- unimplemented!("Big Endian is not supported for Decimal!")
- }
- _ => (),
- };
- fields.push(c_field.into());
- }
-
- let mut metadata: HashMap<String, String> = HashMap::default();
- if let Some(md_fields) = fb.custom_metadata() {
- let len = md_fields.len();
- for i in 0..len {
- let kv = md_fields.get(i);
- let k_str = kv.key();
- let v_str = kv.value();
- if let Some(k) = k_str {
- if let Some(v) = v_str {
- metadata.insert(k.to_string(), v.to_string());
- }
- }
- }
- }
- Schema::new_with_metadata(fields, metadata)
-}
-
-/// Deserialize an IPC message into a schema
-pub fn schema_from_bytes(bytes: &[u8]) -> Result<Schema> {
- if let Ok(ipc) = ipc::root_as_message(bytes) {
- if let Some(schema) = ipc.header_as_schema().map(fb_to_schema) {
- Ok(schema)
- } else {
- Err(ArrowError::IoError(
- "Unable to get head as schema".to_string(),
- ))
- }
- } else {
- Err(ArrowError::IoError(
- "Unable to get root as message".to_string(),
- ))
- }
-}
-
-/// Get the Arrow data type from the flatbuffer Field table
-pub(crate) fn get_data_type(field: ipc::Field, may_be_dictionary: bool) -> DataType {
- if let Some(dictionary) = field.dictionary() {
- if may_be_dictionary {
- let int = dictionary.indexType().unwrap();
- let index_type = match (int.bitWidth(), int.is_signed()) {
- (8, true) => DataType::Int8,
- (8, false) => DataType::UInt8,
- (16, true) => DataType::Int16,
- (16, false) => DataType::UInt16,
- (32, true) => DataType::Int32,
- (32, false) => DataType::UInt32,
- (64, true) => DataType::Int64,
- (64, false) => DataType::UInt64,
- _ => panic!("Unexpected bitwidth and signed"),
- };
- return DataType::Dictionary(
- Box::new(index_type),
- Box::new(get_data_type(field, false)),
- );
- }
- }
-
- match field.type_type() {
- ipc::Type::Null => DataType::Null,
- ipc::Type::Bool => DataType::Boolean,
- ipc::Type::Int => {
- let int = field.type_as_int().unwrap();
- match (int.bitWidth(), int.is_signed()) {
- (8, true) => DataType::Int8,
- (8, false) => DataType::UInt8,
- (16, true) => DataType::Int16,
- (16, false) => DataType::UInt16,
- (32, true) => DataType::Int32,
- (32, false) => DataType::UInt32,
- (64, true) => DataType::Int64,
- (64, false) => DataType::UInt64,
- z => panic!(
- "Int type with bit width of {} and signed of {} not supported",
- z.0, z.1
- ),
- }
- }
- ipc::Type::Binary => DataType::Binary,
- ipc::Type::LargeBinary => DataType::LargeBinary,
- ipc::Type::Utf8 => DataType::Utf8,
- ipc::Type::LargeUtf8 => DataType::LargeUtf8,
- ipc::Type::FixedSizeBinary => {
- let fsb = field.type_as_fixed_size_binary().unwrap();
- DataType::FixedSizeBinary(fsb.byteWidth())
- }
- ipc::Type::FloatingPoint => {
- let float = field.type_as_floating_point().unwrap();
- match float.precision() {
- ipc::Precision::HALF => DataType::Float16,
- ipc::Precision::SINGLE => DataType::Float32,
- ipc::Precision::DOUBLE => DataType::Float64,
- z => panic!("FloatingPoint type with precision of {:?} not supported", z),
- }
- }
- ipc::Type::Date => {
- let date = field.type_as_date().unwrap();
- match date.unit() {
- ipc::DateUnit::DAY => DataType::Date32,
- ipc::DateUnit::MILLISECOND => DataType::Date64,
- z => panic!("Date type with unit of {:?} not supported", z),
- }
- }
- ipc::Type::Time => {
- let time = field.type_as_time().unwrap();
- match (time.bitWidth(), time.unit()) {
- (32, ipc::TimeUnit::SECOND) => DataType::Time32(TimeUnit::Second),
- (32, ipc::TimeUnit::MILLISECOND) => {
- DataType::Time32(TimeUnit::Millisecond)
- }
- (64, ipc::TimeUnit::MICROSECOND) => {
- DataType::Time64(TimeUnit::Microsecond)
- }
- (64, ipc::TimeUnit::NANOSECOND) => DataType::Time64(TimeUnit::Nanosecond),
- z => panic!(
- "Time type with bit width of {} and unit of {:?} not supported",
- z.0, z.1
- ),
- }
- }
- ipc::Type::Timestamp => {
- let timestamp = field.type_as_timestamp().unwrap();
- let timezone: Option<String> = timestamp.timezone().map(|tz| tz.to_string());
- match timestamp.unit() {
- ipc::TimeUnit::SECOND => DataType::Timestamp(TimeUnit::Second, timezone),
- ipc::TimeUnit::MILLISECOND => {
- DataType::Timestamp(TimeUnit::Millisecond, timezone)
- }
- ipc::TimeUnit::MICROSECOND => {
- DataType::Timestamp(TimeUnit::Microsecond, timezone)
- }
- ipc::TimeUnit::NANOSECOND => {
- DataType::Timestamp(TimeUnit::Nanosecond, timezone)
- }
- z => panic!("Timestamp type with unit of {:?} not supported", z),
- }
- }
- ipc::Type::Interval => {
- let interval = field.type_as_interval().unwrap();
- match interval.unit() {
- ipc::IntervalUnit::YEAR_MONTH => {
- DataType::Interval(IntervalUnit::YearMonth)
- }
- ipc::IntervalUnit::DAY_TIME => DataType::Interval(IntervalUnit::DayTime),
- z => panic!("Interval type with unit of {:?} unsupported", z),
- }
- }
- ipc::Type::Duration => {
- let duration = field.type_as_duration().unwrap();
- match duration.unit() {
- ipc::TimeUnit::SECOND => DataType::Duration(TimeUnit::Second),
- ipc::TimeUnit::MILLISECOND => DataType::Duration(TimeUnit::Millisecond),
- ipc::TimeUnit::MICROSECOND => DataType::Duration(TimeUnit::Microsecond),
- ipc::TimeUnit::NANOSECOND => DataType::Duration(TimeUnit::Nanosecond),
- z => panic!("Duration type with unit of {:?} unsupported", z),
- }
- }
- ipc::Type::List => {
- let children = field.children().unwrap();
- if children.len() != 1 {
- panic!("expect a list to have one child")
- }
- DataType::List(Box::new(children.get(0).into()))
- }
- ipc::Type::LargeList => {
- let children = field.children().unwrap();
- if children.len() != 1 {
- panic!("expect a large list to have one child")
- }
- DataType::LargeList(Box::new(children.get(0).into()))
- }
- ipc::Type::FixedSizeList => {
- let children = field.children().unwrap();
- if children.len() != 1 {
- panic!("expect a list to have one child")
- }
- let fsl = field.type_as_fixed_size_list().unwrap();
- DataType::FixedSizeList(Box::new(children.get(0).into()), fsl.listSize())
- }
- ipc::Type::Struct_ => {
- let mut fields = vec![];
- if let Some(children) = field.children() {
- for i in 0..children.len() {
- fields.push(children.get(i).into());
- }
- };
-
- DataType::Struct(fields)
- }
- ipc::Type::Decimal => {
- let fsb = field.type_as_decimal().unwrap();
- DataType::Decimal(fsb.precision() as usize, fsb.scale() as usize)
- }
- t => unimplemented!("Type {:?} not supported", t),
- }
-}
-
-pub(crate) struct FBFieldType<'b> {
- pub(crate) type_type: ipc::Type,
- pub(crate) type_: WIPOffset<UnionWIPOffset>,
- pub(crate) children: Option<WIPOffset<Vector<'b, ForwardsUOffset<ipc::Field<'b>>>>>,
-}
-
-/// Create an IPC Field from an Arrow Field
-pub(crate) fn build_field<'a>(
- fbb: &mut FlatBufferBuilder<'a>,
- field: &Field,
-) -> WIPOffset<ipc::Field<'a>> {
- // Optional custom metadata.
- let mut fb_metadata = None;
- if let Some(metadata) = field.metadata() {
- if !metadata.is_empty() {
- let mut kv_vec = vec![];
- for (k, v) in metadata {
- let kv_args = ipc::KeyValueArgs {
- key: Some(fbb.create_string(k.as_str())),
- value: Some(fbb.create_string(v.as_str())),
- };
- let kv_offset = ipc::KeyValue::create(fbb, &kv_args);
- kv_vec.push(kv_offset);
- }
- fb_metadata = Some(fbb.create_vector(&kv_vec));
- }
- };
-
- let fb_field_name = fbb.create_string(field.name().as_str());
- let field_type = get_fb_field_type(field.data_type(), field.is_nullable(), fbb);
-
- let fb_dictionary = if let Dictionary(index_type, _) = field.data_type() {
- Some(get_fb_dictionary(
- index_type,
- field
- .dict_id()
- .expect("All Dictionary types have `dict_id`"),
- field
- .dict_is_ordered()
- .expect("All Dictionary types have `dict_is_ordered`"),
- fbb,
- ))
- } else {
- None
- };
-
- let mut field_builder = ipc::FieldBuilder::new(fbb);
- field_builder.add_name(fb_field_name);
- if let Some(dictionary) = fb_dictionary {
- field_builder.add_dictionary(dictionary)
- }
- field_builder.add_type_type(field_type.type_type);
- field_builder.add_nullable(field.is_nullable());
- match field_type.children {
- None => {}
- Some(children) => field_builder.add_children(children),
- };
- field_builder.add_type_(field_type.type_);
-
- if let Some(fb_metadata) = fb_metadata {
- field_builder.add_custom_metadata(fb_metadata);
- }
-
- field_builder.finish()
-}
-
-/// Get the IPC type of a data type
-pub(crate) fn get_fb_field_type<'a>(
- data_type: &DataType,
- is_nullable: bool,
- fbb: &mut FlatBufferBuilder<'a>,
-) -> FBFieldType<'a> {
- // some IPC implementations expect an empty list for child data, instead of a null value.
- // An empty field list is thus returned for primitive types
- let empty_fields: Vec<WIPOffset<ipc::Field>> = vec![];
- match data_type {
- Null => FBFieldType {
- type_type: ipc::Type::Null,
- type_: ipc::NullBuilder::new(fbb).finish().as_union_value(),
- children: Some(fbb.create_vector(&empty_fields[..])),
- },
- Boolean => FBFieldType {
- type_type: ipc::Type::Bool,
- type_: ipc::BoolBuilder::new(fbb).finish().as_union_value(),
- children: Some(fbb.create_vector(&empty_fields[..])),
- },
- UInt8 | UInt16 | UInt32 | UInt64 => {
- let children = fbb.create_vector(&empty_fields[..]);
- let mut builder = ipc::IntBuilder::new(fbb);
- builder.add_is_signed(false);
- match data_type {
- UInt8 => builder.add_bitWidth(8),
- UInt16 => builder.add_bitWidth(16),
- UInt32 => builder.add_bitWidth(32),
- UInt64 => builder.add_bitWidth(64),
- _ => {}
- };
- FBFieldType {
- type_type: ipc::Type::Int,
- type_: builder.finish().as_union_value(),
- children: Some(children),
- }
- }
- Int8 | Int16 | Int32 | Int64 => {
- let children = fbb.create_vector(&empty_fields[..]);
- let mut builder = ipc::IntBuilder::new(fbb);
- builder.add_is_signed(true);
- match data_type {
- Int8 => builder.add_bitWidth(8),
- Int16 => builder.add_bitWidth(16),
- Int32 => builder.add_bitWidth(32),
- Int64 => builder.add_bitWidth(64),
- _ => {}
- };
- FBFieldType {
- type_type: ipc::Type::Int,
- type_: builder.finish().as_union_value(),
- children: Some(children),
- }
- }
- Float16 | Float32 | Float64 => {
- let children = fbb.create_vector(&empty_fields[..]);
- let mut builder = ipc::FloatingPointBuilder::new(fbb);
- match data_type {
- Float16 => builder.add_precision(ipc::Precision::HALF),
- Float32 => builder.add_precision(ipc::Precision::SINGLE),
- Float64 => builder.add_precision(ipc::Precision::DOUBLE),
- _ => {}
- };
- FBFieldType {
- type_type: ipc::Type::FloatingPoint,
- type_: builder.finish().as_union_value(),
- children: Some(children),
- }
- }
- Binary => FBFieldType {
- type_type: ipc::Type::Binary,
- type_: ipc::BinaryBuilder::new(fbb).finish().as_union_value(),
- children: Some(fbb.create_vector(&empty_fields[..])),
- },
- LargeBinary => FBFieldType {
- type_type: ipc::Type::LargeBinary,
- type_: ipc::LargeBinaryBuilder::new(fbb).finish().as_union_value(),
- children: Some(fbb.create_vector(&empty_fields[..])),
- },
- Utf8 => FBFieldType {
- type_type: ipc::Type::Utf8,
- type_: ipc::Utf8Builder::new(fbb).finish().as_union_value(),
- children: Some(fbb.create_vector(&empty_fields[..])),
- },
- LargeUtf8 => FBFieldType {
- type_type: ipc::Type::LargeUtf8,
- type_: ipc::LargeUtf8Builder::new(fbb).finish().as_union_value(),
- children: Some(fbb.create_vector(&empty_fields[..])),
- },
- FixedSizeBinary(len) => {
- let mut builder = ipc::FixedSizeBinaryBuilder::new(fbb);
- builder.add_byteWidth(*len as i32);
- FBFieldType {
- type_type: ipc::Type::FixedSizeBinary,
- type_: builder.finish().as_union_value(),
- children: Some(fbb.create_vector(&empty_fields[..])),
- }
- }
- Date32 => {
- let mut builder = ipc::DateBuilder::new(fbb);
- builder.add_unit(ipc::DateUnit::DAY);
- FBFieldType {
- type_type: ipc::Type::Date,
- type_: builder.finish().as_union_value(),
- children: Some(fbb.create_vector(&empty_fields[..])),
- }
- }
- Date64 => {
- let mut builder = ipc::DateBuilder::new(fbb);
- builder.add_unit(ipc::DateUnit::MILLISECOND);
- FBFieldType {
- type_type: ipc::Type::Date,
- type_: builder.finish().as_union_value(),
- children: Some(fbb.create_vector(&empty_fields[..])),
- }
- }
- Time32(unit) | Time64(unit) => {
- let mut builder = ipc::TimeBuilder::new(fbb);
- match unit {
- TimeUnit::Second => {
- builder.add_bitWidth(32);
- builder.add_unit(ipc::TimeUnit::SECOND);
- }
- TimeUnit::Millisecond => {
- builder.add_bitWidth(32);
- builder.add_unit(ipc::TimeUnit::MILLISECOND);
- }
- TimeUnit::Microsecond => {
- builder.add_bitWidth(64);
- builder.add_unit(ipc::TimeUnit::MICROSECOND);
- }
- TimeUnit::Nanosecond => {
- builder.add_bitWidth(64);
- builder.add_unit(ipc::TimeUnit::NANOSECOND);
- }
- }
- FBFieldType {
- type_type: ipc::Type::Time,
- type_: builder.finish().as_union_value(),
- children: Some(fbb.create_vector(&empty_fields[..])),
- }
- }
- Timestamp(unit, tz) => {
- let tz = tz.clone().unwrap_or_else(String::new);
- let tz_str = fbb.create_string(tz.as_str());
- let mut builder = ipc::TimestampBuilder::new(fbb);
- let time_unit = match unit {
- TimeUnit::Second => ipc::TimeUnit::SECOND,
- TimeUnit::Millisecond => ipc::TimeUnit::MILLISECOND,
- TimeUnit::Microsecond => ipc::TimeUnit::MICROSECOND,
- TimeUnit::Nanosecond => ipc::TimeUnit::NANOSECOND,
- };
- builder.add_unit(time_unit);
- if !tz.is_empty() {
- builder.add_timezone(tz_str);
- }
- FBFieldType {
- type_type: ipc::Type::Timestamp,
- type_: builder.finish().as_union_value(),
- children: Some(fbb.create_vector(&empty_fields[..])),
- }
- }
- Interval(unit) => {
- let mut builder = ipc::IntervalBuilder::new(fbb);
- let interval_unit = match unit {
- IntervalUnit::YearMonth => ipc::IntervalUnit::YEAR_MONTH,
- IntervalUnit::DayTime => ipc::IntervalUnit::DAY_TIME,
- };
- builder.add_unit(interval_unit);
- FBFieldType {
- type_type: ipc::Type::Interval,
- type_: builder.finish().as_union_value(),
- children: Some(fbb.create_vector(&empty_fields[..])),
- }
- }
- Duration(unit) => {
- let mut builder = ipc::DurationBuilder::new(fbb);
- let time_unit = match unit {
- TimeUnit::Second => ipc::TimeUnit::SECOND,
- TimeUnit::Millisecond => ipc::TimeUnit::MILLISECOND,
- TimeUnit::Microsecond => ipc::TimeUnit::MICROSECOND,
- TimeUnit::Nanosecond => ipc::TimeUnit::NANOSECOND,
- };
- builder.add_unit(time_unit);
- FBFieldType {
- type_type: ipc::Type::Duration,
- type_: builder.finish().as_union_value(),
- children: Some(fbb.create_vector(&empty_fields[..])),
- }
- }
- List(ref list_type) => {
- let child = build_field(fbb, list_type);
- FBFieldType {
- type_type: ipc::Type::List,
- type_: ipc::ListBuilder::new(fbb).finish().as_union_value(),
- children: Some(fbb.create_vector(&[child])),
- }
- }
- LargeList(ref list_type) => {
- let child = build_field(fbb, list_type);
- FBFieldType {
- type_type: ipc::Type::LargeList,
- type_: ipc::LargeListBuilder::new(fbb).finish().as_union_value(),
- children: Some(fbb.create_vector(&[child])),
- }
- }
- FixedSizeList(ref list_type, len) => {
- let child = build_field(fbb, list_type);
- let mut builder = ipc::FixedSizeListBuilder::new(fbb);
- builder.add_listSize(*len as i32);
- FBFieldType {
- type_type: ipc::Type::FixedSizeList,
- type_: builder.finish().as_union_value(),
- children: Some(fbb.create_vector(&[child])),
- }
- }
- Struct(fields) => {
- // struct's fields are children
- let mut children = vec![];
- for field in fields {
- let inner_types =
- get_fb_field_type(field.data_type(), field.is_nullable(), fbb);
- let field_name = fbb.create_string(field.name());
- children.push(ipc::Field::create(
- fbb,
- &ipc::FieldArgs {
- name: Some(field_name),
- nullable: field.is_nullable(),
- type_type: inner_types.type_type,
- type_: Some(inner_types.type_),
- dictionary: None,
- children: inner_types.children,
- custom_metadata: None,
- },
- ));
- }
- FBFieldType {
- type_type: ipc::Type::Struct_,
- type_: ipc::Struct_Builder::new(fbb).finish().as_union_value(),
- children: Some(fbb.create_vector(&children[..])),
- }
- }
- Dictionary(_, value_type) => {
- // In this library, the dictionary "type" is a logical construct. Here we
- // pass through to the value type, as we've already captured the index
- // type in the DictionaryEncoding metadata in the parent field
- get_fb_field_type(value_type, is_nullable, fbb)
- }
- Decimal(precision, scale) => {
- let mut builder = ipc::DecimalBuilder::new(fbb);
- builder.add_precision(*precision as i32);
- builder.add_scale(*scale as i32);
- builder.add_bitWidth(128);
- FBFieldType {
- type_type: ipc::Type::Decimal,
- type_: builder.finish().as_union_value(),
- children: Some(fbb.create_vector(&empty_fields[..])),
- }
- }
- t => unimplemented!("Type {:?} not supported", t),
- }
-}
-
-/// Create an IPC dictionary encoding
-pub(crate) fn get_fb_dictionary<'a>(
- index_type: &DataType,
- dict_id: i64,
- dict_is_ordered: bool,
- fbb: &mut FlatBufferBuilder<'a>,
-) -> WIPOffset<ipc::DictionaryEncoding<'a>> {
- // We assume that the dictionary index type (as an integer) has already been
- // validated elsewhere, and can safely assume we are dealing with integers
- let mut index_builder = ipc::IntBuilder::new(fbb);
-
- match *index_type {
- Int8 | Int16 | Int32 | Int64 => index_builder.add_is_signed(true),
- UInt8 | UInt16 | UInt32 | UInt64 => index_builder.add_is_signed(false),
- _ => {}
- }
-
- match *index_type {
- Int8 | UInt8 => index_builder.add_bitWidth(8),
- Int16 | UInt16 => index_builder.add_bitWidth(16),
- Int32 | UInt32 => index_builder.add_bitWidth(32),
- Int64 | UInt64 => index_builder.add_bitWidth(64),
- _ => {}
- }
-
- let index_builder = index_builder.finish();
-
- let mut builder = ipc::DictionaryEncodingBuilder::new(fbb);
- builder.add_id(dict_id);
- builder.add_indexType(index_builder);
- builder.add_isOrdered(dict_is_ordered);
-
- builder.finish()
-}
-
-#[cfg(test)]
-mod tests {
- use super::*;
- use crate::datatypes::{DataType, Field, Schema};
-
- #[test]
- fn convert_schema_round_trip() {
- let md: HashMap<String, String> = [("Key".to_string(), "value".to_string())]
- .iter()
- .cloned()
- .collect();
- let field_md: BTreeMap<String, String> = [("k".to_string(), "v".to_string())]
- .iter()
- .cloned()
- .collect();
- let schema = Schema::new_with_metadata(
- vec![
- {
- let mut f = Field::new("uint8", DataType::UInt8, false);
- f.set_metadata(Some(field_md));
- f
- },
- Field::new("uint16", DataType::UInt16, true),
- Field::new("uint32", DataType::UInt32, false),
- Field::new("uint64", DataType::UInt64, true),
- Field::new("int8", DataType::Int8, true),
- Field::new("int16", DataType::Int16, false),
- Field::new("int32", DataType::Int32, true),
- Field::new("int64", DataType::Int64, false),
- Field::new("float16", DataType::Float16, true),
- Field::new("float32", DataType::Float32, false),
- Field::new("float64", DataType::Float64, true),
- Field::new("null", DataType::Null, false),
- Field::new("bool", DataType::Boolean, false),
- Field::new("date32", DataType::Date32, false),
- Field::new("date64", DataType::Date64, true),
- Field::new("time32[s]", DataType::Time32(TimeUnit::Second), true),
- Field::new("time32[ms]", DataType::Time32(TimeUnit::Millisecond), false),
- Field::new("time64[us]", DataType::Time64(TimeUnit::Microsecond), false),
- Field::new("time64[ns]", DataType::Time64(TimeUnit::Nanosecond), true),
- Field::new(
- "timestamp[s]",
- DataType::Timestamp(TimeUnit::Second, None),
- false,
- ),
- Field::new(
- "timestamp[ms]",
- DataType::Timestamp(TimeUnit::Millisecond, None),
- true,
- ),
- Field::new(
- "timestamp[us]",
- DataType::Timestamp(
- TimeUnit::Microsecond,
- Some("Africa/Johannesburg".to_string()),
- ),
- false,
- ),
- Field::new(
- "timestamp[ns]",
- DataType::Timestamp(TimeUnit::Nanosecond, None),
- true,
- ),
- Field::new(
- "interval[ym]",
- DataType::Interval(IntervalUnit::YearMonth),
- true,
- ),
- Field::new(
- "interval[dt]",
- DataType::Interval(IntervalUnit::DayTime),
- true,
- ),
- Field::new("utf8", DataType::Utf8, false),
- Field::new("binary", DataType::Binary, false),
- Field::new(
- "list[u8]",
- DataType::List(Box::new(Field::new("item", DataType::UInt8, false))),
- true,
- ),
- Field::new(
- "list[struct<float32, int32, bool>]",
- DataType::List(Box::new(Field::new(
- "struct",
- DataType::Struct(vec![
- Field::new("float32", DataType::UInt8, false),
- Field::new("int32", DataType::Int32, true),
- Field::new("bool", DataType::Boolean, true),
- ]),
- true,
- ))),
- false,
- ),
- Field::new(
- "struct<int64, list[struct<date32, list[struct<>]>]>",
- DataType::Struct(vec![
- Field::new("int64", DataType::Int64, true),
- Field::new(
- "list[struct<date32, list[struct<>]>]",
- DataType::List(Box::new(Field::new(
- "struct",
- DataType::Struct(vec![
- Field::new("date32", DataType::Date32, true),
- Field::new(
- "list[struct<>]",
- DataType::List(Box::new(Field::new(
- "struct",
- DataType::Struct(vec![]),
- false,
- ))),
- false,
- ),
- ]),
- false,
- ))),
- false,
- ),
- ]),
- false,
- ),
- Field::new("struct<>", DataType::Struct(vec![]), true),
- Field::new_dict(
- "dictionary<int32, utf8>",
- DataType::Dictionary(
- Box::new(DataType::Int32),
- Box::new(DataType::Utf8),
- ),
- true,
- 123,
- true,
- ),
- Field::new_dict(
- "dictionary<uint8, uint32>",
- DataType::Dictionary(
- Box::new(DataType::UInt8),
- Box::new(DataType::UInt32),
- ),
- true,
- 123,
- true,
- ),
- Field::new("decimal<usize, usize>", DataType::Decimal(10, 6), false),
- ],
- md,
- );
-
- let fb = schema_to_fb(&schema);
-
- // read back fields
- let ipc = ipc::root_as_schema(fb.finished_data()).unwrap();
- let schema2 = fb_to_schema(ipc);
- assert_eq!(schema, schema2);
- }
-
- #[test]
- fn schema_from_bytes() {
- // bytes of a schema generated from python (0.14.0), saved as an `ipc::Message`.
- // the schema is: Field("field1", DataType::UInt32, false)
- let bytes: Vec<u8> = vec![
- 16, 0, 0, 0, 0, 0, 10, 0, 12, 0, 6, 0, 5, 0, 8, 0, 10, 0, 0, 0, 0, 1, 3, 0,
- 12, 0, 0, 0, 8, 0, 8, 0, 0, 0, 4, 0, 8, 0, 0, 0, 4, 0, 0, 0, 1, 0, 0, 0, 20,
- 0, 0, 0, 16, 0, 20, 0, 8, 0, 0, 0, 7, 0, 12, 0, 0, 0, 16, 0, 16, 0, 0, 0, 0,
- 0, 0, 2, 32, 0, 0, 0, 20, 0, 0, 0, 4, 0, 0, 0, 0, 0, 0, 0, 0, 0, 6, 0, 8, 0,
- 4, 0, 6, 0, 0, 0, 32, 0, 0, 0, 6, 0, 0, 0, 102, 105, 101, 108, 100, 49, 0, 0,
- 0, 0, 0, 0,
- ];
- let ipc = ipc::root_as_message(&bytes[..]).unwrap();
- let schema = ipc.header_as_schema().unwrap();
-
- // a message generated from Rust, same as the Python one
- let bytes: Vec<u8> = vec![
- 16, 0, 0, 0, 0, 0, 10, 0, 14, 0, 12, 0, 11, 0, 4, 0, 10, 0, 0, 0, 20, 0, 0,
- 0, 0, 0, 0, 1, 3, 0, 10, 0, 12, 0, 0, 0, 8, 0, 4, 0, 10, 0, 0, 0, 8, 0, 0, 0,
- 8, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 16, 0, 0, 0, 12, 0, 18, 0, 12, 0, 0, 0,
- 11, 0, 4, 0, 12, 0, 0, 0, 20, 0, 0, 0, 0, 0, 0, 2, 20, 0, 0, 0, 0, 0, 6, 0,
- 8, 0, 4, 0, 6, 0, 0, 0, 32, 0, 0, 0, 6, 0, 0, 0, 102, 105, 101, 108, 100, 49,
- 0, 0,
- ];
- let ipc2 = ipc::root_as_message(&bytes[..]).unwrap();
- let schema2 = ipc.header_as_schema().unwrap();
-
- assert_eq!(schema, schema2);
- assert_eq!(ipc.version(), ipc2.version());
- assert_eq!(ipc.header_type(), ipc2.header_type());
- assert_eq!(ipc.bodyLength(), ipc2.bodyLength());
- assert!(ipc.custom_metadata().is_none());
- assert!(ipc2.custom_metadata().is_none());
- }
-}
diff --git a/arrow/src/ipc/gen/File.rs b/arrow/src/ipc/gen/File.rs
deleted file mode 100644
index 04cbc64..0000000
--- a/arrow/src/ipc/gen/File.rs
+++ /dev/null
@@ -1,491 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-#![allow(dead_code)]
-#![allow(unused_imports)]
-
-use crate::ipc::gen::Schema::*;
-use flatbuffers::EndianScalar;
-use std::{cmp::Ordering, mem};
-// automatically generated by the FlatBuffers compiler, do not modify
-
-// struct Block, aligned to 8
-#[repr(transparent)]
-#[derive(Clone, Copy, PartialEq)]
-pub struct Block(pub [u8; 24]);
-impl std::fmt::Debug for Block {
- fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
- f.debug_struct("Block")
- .field("offset", &self.offset())
- .field("metaDataLength", &self.metaDataLength())
- .field("bodyLength", &self.bodyLength())
- .finish()
- }
-}
-
-impl flatbuffers::SimpleToVerifyInSlice for Block {}
-impl flatbuffers::SafeSliceAccess for Block {}
-impl<'a> flatbuffers::Follow<'a> for Block {
- type Inner = &'a Block;
- #[inline]
- fn follow(buf: &'a [u8], loc: usize) -> Self::Inner {
- <&'a Block>::follow(buf, loc)
- }
-}
-impl<'a> flatbuffers::Follow<'a> for &'a Block {
- type Inner = &'a Block;
- #[inline]
- fn follow(buf: &'a [u8], loc: usize) -> Self::Inner {
- flatbuffers::follow_cast_ref::<Block>(buf, loc)
- }
-}
-impl<'b> flatbuffers::Push for Block {
- type Output = Block;
- #[inline]
- fn push(&self, dst: &mut [u8], _rest: &[u8]) {
- let src = unsafe {
- ::std::slice::from_raw_parts(self as *const Block as *const u8, Self::size())
- };
- dst.copy_from_slice(src);
- }
-}
-impl<'b> flatbuffers::Push for &'b Block {
- type Output = Block;
-
- #[inline]
- fn push(&self, dst: &mut [u8], _rest: &[u8]) {
- let src = unsafe {
- ::std::slice::from_raw_parts(*self as *const Block as *const u8, Self::size())
- };
- dst.copy_from_slice(src);
- }
-}
-
-impl<'a> flatbuffers::Verifiable for Block {
- #[inline]
- fn run_verifier(
- v: &mut flatbuffers::Verifier,
- pos: usize,
- ) -> Result<(), flatbuffers::InvalidFlatbuffer> {
- use flatbuffers::Verifiable;
- v.in_buffer::<Self>(pos)
- }
-}
-impl Block {
- #[allow(clippy::too_many_arguments)]
- pub fn new(offset: i64, metaDataLength: i32, bodyLength: i64) -> Self {
- let mut s = Self([0; 24]);
- s.set_offset(offset);
- s.set_metaDataLength(metaDataLength);
- s.set_bodyLength(bodyLength);
- s
- }
-
- /// Index to the start of the RecordBlock (note this is past the Message header)
- pub fn offset(&self) -> i64 {
- let mut mem = core::mem::MaybeUninit::<i64>::uninit();
- unsafe {
- core::ptr::copy_nonoverlapping(
- self.0[0..].as_ptr(),
- mem.as_mut_ptr() as *mut u8,
- core::mem::size_of::<i64>(),
- );
- mem.assume_init()
- }
- .from_little_endian()
- }
-
- pub fn set_offset(&mut self, x: i64) {
- let x_le = x.to_little_endian();
- unsafe {
- core::ptr::copy_nonoverlapping(
- &x_le as *const i64 as *const u8,
- self.0[0..].as_mut_ptr(),
- core::mem::size_of::<i64>(),
- );
- }
- }
-
- /// Length of the metadata
- pub fn metaDataLength(&self) -> i32 {
- let mut mem = core::mem::MaybeUninit::<i32>::uninit();
- unsafe {
- core::ptr::copy_nonoverlapping(
- self.0[8..].as_ptr(),
- mem.as_mut_ptr() as *mut u8,
- core::mem::size_of::<i32>(),
- );
- mem.assume_init()
- }
- .from_little_endian()
- }
-
- pub fn set_metaDataLength(&mut self, x: i32) {
- let x_le = x.to_little_endian();
- unsafe {
- core::ptr::copy_nonoverlapping(
- &x_le as *const i32 as *const u8,
- self.0[8..].as_mut_ptr(),
- core::mem::size_of::<i32>(),
- );
- }
- }
-
- /// Length of the data (this is aligned so there can be a gap between this and
- /// the metadata).
- pub fn bodyLength(&self) -> i64 {
- let mut mem = core::mem::MaybeUninit::<i64>::uninit();
- unsafe {
- core::ptr::copy_nonoverlapping(
- self.0[16..].as_ptr(),
- mem.as_mut_ptr() as *mut u8,
- core::mem::size_of::<i64>(),
- );
- mem.assume_init()
- }
- .from_little_endian()
- }
-
- pub fn set_bodyLength(&mut self, x: i64) {
- let x_le = x.to_little_endian();
- unsafe {
- core::ptr::copy_nonoverlapping(
- &x_le as *const i64 as *const u8,
- self.0[16..].as_mut_ptr(),
- core::mem::size_of::<i64>(),
- );
- }
- }
-}
-
-pub enum FooterOffset {}
-#[derive(Copy, Clone, PartialEq)]
-
-/// ----------------------------------------------------------------------
-/// Arrow File metadata
-///
-pub struct Footer<'a> {
- pub _tab: flatbuffers::Table<'a>,
-}
-
-impl<'a> flatbuffers::Follow<'a> for Footer<'a> {
- type Inner = Footer<'a>;
- #[inline]
- fn follow(buf: &'a [u8], loc: usize) -> Self::Inner {
- Self {
- _tab: flatbuffers::Table { buf, loc },
- }
- }
-}
-
-impl<'a> Footer<'a> {
- #[inline]
- pub fn init_from_table(table: flatbuffers::Table<'a>) -> Self {
- Footer { _tab: table }
- }
- #[allow(unused_mut)]
- pub fn create<'bldr: 'args, 'args: 'mut_bldr, 'mut_bldr>(
- _fbb: &'mut_bldr mut flatbuffers::FlatBufferBuilder<'bldr>,
- args: &'args FooterArgs<'args>,
- ) -> flatbuffers::WIPOffset<Footer<'bldr>> {
- let mut builder = FooterBuilder::new(_fbb);
- if let Some(x) = args.custom_metadata {
- builder.add_custom_metadata(x);
- }
- if let Some(x) = args.recordBatches {
- builder.add_recordBatches(x);
- }
- if let Some(x) = args.dictionaries {
- builder.add_dictionaries(x);
- }
- if let Some(x) = args.schema {
- builder.add_schema(x);
- }
- builder.add_version(args.version);
- builder.finish()
- }
-
- pub const VT_VERSION: flatbuffers::VOffsetT = 4;
- pub const VT_SCHEMA: flatbuffers::VOffsetT = 6;
- pub const VT_DICTIONARIES: flatbuffers::VOffsetT = 8;
- pub const VT_RECORDBATCHES: flatbuffers::VOffsetT = 10;
- pub const VT_CUSTOM_METADATA: flatbuffers::VOffsetT = 12;
-
- #[inline]
- pub fn version(&self) -> MetadataVersion {
- self._tab
- .get::<MetadataVersion>(Footer::VT_VERSION, Some(MetadataVersion::V1))
- .unwrap()
- }
- #[inline]
- pub fn schema(&self) -> Option<Schema<'a>> {
- self._tab
- .get::<flatbuffers::ForwardsUOffset<Schema>>(Footer::VT_SCHEMA, None)
- }
- #[inline]
- pub fn dictionaries(&self) -> Option<&'a [Block]> {
- self._tab
- .get::<flatbuffers::ForwardsUOffset<flatbuffers::Vector<'a, Block>>>(
- Footer::VT_DICTIONARIES,
- None,
- )
- .map(|v| v.safe_slice())
- }
- #[inline]
- pub fn recordBatches(&self) -> Option<&'a [Block]> {
- self._tab
- .get::<flatbuffers::ForwardsUOffset<flatbuffers::Vector<'a, Block>>>(
- Footer::VT_RECORDBATCHES,
- None,
- )
- .map(|v| v.safe_slice())
- }
- /// User-defined metadata
- #[inline]
- pub fn custom_metadata(
- &self,
- ) -> Option<flatbuffers::Vector<'a, flatbuffers::ForwardsUOffset<KeyValue<'a>>>> {
- self._tab.get::<flatbuffers::ForwardsUOffset<
- flatbuffers::Vector<'a, flatbuffers::ForwardsUOffset<KeyValue>>,
- >>(Footer::VT_CUSTOM_METADATA, None)
- }
-}
-
-impl flatbuffers::Verifiable for Footer<'_> {
- #[inline]
- fn run_verifier(
- v: &mut flatbuffers::Verifier,
- pos: usize,
- ) -> Result<(), flatbuffers::InvalidFlatbuffer> {
- use flatbuffers::Verifiable;
- v.visit_table(pos)?
- .visit_field::<MetadataVersion>(&"version", Self::VT_VERSION, false)?
- .visit_field::<flatbuffers::ForwardsUOffset<Schema>>(
- &"schema",
- Self::VT_SCHEMA,
- false,
- )?
- .visit_field::<flatbuffers::ForwardsUOffset<flatbuffers::Vector<'_, Block>>>(
- &"dictionaries",
- Self::VT_DICTIONARIES,
- false,
- )?
- .visit_field::<flatbuffers::ForwardsUOffset<flatbuffers::Vector<'_, Block>>>(
- &"recordBatches",
- Self::VT_RECORDBATCHES,
- false,
- )?
- .visit_field::<flatbuffers::ForwardsUOffset<
- flatbuffers::Vector<'_, flatbuffers::ForwardsUOffset<KeyValue>>,
- >>(&"custom_metadata", Self::VT_CUSTOM_METADATA, false)?
- .finish();
- Ok(())
- }
-}
-pub struct FooterArgs<'a> {
- pub version: MetadataVersion,
- pub schema: Option<flatbuffers::WIPOffset<Schema<'a>>>,
- pub dictionaries: Option<flatbuffers::WIPOffset<flatbuffers::Vector<'a, Block>>>,
- pub recordBatches: Option<flatbuffers::WIPOffset<flatbuffers::Vector<'a, Block>>>,
- pub custom_metadata: Option<
- flatbuffers::WIPOffset<
- flatbuffers::Vector<'a, flatbuffers::ForwardsUOffset<KeyValue<'a>>>,
- >,
- >,
-}
-impl<'a> Default for FooterArgs<'a> {
- #[inline]
- fn default() -> Self {
- FooterArgs {
- version: MetadataVersion::V1,
- schema: None,
- dictionaries: None,
- recordBatches: None,
- custom_metadata: None,
- }
- }
-}
-pub struct FooterBuilder<'a: 'b, 'b> {
- fbb_: &'b mut flatbuffers::FlatBufferBuilder<'a>,
- start_: flatbuffers::WIPOffset<flatbuffers::TableUnfinishedWIPOffset>,
-}
-impl<'a: 'b, 'b> FooterBuilder<'a, 'b> {
- #[inline]
- pub fn add_version(&mut self, version: MetadataVersion) {
- self.fbb_.push_slot::<MetadataVersion>(
- Footer::VT_VERSION,
- version,
- MetadataVersion::V1,
- );
- }
- #[inline]
- pub fn add_schema(&mut self, schema: flatbuffers::WIPOffset<Schema<'b>>) {
- self.fbb_
- .push_slot_always::<flatbuffers::WIPOffset<Schema>>(
- Footer::VT_SCHEMA,
- schema,
- );
- }
- #[inline]
- pub fn add_dictionaries(
- &mut self,
- dictionaries: flatbuffers::WIPOffset<flatbuffers::Vector<'b, Block>>,
- ) {
- self.fbb_.push_slot_always::<flatbuffers::WIPOffset<_>>(
- Footer::VT_DICTIONARIES,
- dictionaries,
- );
- }
- #[inline]
- pub fn add_recordBatches(
- &mut self,
- recordBatches: flatbuffers::WIPOffset<flatbuffers::Vector<'b, Block>>,
- ) {
- self.fbb_.push_slot_always::<flatbuffers::WIPOffset<_>>(
- Footer::VT_RECORDBATCHES,
- recordBatches,
- );
- }
- #[inline]
- pub fn add_custom_metadata(
- &mut self,
- custom_metadata: flatbuffers::WIPOffset<
- flatbuffers::Vector<'b, flatbuffers::ForwardsUOffset<KeyValue<'b>>>,
- >,
- ) {
- self.fbb_.push_slot_always::<flatbuffers::WIPOffset<_>>(
- Footer::VT_CUSTOM_METADATA,
- custom_metadata,
- );
- }
- #[inline]
- pub fn new(
- _fbb: &'b mut flatbuffers::FlatBufferBuilder<'a>,
- ) -> FooterBuilder<'a, 'b> {
- let start = _fbb.start_table();
- FooterBuilder {
- fbb_: _fbb,
- start_: start,
- }
- }
- #[inline]
- pub fn finish(self) -> flatbuffers::WIPOffset<Footer<'a>> {
- let o = self.fbb_.end_table(self.start_);
- flatbuffers::WIPOffset::new(o.value())
- }
-}
-
-impl std::fmt::Debug for Footer<'_> {
- fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
- let mut ds = f.debug_struct("Footer");
- ds.field("version", &self.version());
- ds.field("schema", &self.schema());
- ds.field("dictionaries", &self.dictionaries());
- ds.field("recordBatches", &self.recordBatches());
- ds.field("custom_metadata", &self.custom_metadata());
- ds.finish()
- }
-}
-#[inline]
-#[deprecated(since = "2.0.0", note = "Deprecated in favor of `root_as...` methods.")]
-pub fn get_root_as_footer<'a>(buf: &'a [u8]) -> Footer<'a> {
- unsafe { flatbuffers::root_unchecked::<Footer<'a>>(buf) }
-}
-
-#[inline]
-#[deprecated(since = "2.0.0", note = "Deprecated in favor of `root_as...` methods.")]
-pub fn get_size_prefixed_root_as_footer<'a>(buf: &'a [u8]) -> Footer<'a> {
- unsafe { flatbuffers::size_prefixed_root_unchecked::<Footer<'a>>(buf) }
-}
-
-#[inline]
-/// Verifies that a buffer of bytes contains a `Footer`
-/// and returns it.
-/// Note that verification is still experimental and may not
-/// catch every error, or be maximally performant. For the
-/// previous, unchecked, behavior use
-/// `root_as_footer_unchecked`.
-pub fn root_as_footer(buf: &[u8]) -> Result<Footer, flatbuffers::InvalidFlatbuffer> {
- flatbuffers::root::<Footer>(buf)
-}
-#[inline]
-/// Verifies that a buffer of bytes contains a size prefixed
-/// `Footer` and returns it.
-/// Note that verification is still experimental and may not
-/// catch every error, or be maximally performant. For the
-/// previous, unchecked, behavior use
-/// `size_prefixed_root_as_footer_unchecked`.
-pub fn size_prefixed_root_as_footer(
- buf: &[u8],
-) -> Result<Footer, flatbuffers::InvalidFlatbuffer> {
- flatbuffers::size_prefixed_root::<Footer>(buf)
-}
-#[inline]
-/// Verifies, with the given options, that a buffer of bytes
-/// contains a `Footer` and returns it.
-/// Note that verification is still experimental and may not
-/// catch every error, or be maximally performant. For the
-/// previous, unchecked, behavior use
-/// `root_as_footer_unchecked`.
-pub fn root_as_footer_with_opts<'b, 'o>(
- opts: &'o flatbuffers::VerifierOptions,
- buf: &'b [u8],
-) -> Result<Footer<'b>, flatbuffers::InvalidFlatbuffer> {
- flatbuffers::root_with_opts::<Footer<'b>>(opts, buf)
-}
-#[inline]
-/// Verifies, with the given verifier options, that a buffer of
-/// bytes contains a size prefixed `Footer` and returns
-/// it. Note that verification is still experimental and may not
-/// catch every error, or be maximally performant. For the
-/// previous, unchecked, behavior use
-/// `root_as_footer_unchecked`.
-pub fn size_prefixed_root_as_footer_with_opts<'b, 'o>(
- opts: &'o flatbuffers::VerifierOptions,
- buf: &'b [u8],
-) -> Result<Footer<'b>, flatbuffers::InvalidFlatbuffer> {
- flatbuffers::size_prefixed_root_with_opts::<Footer<'b>>(opts, buf)
-}
-#[inline]
-/// Assumes, without verification, that a buffer of bytes contains a Footer and returns it.
-/// # Safety
-/// Callers must trust the given bytes do indeed contain a valid `Footer`.
-pub unsafe fn root_as_footer_unchecked(buf: &[u8]) -> Footer {
- flatbuffers::root_unchecked::<Footer>(buf)
-}
-#[inline]
-/// Assumes, without verification, that a buffer of bytes contains a size prefixed Footer and returns it.
-/// # Safety
-/// Callers must trust the given bytes do indeed contain a valid size prefixed `Footer`.
-pub unsafe fn size_prefixed_root_as_footer_unchecked(buf: &[u8]) -> Footer {
- flatbuffers::size_prefixed_root_unchecked::<Footer>(buf)
-}
-#[inline]
-pub fn finish_footer_buffer<'a, 'b>(
- fbb: &'b mut flatbuffers::FlatBufferBuilder<'a>,
- root: flatbuffers::WIPOffset<Footer<'a>>,
-) {
- fbb.finish(root, None);
-}
-
-#[inline]
-pub fn finish_size_prefixed_footer_buffer<'a, 'b>(
- fbb: &'b mut flatbuffers::FlatBufferBuilder<'a>,
- root: flatbuffers::WIPOffset<Footer<'a>>,
-) {
- fbb.finish_size_prefixed(root, None);
-}
diff --git a/arrow/src/ipc/gen/Message.rs b/arrow/src/ipc/gen/Message.rs
deleted file mode 100644
index 7903844..0000000
--- a/arrow/src/ipc/gen/Message.rs
+++ /dev/null
@@ -1,1346 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-#![allow(dead_code)]
-#![allow(unused_imports)]
-
-use crate::ipc::gen::Schema::*;
-use crate::ipc::gen::SparseTensor::*;
-use crate::ipc::gen::Tensor::*;
-use flatbuffers::EndianScalar;
-use std::{cmp::Ordering, mem};
-// automatically generated by the FlatBuffers compiler, do not modify
-
-#[deprecated(
- since = "2.0.0",
- note = "Use associated constants instead. This will no longer be generated in 2021."
-)]
-pub const ENUM_MIN_COMPRESSION_TYPE: i8 = 0;
-#[deprecated(
- since = "2.0.0",
- note = "Use associated constants instead. This will no longer be generated in 2021."
-)]
-pub const ENUM_MAX_COMPRESSION_TYPE: i8 = 1;
-#[deprecated(
- since = "2.0.0",
- note = "Use associated constants instead. This will no longer be generated in 2021."
-)]
-#[allow(non_camel_case_types)]
-pub const ENUM_VALUES_COMPRESSION_TYPE: [CompressionType; 2] =
- [CompressionType::LZ4_FRAME, CompressionType::ZSTD];
-
-#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
-#[repr(transparent)]
-pub struct CompressionType(pub i8);
-#[allow(non_upper_case_globals)]
-impl CompressionType {
- pub const LZ4_FRAME: Self = Self(0);
- pub const ZSTD: Self = Self(1);
-
- pub const ENUM_MIN: i8 = 0;
- pub const ENUM_MAX: i8 = 1;
- pub const ENUM_VALUES: &'static [Self] = &[Self::LZ4_FRAME, Self::ZSTD];
- /// Returns the variant's name or "" if unknown.
- pub fn variant_name(self) -> Option<&'static str> {
- match self {
- Self::LZ4_FRAME => Some("LZ4_FRAME"),
- Self::ZSTD => Some("ZSTD"),
- _ => None,
- }
- }
-}
-impl std::fmt::Debug for CompressionType {
- fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
- if let Some(name) = self.variant_name() {
- f.write_str(name)
- } else {
- f.write_fmt(format_args!("<UNKNOWN {:?}>", self.0))
- }
- }
-}
-impl<'a> flatbuffers::Follow<'a> for CompressionType {
- type Inner = Self;
- #[inline]
- fn follow(buf: &'a [u8], loc: usize) -> Self::Inner {
- let b = flatbuffers::read_scalar_at::<i8>(buf, loc);
- Self(b)
- }
-}
-
-impl flatbuffers::Push for CompressionType {
- type Output = CompressionType;
- #[inline]
- fn push(&self, dst: &mut [u8], _rest: &[u8]) {
- flatbuffers::emplace_scalar::<i8>(dst, self.0);
- }
-}
-
-impl flatbuffers::EndianScalar for CompressionType {
- #[inline]
- fn to_little_endian(self) -> Self {
- let b = i8::to_le(self.0);
- Self(b)
- }
- #[inline]
- fn from_little_endian(self) -> Self {
- let b = i8::from_le(self.0);
- Self(b)
- }
-}
-
-impl<'a> flatbuffers::Verifiable for CompressionType {
- #[inline]
- fn run_verifier(
- v: &mut flatbuffers::Verifier,
- pos: usize,
- ) -> Result<(), flatbuffers::InvalidFlatbuffer> {
- use flatbuffers::Verifiable;
- i8::run_verifier(v, pos)
- }
-}
-
-impl flatbuffers::SimpleToVerifyInSlice for CompressionType {}
-#[deprecated(
- since = "2.0.0",
- note = "Use associated constants instead. This will no longer be generated in 2021."
-)]
-pub const ENUM_MIN_BODY_COMPRESSION_METHOD: i8 = 0;
-#[deprecated(
- since = "2.0.0",
- note = "Use associated constants instead. This will no longer be generated in 2021."
-)]
-pub const ENUM_MAX_BODY_COMPRESSION_METHOD: i8 = 0;
-#[deprecated(
- since = "2.0.0",
- note = "Use associated constants instead. This will no longer be generated in 2021."
-)]
-#[allow(non_camel_case_types)]
-pub const ENUM_VALUES_BODY_COMPRESSION_METHOD: [BodyCompressionMethod; 1] =
- [BodyCompressionMethod::BUFFER];
-
-/// 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
-#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
-#[repr(transparent)]
-pub struct BodyCompressionMethod(pub i8);
-#[allow(non_upper_case_globals)]
-impl 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.
- pub const BUFFER: Self = Self(0);
-
- pub const ENUM_MIN: i8 = 0;
- pub const ENUM_MAX: i8 = 0;
- pub const ENUM_VALUES: &'static [Self] = &[Self::BUFFER];
- /// Returns the variant's name or "" if unknown.
- pub fn variant_name(self) -> Option<&'static str> {
- match self {
- Self::BUFFER => Some("BUFFER"),
- _ => None,
- }
- }
-}
-impl std::fmt::Debug for BodyCompressionMethod {
- fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
- if let Some(name) = self.variant_name() {
- f.write_str(name)
- } else {
- f.write_fmt(format_args!("<UNKNOWN {:?}>", self.0))
- }
- }
-}
-impl<'a> flatbuffers::Follow<'a> for BodyCompressionMethod {
- type Inner = Self;
- #[inline]
- fn follow(buf: &'a [u8], loc: usize) -> Self::Inner {
- let b = flatbuffers::read_scalar_at::<i8>(buf, loc);
- Self(b)
- }
-}
-
-impl flatbuffers::Push for BodyCompressionMethod {
- type Output = BodyCompressionMethod;
- #[inline]
- fn push(&self, dst: &mut [u8], _rest: &[u8]) {
- flatbuffers::emplace_scalar::<i8>(dst, self.0);
- }
-}
-
-impl flatbuffers::EndianScalar for BodyCompressionMethod {
- #[inline]
- fn to_little_endian(self) -> Self {
- let b = i8::to_le(self.0);
- Self(b)
- }
- #[inline]
- fn from_little_endian(self) -> Self {
- let b = i8::from_le(self.0);
- Self(b)
- }
-}
-
-impl<'a> flatbuffers::Verifiable for BodyCompressionMethod {
- #[inline]
- fn run_verifier(
- v: &mut flatbuffers::Verifier,
- pos: usize,
- ) -> Result<(), flatbuffers::InvalidFlatbuffer> {
- use flatbuffers::Verifiable;
- i8::run_verifier(v, pos)
- }
-}
-
-impl flatbuffers::SimpleToVerifyInSlice for BodyCompressionMethod {}
-#[deprecated(
- since = "2.0.0",
- note = "Use associated constants instead. This will no longer be generated in 2021."
-)]
-pub const ENUM_MIN_MESSAGE_HEADER: u8 = 0;
-#[deprecated(
- since = "2.0.0",
- note = "Use associated constants instead. This will no longer be generated in 2021."
-)]
-pub const ENUM_MAX_MESSAGE_HEADER: u8 = 5;
-#[deprecated(
- since = "2.0.0",
- note = "Use associated constants instead. This will no longer be generated in 2021."
-)]
-#[allow(non_camel_case_types)]
-pub const ENUM_VALUES_MESSAGE_HEADER: [MessageHeader; 6] = [
- MessageHeader::NONE,
- MessageHeader::Schema,
- MessageHeader::DictionaryBatch,
- MessageHeader::RecordBatch,
- MessageHeader::Tensor,
- MessageHeader::SparseTensor,
-];
-
-/// ----------------------------------------------------------------------
-/// 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
-#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
-#[repr(transparent)]
-pub struct MessageHeader(pub u8);
-#[allow(non_upper_case_globals)]
-impl MessageHeader {
- pub const NONE: Self = Self(0);
- pub const Schema: Self = Self(1);
- pub const DictionaryBatch: Self = Self(2);
- pub const RecordBatch: Self = Self(3);
- pub const Tensor: Self = Self(4);
- pub const SparseTensor: Self = Self(5);
-
- pub const ENUM_MIN: u8 = 0;
- pub const ENUM_MAX: u8 = 5;
- pub const ENUM_VALUES: &'static [Self] = &[
- Self::NONE,
- Self::Schema,
- Self::DictionaryBatch,
- Self::RecordBatch,
- Self::Tensor,
- Self::SparseTensor,
- ];
- /// Returns the variant's name or "" if unknown.
- pub fn variant_name(self) -> Option<&'static str> {
- match self {
- Self::NONE => Some("NONE"),
- Self::Schema => Some("Schema"),
- Self::DictionaryBatch => Some("DictionaryBatch"),
- Self::RecordBatch => Some("RecordBatch"),
- Self::Tensor => Some("Tensor"),
- Self::SparseTensor => Some("SparseTensor"),
- _ => None,
- }
- }
-}
-impl std::fmt::Debug for MessageHeader {
- fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
- if let Some(name) = self.variant_name() {
- f.write_str(name)
- } else {
- f.write_fmt(format_args!("<UNKNOWN {:?}>", self.0))
- }
- }
-}
-pub struct MessageHeaderUnionTableOffset {}
-impl<'a> flatbuffers::Follow<'a> for MessageHeader {
- type Inner = Self;
- #[inline]
- fn follow(buf: &'a [u8], loc: usize) -> Self::Inner {
- let b = flatbuffers::read_scalar_at::<u8>(buf, loc);
- Self(b)
- }
-}
-
-impl flatbuffers::Push for MessageHeader {
- type Output = MessageHeader;
- #[inline]
- fn push(&self, dst: &mut [u8], _rest: &[u8]) {
- flatbuffers::emplace_scalar::<u8>(dst, self.0);
- }
-}
-
-impl flatbuffers::EndianScalar for MessageHeader {
- #[inline]
- fn to_little_endian(self) -> Self {
- let b = u8::to_le(self.0);
- Self(b)
- }
- #[inline]
- fn from_little_endian(self) -> Self {
- let b = u8::from_le(self.0);
- Self(b)
- }
-}
-
-impl<'a> flatbuffers::Verifiable for MessageHeader {
- #[inline]
- fn run_verifier(
- v: &mut flatbuffers::Verifier,
- pos: usize,
- ) -> Result<(), flatbuffers::InvalidFlatbuffer> {
- use flatbuffers::Verifiable;
- u8::run_verifier(v, pos)
- }
-}
-
-impl flatbuffers::SimpleToVerifyInSlice for MessageHeader {}
-/// ----------------------------------------------------------------------
-/// 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<Int16> 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, aligned to 8
-#[repr(transparent)]
-#[derive(Clone, Copy, PartialEq)]
-pub struct FieldNode(pub [u8; 16]);
-impl std::fmt::Debug for FieldNode {
- fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
- f.debug_struct("FieldNode")
- .field("length", &self.length())
- .field("null_count", &self.null_count())
- .finish()
- }
-}
-
-impl flatbuffers::SimpleToVerifyInSlice for FieldNode {}
-impl flatbuffers::SafeSliceAccess for FieldNode {}
-impl<'a> flatbuffers::Follow<'a> for FieldNode {
- type Inner = &'a FieldNode;
- #[inline]
- fn follow(buf: &'a [u8], loc: usize) -> Self::Inner {
- <&'a FieldNode>::follow(buf, loc)
- }
-}
-impl<'a> flatbuffers::Follow<'a> for &'a FieldNode {
- type Inner = &'a FieldNode;
- #[inline]
- fn follow(buf: &'a [u8], loc: usize) -> Self::Inner {
- flatbuffers::follow_cast_ref::<FieldNode>(buf, loc)
- }
-}
-impl<'b> flatbuffers::Push for FieldNode {
- type Output = FieldNode;
- #[inline]
- fn push(&self, dst: &mut [u8], _rest: &[u8]) {
- let src = unsafe {
- ::std::slice::from_raw_parts(
- self as *const FieldNode as *const u8,
- Self::size(),
- )
- };
- dst.copy_from_slice(src);
- }
-}
-impl<'b> flatbuffers::Push for &'b FieldNode {
- type Output = FieldNode;
-
- #[inline]
- fn push(&self, dst: &mut [u8], _rest: &[u8]) {
- let src = unsafe {
- ::std::slice::from_raw_parts(
- *self as *const FieldNode as *const u8,
- Self::size(),
- )
- };
- dst.copy_from_slice(src);
- }
-}
-
-impl<'a> flatbuffers::Verifiable for FieldNode {
- #[inline]
- fn run_verifier(
- v: &mut flatbuffers::Verifier,
- pos: usize,
- ) -> Result<(), flatbuffers::InvalidFlatbuffer> {
- use flatbuffers::Verifiable;
- v.in_buffer::<Self>(pos)
- }
-}
-impl FieldNode {
- #[allow(clippy::too_many_arguments)]
- pub fn new(length: i64, null_count: i64) -> Self {
- let mut s = Self([0; 16]);
- s.set_length(length);
- s.set_null_count(null_count);
- s
- }
-
- /// The number of value slots in the Arrow array at this level of a nested
- /// tree
- pub fn length(&self) -> i64 {
- let mut mem = core::mem::MaybeUninit::<i64>::uninit();
- unsafe {
- core::ptr::copy_nonoverlapping(
- self.0[0..].as_ptr(),
- mem.as_mut_ptr() as *mut u8,
- core::mem::size_of::<i64>(),
- );
- mem.assume_init()
- }
- .from_little_endian()
- }
-
- pub fn set_length(&mut self, x: i64) {
- let x_le = x.to_little_endian();
- unsafe {
- core::ptr::copy_nonoverlapping(
- &x_le as *const i64 as *const u8,
- self.0[0..].as_mut_ptr(),
- core::mem::size_of::<i64>(),
- );
- }
- }
-
- /// 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.
- pub fn null_count(&self) -> i64 {
- let mut mem = core::mem::MaybeUninit::<i64>::uninit();
- unsafe {
- core::ptr::copy_nonoverlapping(
- self.0[8..].as_ptr(),
- mem.as_mut_ptr() as *mut u8,
- core::mem::size_of::<i64>(),
- );
- mem.assume_init()
- }
- .from_little_endian()
- }
-
- pub fn set_null_count(&mut self, x: i64) {
- let x_le = x.to_little_endian();
- unsafe {
- core::ptr::copy_nonoverlapping(
- &x_le as *const i64 as *const u8,
- self.0[8..].as_mut_ptr(),
- core::mem::size_of::<i64>(),
- );
- }
- }
-}
-
-pub enum BodyCompressionOffset {}
-#[derive(Copy, Clone, PartialEq)]
-
-/// Optional compression for the memory buffers constituting IPC message
-/// bodies. Intended for use with RecordBatch but could be used for other
-/// message types
-pub struct BodyCompression<'a> {
- pub _tab: flatbuffers::Table<'a>,
-}
-
-impl<'a> flatbuffers::Follow<'a> for BodyCompression<'a> {
- type Inner = BodyCompression<'a>;
- #[inline]
- fn follow(buf: &'a [u8], loc: usize) -> Self::Inner {
- Self {
- _tab: flatbuffers::Table { buf, loc },
- }
- }
-}
-
-impl<'a> BodyCompression<'a> {
- #[inline]
- pub fn init_from_table(table: flatbuffers::Table<'a>) -> Self {
- BodyCompression { _tab: table }
- }
- #[allow(unused_mut)]
- pub fn create<'bldr: 'args, 'args: 'mut_bldr, 'mut_bldr>(
- _fbb: &'mut_bldr mut flatbuffers::FlatBufferBuilder<'bldr>,
- args: &'args BodyCompressionArgs,
- ) -> flatbuffers::WIPOffset<BodyCompression<'bldr>> {
- let mut builder = BodyCompressionBuilder::new(_fbb);
- builder.add_method(args.method);
- builder.add_codec(args.codec);
- builder.finish()
- }
-
- pub const VT_CODEC: flatbuffers::VOffsetT = 4;
- pub const VT_METHOD: flatbuffers::VOffsetT = 6;
-
- /// Compressor library
- #[inline]
- pub fn codec(&self) -> CompressionType {
- self._tab
- .get::<CompressionType>(
- BodyCompression::VT_CODEC,
- Some(CompressionType::LZ4_FRAME),
- )
- .unwrap()
- }
- /// Indicates the way the record batch body was compressed
- #[inline]
- pub fn method(&self) -> BodyCompressionMethod {
- self._tab
- .get::<BodyCompressionMethod>(
- BodyCompression::VT_METHOD,
- Some(BodyCompressionMethod::BUFFER),
- )
- .unwrap()
- }
-}
-
-impl flatbuffers::Verifiable for BodyCompression<'_> {
- #[inline]
- fn run_verifier(
- v: &mut flatbuffers::Verifier,
- pos: usize,
- ) -> Result<(), flatbuffers::InvalidFlatbuffer> {
- use flatbuffers::Verifiable;
- v.visit_table(pos)?
- .visit_field::<CompressionType>(&"codec", Self::VT_CODEC, false)?
- .visit_field::<BodyCompressionMethod>(&"method", Self::VT_METHOD, false)?
- .finish();
- Ok(())
- }
-}
-pub struct BodyCompressionArgs {
- pub codec: CompressionType,
- pub method: BodyCompressionMethod,
-}
-impl<'a> Default for BodyCompressionArgs {
- #[inline]
- fn default() -> Self {
- BodyCompressionArgs {
- codec: CompressionType::LZ4_FRAME,
- method: BodyCompressionMethod::BUFFER,
- }
- }
-}
-pub struct BodyCompressionBuilder<'a: 'b, 'b> {
- fbb_: &'b mut flatbuffers::FlatBufferBuilder<'a>,
- start_: flatbuffers::WIPOffset<flatbuffers::TableUnfinishedWIPOffset>,
-}
-impl<'a: 'b, 'b> BodyCompressionBuilder<'a, 'b> {
- #[inline]
- pub fn add_codec(&mut self, codec: CompressionType) {
- self.fbb_.push_slot::<CompressionType>(
- BodyCompression::VT_CODEC,
- codec,
- CompressionType::LZ4_FRAME,
- );
- }
- #[inline]
- pub fn add_method(&mut self, method: BodyCompressionMethod) {
- self.fbb_.push_slot::<BodyCompressionMethod>(
- BodyCompression::VT_METHOD,
- method,
- BodyCompressionMethod::BUFFER,
- );
- }
- #[inline]
- pub fn new(
- _fbb: &'b mut flatbuffers::FlatBufferBuilder<'a>,
- ) -> BodyCompressionBuilder<'a, 'b> {
- let start = _fbb.start_table();
- BodyCompressionBuilder {
- fbb_: _fbb,
- start_: start,
- }
- }
- #[inline]
- pub fn finish(self) -> flatbuffers::WIPOffset<BodyCompression<'a>> {
- let o = self.fbb_.end_table(self.start_);
- flatbuffers::WIPOffset::new(o.value())
- }
-}
-
-impl std::fmt::Debug for BodyCompression<'_> {
- fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
- let mut ds = f.debug_struct("BodyCompression");
- ds.field("codec", &self.codec());
- ds.field("method", &self.method());
- ds.finish()
- }
-}
-pub enum RecordBatchOffset {}
-#[derive(Copy, Clone, PartialEq)]
-
-/// 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".
-pub struct RecordBatch<'a> {
- pub _tab: flatbuffers::Table<'a>,
-}
-
-impl<'a> flatbuffers::Follow<'a> for RecordBatch<'a> {
- type Inner = RecordBatch<'a>;
- #[inline]
- fn follow(buf: &'a [u8], loc: usize) -> Self::Inner {
- Self {
- _tab: flatbuffers::Table { buf, loc },
- }
- }
-}
-
-impl<'a> RecordBatch<'a> {
- #[inline]
- pub fn init_from_table(table: flatbuffers::Table<'a>) -> Self {
- RecordBatch { _tab: table }
- }
- #[allow(unused_mut)]
- pub fn create<'bldr: 'args, 'args: 'mut_bldr, 'mut_bldr>(
- _fbb: &'mut_bldr mut flatbuffers::FlatBufferBuilder<'bldr>,
- args: &'args RecordBatchArgs<'args>,
- ) -> flatbuffers::WIPOffset<RecordBatch<'bldr>> {
- let mut builder = RecordBatchBuilder::new(_fbb);
- builder.add_length(args.length);
- if let Some(x) = args.compression {
- builder.add_compression(x);
- }
- if let Some(x) = args.buffers {
- builder.add_buffers(x);
- }
- if let Some(x) = args.nodes {
- builder.add_nodes(x);
- }
- builder.finish()
- }
-
- pub const VT_LENGTH: flatbuffers::VOffsetT = 4;
- pub const VT_NODES: flatbuffers::VOffsetT = 6;
- pub const VT_BUFFERS: flatbuffers::VOffsetT = 8;
- pub const VT_COMPRESSION: flatbuffers::VOffsetT = 10;
-
- /// number of records / rows. The arrays in the batch should all have this
- /// length
- #[inline]
- pub fn length(&self) -> i64 {
- self._tab
- .get::<i64>(RecordBatch::VT_LENGTH, Some(0))
- .unwrap()
- }
- /// Nodes correspond to the pre-ordered flattened logical schema
- #[inline]
- pub fn nodes(&self) -> Option<&'a [FieldNode]> {
- self._tab
- .get::<flatbuffers::ForwardsUOffset<flatbuffers::Vector<'a, FieldNode>>>(
- RecordBatch::VT_NODES,
- None,
- )
- .map(|v| v.safe_slice())
- }
- /// 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
- #[inline]
- pub fn buffers(&self) -> Option<&'a [Buffer]> {
- self._tab
- .get::<flatbuffers::ForwardsUOffset<flatbuffers::Vector<'a, Buffer>>>(
- RecordBatch::VT_BUFFERS,
- None,
- )
- .map(|v| v.safe_slice())
- }
- /// Optional compression of the message body
- #[inline]
- pub fn compression(&self) -> Option<BodyCompression<'a>> {
- self._tab
- .get::<flatbuffers::ForwardsUOffset<BodyCompression>>(
- RecordBatch::VT_COMPRESSION,
- None,
- )
- }
-}
-
-impl flatbuffers::Verifiable for RecordBatch<'_> {
- #[inline]
- fn run_verifier(
- v: &mut flatbuffers::Verifier,
- pos: usize,
- ) -> Result<(), flatbuffers::InvalidFlatbuffer> {
- use flatbuffers::Verifiable;
- v.visit_table(pos)?
- .visit_field::<i64>(&"length", Self::VT_LENGTH, false)?
- .visit_field::<flatbuffers::ForwardsUOffset<flatbuffers::Vector<'_, FieldNode>>>(&"nodes", Self::VT_NODES, false)?
- .visit_field::<flatbuffers::ForwardsUOffset<flatbuffers::Vector<'_, Buffer>>>(&"buffers", Self::VT_BUFFERS, false)?
- .visit_field::<flatbuffers::ForwardsUOffset<BodyCompression>>(&"compression", Self::VT_COMPRESSION, false)?
- .finish();
- Ok(())
- }
-}
-pub struct RecordBatchArgs<'a> {
- pub length: i64,
- pub nodes: Option<flatbuffers::WIPOffset<flatbuffers::Vector<'a, FieldNode>>>,
- pub buffers: Option<flatbuffers::WIPOffset<flatbuffers::Vector<'a, Buffer>>>,
- pub compression: Option<flatbuffers::WIPOffset<BodyCompression<'a>>>,
-}
-impl<'a> Default for RecordBatchArgs<'a> {
- #[inline]
- fn default() -> Self {
- RecordBatchArgs {
- length: 0,
- nodes: None,
- buffers: None,
- compression: None,
- }
- }
-}
-pub struct RecordBatchBuilder<'a: 'b, 'b> {
- fbb_: &'b mut flatbuffers::FlatBufferBuilder<'a>,
- start_: flatbuffers::WIPOffset<flatbuffers::TableUnfinishedWIPOffset>,
-}
-impl<'a: 'b, 'b> RecordBatchBuilder<'a, 'b> {
- #[inline]
- pub fn add_length(&mut self, length: i64) {
- self.fbb_
- .push_slot::<i64>(RecordBatch::VT_LENGTH, length, 0);
- }
- #[inline]
- pub fn add_nodes(
- &mut self,
- nodes: flatbuffers::WIPOffset<flatbuffers::Vector<'b, FieldNode>>,
- ) {
- self.fbb_
- .push_slot_always::<flatbuffers::WIPOffset<_>>(RecordBatch::VT_NODES, nodes);
- }
- #[inline]
- pub fn add_buffers(
- &mut self,
- buffers: flatbuffers::WIPOffset<flatbuffers::Vector<'b, Buffer>>,
- ) {
- self.fbb_.push_slot_always::<flatbuffers::WIPOffset<_>>(
- RecordBatch::VT_BUFFERS,
- buffers,
- );
- }
- #[inline]
- pub fn add_compression(
- &mut self,
- compression: flatbuffers::WIPOffset<BodyCompression<'b>>,
- ) {
- self.fbb_
- .push_slot_always::<flatbuffers::WIPOffset<BodyCompression>>(
- RecordBatch::VT_COMPRESSION,
- compression,
- );
- }
- #[inline]
- pub fn new(
- _fbb: &'b mut flatbuffers::FlatBufferBuilder<'a>,
- ) -> RecordBatchBuilder<'a, 'b> {
- let start = _fbb.start_table();
- RecordBatchBuilder {
- fbb_: _fbb,
- start_: start,
- }
- }
- #[inline]
- pub fn finish(self) -> flatbuffers::WIPOffset<RecordBatch<'a>> {
- let o = self.fbb_.end_table(self.start_);
- flatbuffers::WIPOffset::new(o.value())
- }
-}
-
-impl std::fmt::Debug for RecordBatch<'_> {
- fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
- let mut ds = f.debug_struct("RecordBatch");
- ds.field("length", &self.length());
- ds.field("nodes", &self.nodes());
- ds.field("buffers", &self.buffers());
- ds.field("compression", &self.compression());
- ds.finish()
- }
-}
-pub enum DictionaryBatchOffset {}
-#[derive(Copy, Clone, PartialEq)]
-
-/// 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
-pub struct DictionaryBatch<'a> {
- pub _tab: flatbuffers::Table<'a>,
-}
-
-impl<'a> flatbuffers::Follow<'a> for DictionaryBatch<'a> {
- type Inner = DictionaryBatch<'a>;
- #[inline]
- fn follow(buf: &'a [u8], loc: usize) -> Self::Inner {
- Self {
- _tab: flatbuffers::Table { buf, loc },
- }
- }
-}
-
-impl<'a> DictionaryBatch<'a> {
- #[inline]
- pub fn init_from_table(table: flatbuffers::Table<'a>) -> Self {
- DictionaryBatch { _tab: table }
- }
- #[allow(unused_mut)]
- pub fn create<'bldr: 'args, 'args: 'mut_bldr, 'mut_bldr>(
- _fbb: &'mut_bldr mut flatbuffers::FlatBufferBuilder<'bldr>,
- args: &'args DictionaryBatchArgs<'args>,
- ) -> flatbuffers::WIPOffset<DictionaryBatch<'bldr>> {
- let mut builder = DictionaryBatchBuilder::new(_fbb);
- builder.add_id(args.id);
- if let Some(x) = args.data {
- builder.add_data(x);
- }
- builder.add_isDelta(args.isDelta);
- builder.finish()
- }
-
- pub const VT_ID: flatbuffers::VOffsetT = 4;
- pub const VT_DATA: flatbuffers::VOffsetT = 6;
- pub const VT_ISDELTA: flatbuffers::VOffsetT = 8;
-
- #[inline]
- pub fn id(&self) -> i64 {
- self._tab
- .get::<i64>(DictionaryBatch::VT_ID, Some(0))
- .unwrap()
- }
- #[inline]
- pub fn data(&self) -> Option<RecordBatch<'a>> {
- self._tab.get::<flatbuffers::ForwardsUOffset<RecordBatch>>(
- DictionaryBatch::VT_DATA,
- None,
- )
- }
- /// 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.
- #[inline]
- pub fn isDelta(&self) -> bool {
- self._tab
- .get::<bool>(DictionaryBatch::VT_ISDELTA, Some(false))
- .unwrap()
- }
-}
-
-impl flatbuffers::Verifiable for DictionaryBatch<'_> {
- #[inline]
- fn run_verifier(
- v: &mut flatbuffers::Verifier,
- pos: usize,
- ) -> Result<(), flatbuffers::InvalidFlatbuffer> {
- use flatbuffers::Verifiable;
- v.visit_table(pos)?
- .visit_field::<i64>(&"id", Self::VT_ID, false)?
- .visit_field::<flatbuffers::ForwardsUOffset<RecordBatch>>(
- &"data",
- Self::VT_DATA,
- false,
- )?
- .visit_field::<bool>(&"isDelta", Self::VT_ISDELTA, false)?
- .finish();
- Ok(())
- }
-}
-pub struct DictionaryBatchArgs<'a> {
- pub id: i64,
- pub data: Option<flatbuffers::WIPOffset<RecordBatch<'a>>>,
- pub isDelta: bool,
-}
-impl<'a> Default for DictionaryBatchArgs<'a> {
- #[inline]
- fn default() -> Self {
- DictionaryBatchArgs {
- id: 0,
- data: None,
- isDelta: false,
- }
- }
-}
-pub struct DictionaryBatchBuilder<'a: 'b, 'b> {
- fbb_: &'b mut flatbuffers::FlatBufferBuilder<'a>,
- start_: flatbuffers::WIPOffset<flatbuffers::TableUnfinishedWIPOffset>,
-}
-impl<'a: 'b, 'b> DictionaryBatchBuilder<'a, 'b> {
- #[inline]
- pub fn add_id(&mut self, id: i64) {
- self.fbb_.push_slot::<i64>(DictionaryBatch::VT_ID, id, 0);
- }
- #[inline]
- pub fn add_data(&mut self, data: flatbuffers::WIPOffset<RecordBatch<'b>>) {
- self.fbb_
- .push_slot_always::<flatbuffers::WIPOffset<RecordBatch>>(
- DictionaryBatch::VT_DATA,
- data,
- );
- }
- #[inline]
- pub fn add_isDelta(&mut self, isDelta: bool) {
- self.fbb_
- .push_slot::<bool>(DictionaryBatch::VT_ISDELTA, isDelta, false);
- }
- #[inline]
- pub fn new(
- _fbb: &'b mut flatbuffers::FlatBufferBuilder<'a>,
- ) -> DictionaryBatchBuilder<'a, 'b> {
- let start = _fbb.start_table();
- DictionaryBatchBuilder {
- fbb_: _fbb,
- start_: start,
- }
- }
- #[inline]
- pub fn finish(self) -> flatbuffers::WIPOffset<DictionaryBatch<'a>> {
- let o = self.fbb_.end_table(self.start_);
- flatbuffers::WIPOffset::new(o.value())
- }
-}
-
-impl std::fmt::Debug for DictionaryBatch<'_> {
- fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
- let mut ds = f.debug_struct("DictionaryBatch");
- ds.field("id", &self.id());
- ds.field("data", &self.data());
- ds.field("isDelta", &self.isDelta());
- ds.finish()
- }
-}
-pub enum MessageOffset {}
-#[derive(Copy, Clone, PartialEq)]
-
-pub struct Message<'a> {
- pub _tab: flatbuffers::Table<'a>,
-}
-
-impl<'a> flatbuffers::Follow<'a> for Message<'a> {
- type Inner = Message<'a>;
- #[inline]
- fn follow(buf: &'a [u8], loc: usize) -> Self::Inner {
- Self {
- _tab: flatbuffers::Table { buf, loc },
- }
- }
-}
-
-impl<'a> Message<'a> {
- #[inline]
- pub fn init_from_table(table: flatbuffers::Table<'a>) -> Self {
- Message { _tab: table }
- }
- #[allow(unused_mut)]
- pub fn create<'bldr: 'args, 'args: 'mut_bldr, 'mut_bldr>(
- _fbb: &'mut_bldr mut flatbuffers::FlatBufferBuilder<'bldr>,
- args: &'args MessageArgs<'args>,
- ) -> flatbuffers::WIPOffset<Message<'bldr>> {
- let mut builder = MessageBuilder::new(_fbb);
- builder.add_bodyLength(args.bodyLength);
- if let Some(x) = args.custom_metadata {
- builder.add_custom_metadata(x);
- }
- if let Some(x) = args.header {
- builder.add_header(x);
- }
- builder.add_version(args.version);
- builder.add_header_type(args.header_type);
- builder.finish()
- }
-
- pub const VT_VERSION: flatbuffers::VOffsetT = 4;
- pub const VT_HEADER_TYPE: flatbuffers::VOffsetT = 6;
- pub const VT_HEADER: flatbuffers::VOffsetT = 8;
- pub const VT_BODYLENGTH: flatbuffers::VOffsetT = 10;
- pub const VT_CUSTOM_METADATA: flatbuffers::VOffsetT = 12;
-
- #[inline]
- pub fn version(&self) -> MetadataVersion {
- self._tab
- .get::<MetadataVersion>(Message::VT_VERSION, Some(MetadataVersion::V1))
- .unwrap()
- }
- #[inline]
- pub fn header_type(&self) -> MessageHeader {
- self._tab
- .get::<MessageHeader>(Message::VT_HEADER_TYPE, Some(MessageHeader::NONE))
- .unwrap()
- }
- #[inline]
- pub fn header(&self) -> Option<flatbuffers::Table<'a>> {
- self._tab
- .get::<flatbuffers::ForwardsUOffset<flatbuffers::Table<'a>>>(
- Message::VT_HEADER,
- None,
- )
- }
- #[inline]
- pub fn bodyLength(&self) -> i64 {
- self._tab
- .get::<i64>(Message::VT_BODYLENGTH, Some(0))
- .unwrap()
- }
- #[inline]
- pub fn custom_metadata(
- &self,
- ) -> Option<flatbuffers::Vector<'a, flatbuffers::ForwardsUOffset<KeyValue<'a>>>> {
- self._tab.get::<flatbuffers::ForwardsUOffset<
- flatbuffers::Vector<'a, flatbuffers::ForwardsUOffset<KeyValue>>,
- >>(Message::VT_CUSTOM_METADATA, None)
- }
- #[inline]
- #[allow(non_snake_case)]
- pub fn header_as_schema(&self) -> Option<Schema<'a>> {
- if self.header_type() == MessageHeader::Schema {
- self.header().map(Schema::init_from_table)
- } else {
- None
- }
- }
-
- #[inline]
- #[allow(non_snake_case)]
- pub fn header_as_dictionary_batch(&self) -> Option<DictionaryBatch<'a>> {
- if self.header_type() == MessageHeader::DictionaryBatch {
- self.header().map(DictionaryBatch::init_from_table)
- } else {
- None
- }
- }
-
- #[inline]
- #[allow(non_snake_case)]
- pub fn header_as_record_batch(&self) -> Option<RecordBatch<'a>> {
- if self.header_type() == MessageHeader::RecordBatch {
- self.header().map(RecordBatch::init_from_table)
- } else {
- None
- }
- }
-
- #[inline]
- #[allow(non_snake_case)]
- pub fn header_as_tensor(&self) -> Option<Tensor<'a>> {
- if self.header_type() == MessageHeader::Tensor {
- self.header().map(Tensor::init_from_table)
- } else {
- None
- }
- }
-
- #[inline]
- #[allow(non_snake_case)]
- pub fn header_as_sparse_tensor(&self) -> Option<SparseTensor<'a>> {
- if self.header_type() == MessageHeader::SparseTensor {
- self.header().map(SparseTensor::init_from_table)
- } else {
- None
- }
- }
-}
-
-impl flatbuffers::Verifiable for Message<'_> {
- #[inline]
- fn run_verifier(
- v: &mut flatbuffers::Verifier,
- pos: usize,
- ) -> Result<(), flatbuffers::InvalidFlatbuffer> {
- use flatbuffers::Verifiable;
- v.visit_table(pos)?
- .visit_field::<MetadataVersion>(&"version", Self::VT_VERSION, false)?
- .visit_union::<MessageHeader, _>(&"header_type", Self::VT_HEADER_TYPE, &"header", Self::VT_HEADER, false, |key, v, pos| {
- match key {
- MessageHeader::Schema => v.verify_union_variant::<flatbuffers::ForwardsUOffset<Schema>>("MessageHeader::Schema", pos),
- MessageHeader::DictionaryBatch => v.verify_union_variant::<flatbuffers::ForwardsUOffset<DictionaryBatch>>("MessageHeader::DictionaryBatch", pos),
- MessageHeader::RecordBatch => v.verify_union_variant::<flatbuffers::ForwardsUOffset<RecordBatch>>("MessageHeader::RecordBatch", pos),
- MessageHeader::Tensor => v.verify_union_variant::<flatbuffers::ForwardsUOffset<Tensor>>("MessageHeader::Tensor", pos),
- MessageHeader::SparseTensor => v.verify_union_variant::<flatbuffers::ForwardsUOffset<SparseTensor>>("MessageHeader::SparseTensor", pos),
- _ => Ok(()),
- }
- })?
- .visit_field::<i64>(&"bodyLength", Self::VT_BODYLENGTH, false)?
- .visit_field::<flatbuffers::ForwardsUOffset<flatbuffers::Vector<'_, flatbuffers::ForwardsUOffset<KeyValue>>>>(&"custom_metadata", Self::VT_CUSTOM_METADATA, false)?
- .finish();
- Ok(())
- }
-}
-pub struct MessageArgs<'a> {
- pub version: MetadataVersion,
- pub header_type: MessageHeader,
- pub header: Option<flatbuffers::WIPOffset<flatbuffers::UnionWIPOffset>>,
- pub bodyLength: i64,
- pub custom_metadata: Option<
- flatbuffers::WIPOffset<
- flatbuffers::Vector<'a, flatbuffers::ForwardsUOffset<KeyValue<'a>>>,
- >,
- >,
-}
-impl<'a> Default for MessageArgs<'a> {
- #[inline]
- fn default() -> Self {
- MessageArgs {
- version: MetadataVersion::V1,
- header_type: MessageHeader::NONE,
- header: None,
- bodyLength: 0,
- custom_metadata: None,
- }
- }
-}
-pub struct MessageBuilder<'a: 'b, 'b> {
- fbb_: &'b mut flatbuffers::FlatBufferBuilder<'a>,
- start_: flatbuffers::WIPOffset<flatbuffers::TableUnfinishedWIPOffset>,
-}
-impl<'a: 'b, 'b> MessageBuilder<'a, 'b> {
- #[inline]
- pub fn add_version(&mut self, version: MetadataVersion) {
- self.fbb_.push_slot::<MetadataVersion>(
- Message::VT_VERSION,
- version,
- MetadataVersion::V1,
- );
- }
- #[inline]
- pub fn add_header_type(&mut self, header_type: MessageHeader) {
- self.fbb_.push_slot::<MessageHeader>(
- Message::VT_HEADER_TYPE,
- header_type,
- MessageHeader::NONE,
- );
- }
- #[inline]
- pub fn add_header(
- &mut self,
- header: flatbuffers::WIPOffset<flatbuffers::UnionWIPOffset>,
- ) {
- self.fbb_
- .push_slot_always::<flatbuffers::WIPOffset<_>>(Message::VT_HEADER, header);
- }
- #[inline]
- pub fn add_bodyLength(&mut self, bodyLength: i64) {
- self.fbb_
- .push_slot::<i64>(Message::VT_BODYLENGTH, bodyLength, 0);
- }
- #[inline]
- pub fn add_custom_metadata(
- &mut self,
- custom_metadata: flatbuffers::WIPOffset<
- flatbuffers::Vector<'b, flatbuffers::ForwardsUOffset<KeyValue<'b>>>,
- >,
- ) {
- self.fbb_.push_slot_always::<flatbuffers::WIPOffset<_>>(
- Message::VT_CUSTOM_METADATA,
- custom_metadata,
- );
- }
- #[inline]
- pub fn new(
- _fbb: &'b mut flatbuffers::FlatBufferBuilder<'a>,
- ) -> MessageBuilder<'a, 'b> {
- let start = _fbb.start_table();
- MessageBuilder {
- fbb_: _fbb,
- start_: start,
- }
- }
- #[inline]
- pub fn finish(self) -> flatbuffers::WIPOffset<Message<'a>> {
- let o = self.fbb_.end_table(self.start_);
- flatbuffers::WIPOffset::new(o.value())
- }
-}
-
-impl std::fmt::Debug for Message<'_> {
- fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
- let mut ds = f.debug_struct("Message");
- ds.field("version", &self.version());
- ds.field("header_type", &self.header_type());
- match self.header_type() {
- MessageHeader::Schema => {
- if let Some(x) = self.header_as_schema() {
- ds.field("header", &x)
- } else {
- ds.field(
- "header",
- &"InvalidFlatbuffer: Union discriminant does not match value.",
- )
- }
- }
- MessageHeader::DictionaryBatch => {
- if let Some(x) = self.header_as_dictionary_batch() {
- ds.field("header", &x)
- } else {
- ds.field(
- "header",
- &"InvalidFlatbuffer: Union discriminant does not match value.",
- )
- }
- }
- MessageHeader::RecordBatch => {
- if let Some(x) = self.header_as_record_batch() {
- ds.field("header", &x)
- } else {
- ds.field(
- "header",
- &"InvalidFlatbuffer: Union discriminant does not match value.",
- )
- }
- }
- MessageHeader::Tensor => {
- if let Some(x) = self.header_as_tensor() {
- ds.field("header", &x)
- } else {
- ds.field(
- "header",
- &"InvalidFlatbuffer: Union discriminant does not match value.",
- )
- }
- }
- MessageHeader::SparseTensor => {
- if let Some(x) = self.header_as_sparse_tensor() {
- ds.field("header", &x)
- } else {
- ds.field(
- "header",
- &"InvalidFlatbuffer: Union discriminant does not match value.",
- )
- }
- }
- _ => {
- let x: Option<()> = None;
- ds.field("header", &x)
- }
- };
- ds.field("bodyLength", &self.bodyLength());
- ds.field("custom_metadata", &self.custom_metadata());
- ds.finish()
- }
-}
-#[inline]
-#[deprecated(since = "2.0.0", note = "Deprecated in favor of `root_as...` methods.")]
-pub fn get_root_as_message<'a>(buf: &'a [u8]) -> Message<'a> {
- unsafe { flatbuffers::root_unchecked::<Message<'a>>(buf) }
-}
-
-#[inline]
-#[deprecated(since = "2.0.0", note = "Deprecated in favor of `root_as...` methods.")]
-pub fn get_size_prefixed_root_as_message<'a>(buf: &'a [u8]) -> Message<'a> {
- unsafe { flatbuffers::size_prefixed_root_unchecked::<Message<'a>>(buf) }
-}
-
-#[inline]
-/// Verifies that a buffer of bytes contains a `Message`
-/// and returns it.
-/// Note that verification is still experimental and may not
-/// catch every error, or be maximally performant. For the
-/// previous, unchecked, behavior use
-/// `root_as_message_unchecked`.
-pub fn root_as_message(buf: &[u8]) -> Result<Message, flatbuffers::InvalidFlatbuffer> {
- flatbuffers::root::<Message>(buf)
-}
-#[inline]
-/// Verifies that a buffer of bytes contains a size prefixed
-/// `Message` and returns it.
-/// Note that verification is still experimental and may not
-/// catch every error, or be maximally performant. For the
-/// previous, unchecked, behavior use
-/// `size_prefixed_root_as_message_unchecked`.
-pub fn size_prefixed_root_as_message(
- buf: &[u8],
-) -> Result<Message, flatbuffers::InvalidFlatbuffer> {
- flatbuffers::size_prefixed_root::<Message>(buf)
-}
-#[inline]
-/// Verifies, with the given options, that a buffer of bytes
-/// contains a `Message` and returns it.
-/// Note that verification is still experimental and may not
-/// catch every error, or be maximally performant. For the
-/// previous, unchecked, behavior use
-/// `root_as_message_unchecked`.
-pub fn root_as_message_with_opts<'b, 'o>(
- opts: &'o flatbuffers::VerifierOptions,
- buf: &'b [u8],
-) -> Result<Message<'b>, flatbuffers::InvalidFlatbuffer> {
- flatbuffers::root_with_opts::<Message<'b>>(opts, buf)
-}
-#[inline]
-/// Verifies, with the given verifier options, that a buffer of
-/// bytes contains a size prefixed `Message` and returns
-/// it. Note that verification is still experimental and may not
-/// catch every error, or be maximally performant. For the
-/// previous, unchecked, behavior use
-/// `root_as_message_unchecked`.
-pub fn size_prefixed_root_as_message_with_opts<'b, 'o>(
- opts: &'o flatbuffers::VerifierOptions,
- buf: &'b [u8],
-) -> Result<Message<'b>, flatbuffers::InvalidFlatbuffer> {
- flatbuffers::size_prefixed_root_with_opts::<Message<'b>>(opts, buf)
-}
-#[inline]
-/// Assumes, without verification, that a buffer of bytes contains a Message and returns it.
-/// # Safety
-/// Callers must trust the given bytes do indeed contain a valid `Message`.
-pub unsafe fn root_as_message_unchecked(buf: &[u8]) -> Message {
- flatbuffers::root_unchecked::<Message>(buf)
-}
-#[inline]
-/// Assumes, without verification, that a buffer of bytes contains a size prefixed Message and returns it.
-/// # Safety
-/// Callers must trust the given bytes do indeed contain a valid size prefixed `Message`.
-pub unsafe fn size_prefixed_root_as_message_unchecked(buf: &[u8]) -> Message {
- flatbuffers::size_prefixed_root_unchecked::<Message>(buf)
-}
-#[inline]
-pub fn finish_message_buffer<'a, 'b>(
- fbb: &'b mut flatbuffers::FlatBufferBuilder<'a>,
- root: flatbuffers::WIPOffset<Message<'a>>,
-) {
- fbb.finish(root, None);
-}
-
-#[inline]
-pub fn finish_size_prefixed_message_buffer<'a, 'b>(
- fbb: &'b mut flatbuffers::FlatBufferBuilder<'a>,
- root: flatbuffers::WIPOffset<Message<'a>>,
-) {
- fbb.finish_size_prefixed(root, None);
-}
diff --git a/arrow/src/ipc/gen/Schema.rs b/arrow/src/ipc/gen/Schema.rs
deleted file mode 100644
index f37f920..0000000
--- a/arrow/src/ipc/gen/Schema.rs
+++ /dev/null
@@ -1,4586 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-#![allow(dead_code)]
-#![allow(unused_imports)]
-
-use flatbuffers::EndianScalar;
-use std::{cmp::Ordering, mem};
-// automatically generated by the FlatBuffers compiler, do not modify
-
-#[deprecated(
- since = "2.0.0",
- note = "Use associated constants instead. This will no longer be generated in 2021."
-)]
-pub const ENUM_MIN_METADATA_VERSION: i16 = 0;
-#[deprecated(
- since = "2.0.0",
- note = "Use associated constants instead. This will no longer be generated in 2021."
-)]
-pub const ENUM_MAX_METADATA_VERSION: i16 = 4;
-#[deprecated(
- since = "2.0.0",
- note = "Use associated constants instead. This will no longer be generated in 2021."
-)]
-#[allow(non_camel_case_types)]
-pub const ENUM_VALUES_METADATA_VERSION: [MetadataVersion; 5] = [
- MetadataVersion::V1,
- MetadataVersion::V2,
- MetadataVersion::V3,
- MetadataVersion::V4,
- MetadataVersion::V5,
-];
-
-#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
-#[repr(transparent)]
-pub struct MetadataVersion(pub i16);
-#[allow(non_upper_case_globals)]
-impl MetadataVersion {
- /// 0.1.0 (October 2016).
- pub const V1: Self = Self(0);
- /// 0.2.0 (February 2017). Non-backwards compatible with V1.
- pub const V2: Self = Self(1);
- /// 0.3.0 -> 0.7.1 (May - December 2017). Non-backwards compatible with V2.
- pub const V3: Self = Self(2);
- /// >= 0.8.0 (December 2017). Non-backwards compatible with V3.
- pub const V4: Self = Self(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.
- pub const V5: Self = Self(4);
-
- pub const ENUM_MIN: i16 = 0;
- pub const ENUM_MAX: i16 = 4;
- pub const ENUM_VALUES: &'static [Self] =
- &[Self::V1, Self::V2, Self::V3, Self::V4, Self::V5];
- /// Returns the variant's name or "" if unknown.
- pub fn variant_name(self) -> Option<&'static str> {
- match self {
- Self::V1 => Some("V1"),
- Self::V2 => Some("V2"),
- Self::V3 => Some("V3"),
- Self::V4 => Some("V4"),
- Self::V5 => Some("V5"),
- _ => None,
- }
- }
-}
-impl std::fmt::Debug for MetadataVersion {
- fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
- if let Some(name) = self.variant_name() {
- f.write_str(name)
- } else {
- f.write_fmt(format_args!("<UNKNOWN {:?}>", self.0))
- }
- }
-}
-impl<'a> flatbuffers::Follow<'a> for MetadataVersion {
- type Inner = Self;
- #[inline]
- fn follow(buf: &'a [u8], loc: usize) -> Self::Inner {
- let b = flatbuffers::read_scalar_at::<i16>(buf, loc);
- Self(b)
- }
-}
-
-impl flatbuffers::Push for MetadataVersion {
- type Output = MetadataVersion;
- #[inline]
- fn push(&self, dst: &mut [u8], _rest: &[u8]) {
- flatbuffers::emplace_scalar::<i16>(dst, self.0);
- }
-}
-
-impl flatbuffers::EndianScalar for MetadataVersion {
- #[inline]
- fn to_little_endian(self) -> Self {
- let b = i16::to_le(self.0);
- Self(b)
- }
- #[inline]
- fn from_little_endian(self) -> Self {
- let b = i16::from_le(self.0);
- Self(b)
- }
-}
-
-impl<'a> flatbuffers::Verifiable for MetadataVersion {
- #[inline]
- fn run_verifier(
- v: &mut flatbuffers::Verifier,
- pos: usize,
- ) -> Result<(), flatbuffers::InvalidFlatbuffer> {
- use flatbuffers::Verifiable;
- i16::run_verifier(v, pos)
- }
-}
-
-impl flatbuffers::SimpleToVerifyInSlice for MetadataVersion {}
-#[deprecated(
- since = "2.0.0",
- note = "Use associated constants instead. This will no longer be generated in 2021."
-)]
-pub const ENUM_MIN_FEATURE: i64 = 0;
-#[deprecated(
- since = "2.0.0",
- note = "Use associated constants instead. This will no longer be generated in 2021."
-)]
-pub const ENUM_MAX_FEATURE: i64 = 2;
-#[deprecated(
- since = "2.0.0",
- note = "Use associated constants instead. This will no longer be generated in 2021."
-)]
-#[allow(non_camel_case_types)]
-pub const ENUM_VALUES_FEATURE: [Feature; 3] = [
- Feature::UNUSED,
- Feature::DICTIONARY_REPLACEMENT,
- Feature::COMPRESSED_BODY,
-];
-
-/// 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.
-#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
-#[repr(transparent)]
-pub struct Feature(pub i64);
-#[allow(non_upper_case_globals)]
-impl Feature {
- /// Needed to make flatbuffers happy.
- pub const UNUSED: Self = Self(0);
- /// The stream makes use of multiple full dictionaries with the
- /// same ID and assumes clients implement dictionary replacement
- /// correctly.
- pub const DICTIONARY_REPLACEMENT: Self = Self(1);
- /// The stream makes use of compressed bodies as described
- /// in Message.fbs.
- pub const COMPRESSED_BODY: Self = Self(2);
-
- pub const ENUM_MIN: i64 = 0;
- pub const ENUM_MAX: i64 = 2;
- pub const ENUM_VALUES: &'static [Self] = &[
- Self::UNUSED,
- Self::DICTIONARY_REPLACEMENT,
- Self::COMPRESSED_BODY,
- ];
- /// Returns the variant's name or "" if unknown.
- pub fn variant_name(self) -> Option<&'static str> {
- match self {
- Self::UNUSED => Some("UNUSED"),
- Self::DICTIONARY_REPLACEMENT => Some("DICTIONARY_REPLACEMENT"),
- Self::COMPRESSED_BODY => Some("COMPRESSED_BODY"),
- _ => None,
- }
- }
-}
-impl std::fmt::Debug for Feature {
- fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
- if let Some(name) = self.variant_name() {
- f.write_str(name)
- } else {
- f.write_fmt(format_args!("<UNKNOWN {:?}>", self.0))
- }
- }
-}
-impl<'a> flatbuffers::Follow<'a> for Feature {
- type Inner = Self;
- #[inline]
- fn follow(buf: &'a [u8], loc: usize) -> Self::Inner {
- let b = flatbuffers::read_scalar_at::<i64>(buf, loc);
- Self(b)
- }
-}
-
-impl flatbuffers::Push for Feature {
- type Output = Feature;
- #[inline]
- fn push(&self, dst: &mut [u8], _rest: &[u8]) {
- flatbuffers::emplace_scalar::<i64>(dst, self.0);
- }
-}
-
-impl flatbuffers::EndianScalar for Feature {
- #[inline]
- fn to_little_endian(self) -> Self {
- let b = i64::to_le(self.0);
- Self(b)
- }
- #[inline]
- fn from_little_endian(self) -> Self {
- let b = i64::from_le(self.0);
- Self(b)
- }
-}
-
-impl<'a> flatbuffers::Verifiable for Feature {
- #[inline]
- fn run_verifier(
- v: &mut flatbuffers::Verifier,
- pos: usize,
- ) -> Result<(), flatbuffers::InvalidFlatbuffer> {
- use flatbuffers::Verifiable;
- i64::run_verifier(v, pos)
- }
-}
-
-impl flatbuffers::SimpleToVerifyInSlice for Feature {}
-#[deprecated(
- since = "2.0.0",
- note = "Use associated constants instead. This will no longer be generated in 2021."
-)]
-pub const ENUM_MIN_UNION_MODE: i16 = 0;
-#[deprecated(
- since = "2.0.0",
- note = "Use associated constants instead. This will no longer be generated in 2021."
-)]
-pub const ENUM_MAX_UNION_MODE: i16 = 1;
-#[deprecated(
- since = "2.0.0",
- note = "Use associated constants instead. This will no longer be generated in 2021."
-)]
-#[allow(non_camel_case_types)]
-pub const ENUM_VALUES_UNION_MODE: [UnionMode; 2] = [UnionMode::Sparse, UnionMode::Dense];
-
-#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
-#[repr(transparent)]
-pub struct UnionMode(pub i16);
-#[allow(non_upper_case_globals)]
-impl UnionMode {
- pub const Sparse: Self = Self(0);
- pub const Dense: Self = Self(1);
-
- pub const ENUM_MIN: i16 = 0;
- pub const ENUM_MAX: i16 = 1;
- pub const ENUM_VALUES: &'static [Self] = &[Self::Sparse, Self::Dense];
- /// Returns the variant's name or "" if unknown.
- pub fn variant_name(self) -> Option<&'static str> {
- match self {
- Self::Sparse => Some("Sparse"),
- Self::Dense => Some("Dense"),
- _ => None,
- }
- }
-}
-impl std::fmt::Debug for UnionMode {
- fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
- if let Some(name) = self.variant_name() {
- f.write_str(name)
- } else {
- f.write_fmt(format_args!("<UNKNOWN {:?}>", self.0))
- }
- }
-}
-impl<'a> flatbuffers::Follow<'a> for UnionMode {
- type Inner = Self;
- #[inline]
- fn follow(buf: &'a [u8], loc: usize) -> Self::Inner {
- let b = flatbuffers::read_scalar_at::<i16>(buf, loc);
- Self(b)
- }
-}
-
-impl flatbuffers::Push for UnionMode {
- type Output = UnionMode;
- #[inline]
- fn push(&self, dst: &mut [u8], _rest: &[u8]) {
- flatbuffers::emplace_scalar::<i16>(dst, self.0);
- }
-}
-
-impl flatbuffers::EndianScalar for UnionMode {
- #[inline]
- fn to_little_endian(self) -> Self {
- let b = i16::to_le(self.0);
- Self(b)
- }
- #[inline]
- fn from_little_endian(self) -> Self {
- let b = i16::from_le(self.0);
- Self(b)
- }
-}
-
-impl<'a> flatbuffers::Verifiable for UnionMode {
- #[inline]
- fn run_verifier(
- v: &mut flatbuffers::Verifier,
- pos: usize,
- ) -> Result<(), flatbuffers::InvalidFlatbuffer> {
- use flatbuffers::Verifiable;
- i16::run_verifier(v, pos)
- }
-}
-
-impl flatbuffers::SimpleToVerifyInSlice for UnionMode {}
-#[deprecated(
- since = "2.0.0",
- note = "Use associated constants instead. This will no longer be generated in 2021."
-)]
-pub const ENUM_MIN_PRECISION: i16 = 0;
-#[deprecated(
- since = "2.0.0",
- note = "Use associated constants instead. This will no longer be generated in 2021."
-)]
-pub const ENUM_MAX_PRECISION: i16 = 2;
-#[deprecated(
- since = "2.0.0",
- note = "Use associated constants instead. This will no longer be generated in 2021."
-)]
-#[allow(non_camel_case_types)]
-pub const ENUM_VALUES_PRECISION: [Precision; 3] =
- [Precision::HALF, Precision::SINGLE, Precision::DOUBLE];
-
-#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
-#[repr(transparent)]
-pub struct Precision(pub i16);
-#[allow(non_upper_case_globals)]
-impl Precision {
- pub const HALF: Self = Self(0);
- pub const SINGLE: Self = Self(1);
- pub const DOUBLE: Self = Self(2);
-
- pub const ENUM_MIN: i16 = 0;
- pub const ENUM_MAX: i16 = 2;
- pub const ENUM_VALUES: &'static [Self] = &[Self::HALF, Self::SINGLE, Self::DOUBLE];
- /// Returns the variant's name or "" if unknown.
- pub fn variant_name(self) -> Option<&'static str> {
- match self {
- Self::HALF => Some("HALF"),
- Self::SINGLE => Some("SINGLE"),
- Self::DOUBLE => Some("DOUBLE"),
- _ => None,
- }
- }
-}
-impl std::fmt::Debug for Precision {
- fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
- if let Some(name) = self.variant_name() {
- f.write_str(name)
- } else {
- f.write_fmt(format_args!("<UNKNOWN {:?}>", self.0))
- }
- }
-}
-impl<'a> flatbuffers::Follow<'a> for Precision {
- type Inner = Self;
- #[inline]
- fn follow(buf: &'a [u8], loc: usize) -> Self::Inner {
- let b = flatbuffers::read_scalar_at::<i16>(buf, loc);
- Self(b)
- }
-}
-
-impl flatbuffers::Push for Precision {
- type Output = Precision;
- #[inline]
- fn push(&self, dst: &mut [u8], _rest: &[u8]) {
- flatbuffers::emplace_scalar::<i16>(dst, self.0);
- }
-}
-
-impl flatbuffers::EndianScalar for Precision {
- #[inline]
- fn to_little_endian(self) -> Self {
- let b = i16::to_le(self.0);
- Self(b)
- }
- #[inline]
- fn from_little_endian(self) -> Self {
- let b = i16::from_le(self.0);
- Self(b)
- }
-}
-
-impl<'a> flatbuffers::Verifiable for Precision {
- #[inline]
- fn run_verifier(
- v: &mut flatbuffers::Verifier,
- pos: usize,
- ) -> Result<(), flatbuffers::InvalidFlatbuffer> {
- use flatbuffers::Verifiable;
- i16::run_verifier(v, pos)
- }
-}
-
-impl flatbuffers::SimpleToVerifyInSlice for Precision {}
-#[deprecated(
- since = "2.0.0",
- note = "Use associated constants instead. This will no longer be generated in 2021."
-)]
-pub const ENUM_MIN_DATE_UNIT: i16 = 0;
-#[deprecated(
- since = "2.0.0",
- note = "Use associated constants instead. This will no longer be generated in 2021."
-)]
-pub const ENUM_MAX_DATE_UNIT: i16 = 1;
-#[deprecated(
- since = "2.0.0",
- note = "Use associated constants instead. This will no longer be generated in 2021."
-)]
-#[allow(non_camel_case_types)]
-pub const ENUM_VALUES_DATE_UNIT: [DateUnit; 2] = [DateUnit::DAY, DateUnit::MILLISECOND];
-
-#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
-#[repr(transparent)]
-pub struct DateUnit(pub i16);
-#[allow(non_upper_case_globals)]
-impl DateUnit {
- pub const DAY: Self = Self(0);
- pub const MILLISECOND: Self = Self(1);
-
- pub const ENUM_MIN: i16 = 0;
- pub const ENUM_MAX: i16 = 1;
- pub const ENUM_VALUES: &'static [Self] = &[Self::DAY, Self::MILLISECOND];
- /// Returns the variant's name or "" if unknown.
- pub fn variant_name(self) -> Option<&'static str> {
- match self {
- Self::DAY => Some("DAY"),
- Self::MILLISECOND => Some("MILLISECOND"),
- _ => None,
- }
- }
-}
-impl std::fmt::Debug for DateUnit {
- fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
- if let Some(name) = self.variant_name() {
- f.write_str(name)
- } else {
- f.write_fmt(format_args!("<UNKNOWN {:?}>", self.0))
- }
- }
-}
-impl<'a> flatbuffers::Follow<'a> for DateUnit {
- type Inner = Self;
- #[inline]
- fn follow(buf: &'a [u8], loc: usize) -> Self::Inner {
- let b = flatbuffers::read_scalar_at::<i16>(buf, loc);
- Self(b)
- }
-}
-
-impl flatbuffers::Push for DateUnit {
- type Output = DateUnit;
- #[inline]
- fn push(&self, dst: &mut [u8], _rest: &[u8]) {
- flatbuffers::emplace_scalar::<i16>(dst, self.0);
- }
-}
-
-impl flatbuffers::EndianScalar for DateUnit {
- #[inline]
- fn to_little_endian(self) -> Self {
- let b = i16::to_le(self.0);
- Self(b)
- }
- #[inline]
- fn from_little_endian(self) -> Self {
- let b = i16::from_le(self.0);
- Self(b)
- }
-}
-
-impl<'a> flatbuffers::Verifiable for DateUnit {
- #[inline]
- fn run_verifier(
- v: &mut flatbuffers::Verifier,
- pos: usize,
- ) -> Result<(), flatbuffers::InvalidFlatbuffer> {
- use flatbuffers::Verifiable;
- i16::run_verifier(v, pos)
- }
-}
-
-impl flatbuffers::SimpleToVerifyInSlice for DateUnit {}
-#[deprecated(
- since = "2.0.0",
- note = "Use associated constants instead. This will no longer be generated in 2021."
-)]
-pub const ENUM_MIN_TIME_UNIT: i16 = 0;
-#[deprecated(
- since = "2.0.0",
- note = "Use associated constants instead. This will no longer be generated in 2021."
-)]
-pub const ENUM_MAX_TIME_UNIT: i16 = 3;
-#[deprecated(
- since = "2.0.0",
- note = "Use associated constants instead. This will no longer be generated in 2021."
-)]
-#[allow(non_camel_case_types)]
-pub const ENUM_VALUES_TIME_UNIT: [TimeUnit; 4] = [
- TimeUnit::SECOND,
- TimeUnit::MILLISECOND,
- TimeUnit::MICROSECOND,
- TimeUnit::NANOSECOND,
-];
-
-#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
-#[repr(transparent)]
-pub struct TimeUnit(pub i16);
-#[allow(non_upper_case_globals)]
-impl TimeUnit {
- pub const SECOND: Self = Self(0);
- pub const MILLISECOND: Self = Self(1);
- pub const MICROSECOND: Self = Self(2);
- pub const NANOSECOND: Self = Self(3);
-
- pub const ENUM_MIN: i16 = 0;
- pub const ENUM_MAX: i16 = 3;
- pub const ENUM_VALUES: &'static [Self] = &[
- Self::SECOND,
- Self::MILLISECOND,
- Self::MICROSECOND,
- Self::NANOSECOND,
- ];
- /// Returns the variant's name or "" if unknown.
- pub fn variant_name(self) -> Option<&'static str> {
- match self {
- Self::SECOND => Some("SECOND"),
- Self::MILLISECOND => Some("MILLISECOND"),
- Self::MICROSECOND => Some("MICROSECOND"),
- Self::NANOSECOND => Some("NANOSECOND"),
- _ => None,
- }
- }
-}
-impl std::fmt::Debug for TimeUnit {
- fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
- if let Some(name) = self.variant_name() {
- f.write_str(name)
- } else {
- f.write_fmt(format_args!("<UNKNOWN {:?}>", self.0))
- }
- }
-}
-impl<'a> flatbuffers::Follow<'a> for TimeUnit {
- type Inner = Self;
- #[inline]
- fn follow(buf: &'a [u8], loc: usize) -> Self::Inner {
- let b = flatbuffers::read_scalar_at::<i16>(buf, loc);
- Self(b)
- }
-}
-
-impl flatbuffers::Push for TimeUnit {
- type Output = TimeUnit;
- #[inline]
- fn push(&self, dst: &mut [u8], _rest: &[u8]) {
- flatbuffers::emplace_scalar::<i16>(dst, self.0);
- }
-}
-
-impl flatbuffers::EndianScalar for TimeUnit {
- #[inline]
- fn to_little_endian(self) -> Self {
- let b = i16::to_le(self.0);
- Self(b)
- }
- #[inline]
- fn from_little_endian(self) -> Self {
- let b = i16::from_le(self.0);
- Self(b)
- }
-}
-
-impl<'a> flatbuffers::Verifiable for TimeUnit {
- #[inline]
- fn run_verifier(
- v: &mut flatbuffers::Verifier,
- pos: usize,
- ) -> Result<(), flatbuffers::InvalidFlatbuffer> {
- use flatbuffers::Verifiable;
- i16::run_verifier(v, pos)
- }
-}
-
-impl flatbuffers::SimpleToVerifyInSlice for TimeUnit {}
-#[deprecated(
- since = "2.0.0",
- note = "Use associated constants instead. This will no longer be generated in 2021."
-)]
-pub const ENUM_MIN_INTERVAL_UNIT: i16 = 0;
-#[deprecated(
- since = "2.0.0",
- note = "Use associated constants instead. This will no longer be generated in 2021."
-)]
-pub const ENUM_MAX_INTERVAL_UNIT: i16 = 1;
-#[deprecated(
- since = "2.0.0",
- note = "Use associated constants instead. This will no longer be generated in 2021."
-)]
-#[allow(non_camel_case_types)]
-pub const ENUM_VALUES_INTERVAL_UNIT: [IntervalUnit; 2] =
- [IntervalUnit::YEAR_MONTH, IntervalUnit::DAY_TIME];
-
-#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
-#[repr(transparent)]
-pub struct IntervalUnit(pub i16);
-#[allow(non_upper_case_globals)]
-impl IntervalUnit {
- pub const YEAR_MONTH: Self = Self(0);
- pub const DAY_TIME: Self = Self(1);
-
- pub const ENUM_MIN: i16 = 0;
- pub const ENUM_MAX: i16 = 1;
- pub const ENUM_VALUES: &'static [Self] = &[Self::YEAR_MONTH, Self::DAY_TIME];
- /// Returns the variant's name or "" if unknown.
- pub fn variant_name(self) -> Option<&'static str> {
- match self {
- Self::YEAR_MONTH => Some("YEAR_MONTH"),
- Self::DAY_TIME => Some("DAY_TIME"),
- _ => None,
- }
- }
-}
-impl std::fmt::Debug for IntervalUnit {
- fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
- if let Some(name) = self.variant_name() {
- f.write_str(name)
- } else {
- f.write_fmt(format_args!("<UNKNOWN {:?}>", self.0))
- }
- }
-}
-impl<'a> flatbuffers::Follow<'a> for IntervalUnit {
- type Inner = Self;
- #[inline]
- fn follow(buf: &'a [u8], loc: usize) -> Self::Inner {
- let b = flatbuffers::read_scalar_at::<i16>(buf, loc);
- Self(b)
- }
-}
-
-impl flatbuffers::Push for IntervalUnit {
- type Output = IntervalUnit;
- #[inline]
- fn push(&self, dst: &mut [u8], _rest: &[u8]) {
- flatbuffers::emplace_scalar::<i16>(dst, self.0);
- }
-}
-
-impl flatbuffers::EndianScalar for IntervalUnit {
- #[inline]
- fn to_little_endian(self) -> Self {
- let b = i16::to_le(self.0);
- Self(b)
- }
- #[inline]
- fn from_little_endian(self) -> Self {
- let b = i16::from_le(self.0);
- Self(b)
- }
-}
-
-impl<'a> flatbuffers::Verifiable for IntervalUnit {
- #[inline]
- fn run_verifier(
- v: &mut flatbuffers::Verifier,
- pos: usize,
- ) -> Result<(), flatbuffers::InvalidFlatbuffer> {
- use flatbuffers::Verifiable;
- i16::run_verifier(v, pos)
- }
-}
-
-impl flatbuffers::SimpleToVerifyInSlice for IntervalUnit {}
-#[deprecated(
- since = "2.0.0",
- note = "Use associated constants instead. This will no longer be generated in 2021."
-)]
-pub const ENUM_MIN_TYPE: u8 = 0;
-#[deprecated(
- since = "2.0.0",
- note = "Use associated constants instead. This will no longer be generated in 2021."
-)]
-pub const ENUM_MAX_TYPE: u8 = 21;
-#[deprecated(
- since = "2.0.0",
- note = "Use associated constants instead. This will no longer be generated in 2021."
-)]
-#[allow(non_camel_case_types)]
-pub const ENUM_VALUES_TYPE: [Type; 22] = [
- Type::NONE,
- Type::Null,
- Type::Int,
- Type::FloatingPoint,
- Type::Binary,
- Type::Utf8,
- Type::Bool,
- Type::Decimal,
- Type::Date,
- Type::Time,
- Type::Timestamp,
- Type::Interval,
- Type::List,
- Type::Struct_,
- Type::Union,
- Type::FixedSizeBinary,
- Type::FixedSizeList,
- Type::Map,
- Type::Duration,
- Type::LargeBinary,
- Type::LargeUtf8,
- Type::LargeList,
-];
-
-/// ----------------------------------------------------------------------
-/// Top-level Type value, enabling extensible type-specific metadata. We can
-/// add new logical types to Type without breaking backwards compatibility
-#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
-#[repr(transparent)]
-pub struct Type(pub u8);
-#[allow(non_upper_case_globals)]
-impl Type {
- pub const NONE: Self = Self(0);
- pub const Null: Self = Self(1);
- pub const Int: Self = Self(2);
- pub const FloatingPoint: Self = Self(3);
- pub const Binary: Self = Self(4);
- pub const Utf8: Self = Self(5);
- pub const Bool: Self = Self(6);
- pub const Decimal: Self = Self(7);
- pub const Date: Self = Self(8);
- pub const Time: Self = Self(9);
- pub const Timestamp: Self = Self(10);
- pub const Interval: Self = Self(11);
- pub const List: Self = Self(12);
- pub const Struct_: Self = Self(13);
- pub const Union: Self = Self(14);
- pub const FixedSizeBinary: Self = Self(15);
- pub const FixedSizeList: Self = Self(16);
- pub const Map: Self = Self(17);
- pub const Duration: Self = Self(18);
- pub const LargeBinary: Self = Self(19);
- pub const LargeUtf8: Self = Self(20);
- pub const LargeList: Self = Self(21);
-
- pub const ENUM_MIN: u8 = 0;
- pub const ENUM_MAX: u8 = 21;
- pub const ENUM_VALUES: &'static [Self] = &[
- Self::NONE,
- Self::Null,
- Self::Int,
- Self::FloatingPoint,
- Self::Binary,
- Self::Utf8,
- Self::Bool,
- Self::Decimal,
- Self::Date,
- Self::Time,
- Self::Timestamp,
- Self::Interval,
- Self::List,
- Self::Struct_,
- Self::Union,
- Self::FixedSizeBinary,
- Self::FixedSizeList,
- Self::Map,
- Self::Duration,
- Self::LargeBinary,
- Self::LargeUtf8,
- Self::LargeList,
- ];
- /// Returns the variant's name or "" if unknown.
- pub fn variant_name(self) -> Option<&'static str> {
- match self {
- Self::NONE => Some("NONE"),
- Self::Null => Some("Null"),
- Self::Int => Some("Int"),
- Self::FloatingPoint => Some("FloatingPoint"),
- Self::Binary => Some("Binary"),
- Self::Utf8 => Some("Utf8"),
- Self::Bool => Some("Bool"),
- Self::Decimal => Some("Decimal"),
- Self::Date => Some("Date"),
- Self::Time => Some("Time"),
- Self::Timestamp => Some("Timestamp"),
- Self::Interval => Some("Interval"),
- Self::List => Some("List"),
- Self::Struct_ => Some("Struct_"),
- Self::Union => Some("Union"),
- Self::FixedSizeBinary => Some("FixedSizeBinary"),
- Self::FixedSizeList => Some("FixedSizeList"),
- Self::Map => Some("Map"),
- Self::Duration => Some("Duration"),
- Self::LargeBinary => Some("LargeBinary"),
- Self::LargeUtf8 => Some("LargeUtf8"),
- Self::LargeList => Some("LargeList"),
- _ => None,
- }
- }
-}
-impl std::fmt::Debug for Type {
- fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
- if let Some(name) = self.variant_name() {
- f.write_str(name)
- } else {
- f.write_fmt(format_args!("<UNKNOWN {:?}>", self.0))
- }
- }
-}
-pub struct TypeUnionTableOffset {}
-impl<'a> flatbuffers::Follow<'a> for Type {
- type Inner = Self;
- #[inline]
- fn follow(buf: &'a [u8], loc: usize) -> Self::Inner {
- let b = flatbuffers::read_scalar_at::<u8>(buf, loc);
- Self(b)
- }
-}
-
-impl flatbuffers::Push for Type {
- type Output = Type;
- #[inline]
- fn push(&self, dst: &mut [u8], _rest: &[u8]) {
- flatbuffers::emplace_scalar::<u8>(dst, self.0);
- }
-}
-
-impl flatbuffers::EndianScalar for Type {
- #[inline]
- fn to_little_endian(self) -> Self {
- let b = u8::to_le(self.0);
- Self(b)
- }
- #[inline]
- fn from_little_endian(self) -> Self {
- let b = u8::from_le(self.0);
- Self(b)
- }
-}
-
-impl<'a> flatbuffers::Verifiable for Type {
- #[inline]
- fn run_verifier(
- v: &mut flatbuffers::Verifier,
- pos: usize,
- ) -> Result<(), flatbuffers::InvalidFlatbuffer> {
- use flatbuffers::Verifiable;
- u8::run_verifier(v, pos)
- }
-}
-
-impl flatbuffers::SimpleToVerifyInSlice for Type {}
-#[deprecated(
- since = "2.0.0",
- note = "Use associated constants instead. This will no longer be generated in 2021."
-)]
-pub const ENUM_MIN_DICTIONARY_KIND: i16 = 0;
-#[deprecated(
- since = "2.0.0",
- note = "Use associated constants instead. This will no longer be generated in 2021."
-)]
-pub const ENUM_MAX_DICTIONARY_KIND: i16 = 0;
-#[deprecated(
- since = "2.0.0",
- note = "Use associated constants instead. This will no longer be generated in 2021."
-)]
-#[allow(non_camel_case_types)]
-pub const ENUM_VALUES_DICTIONARY_KIND: [DictionaryKind; 1] = [DictionaryKind::DenseArray];
-
-/// ----------------------------------------------------------------------
-/// 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
-#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
-#[repr(transparent)]
-pub struct DictionaryKind(pub i16);
-#[allow(non_upper_case_globals)]
-impl DictionaryKind {
- pub const DenseArray: Self = Self(0);
-
- pub const ENUM_MIN: i16 = 0;
- pub const ENUM_MAX: i16 = 0;
- pub const ENUM_VALUES: &'static [Self] = &[Self::DenseArray];
- /// Returns the variant's name or "" if unknown.
- pub fn variant_name(self) -> Option<&'static str> {
- match self {
- Self::DenseArray => Some("DenseArray"),
- _ => None,
- }
- }
-}
-impl std::fmt::Debug for DictionaryKind {
- fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
- if let Some(name) = self.variant_name() {
- f.write_str(name)
- } else {
- f.write_fmt(format_args!("<UNKNOWN {:?}>", self.0))
- }
- }
-}
-impl<'a> flatbuffers::Follow<'a> for DictionaryKind {
- type Inner = Self;
- #[inline]
- fn follow(buf: &'a [u8], loc: usize) -> Self::Inner {
- let b = flatbuffers::read_scalar_at::<i16>(buf, loc);
- Self(b)
- }
-}
-
-impl flatbuffers::Push for DictionaryKind {
- type Output = DictionaryKind;
- #[inline]
- fn push(&self, dst: &mut [u8], _rest: &[u8]) {
- flatbuffers::emplace_scalar::<i16>(dst, self.0);
- }
-}
-
-impl flatbuffers::EndianScalar for DictionaryKind {
- #[inline]
- fn to_little_endian(self) -> Self {
- let b = i16::to_le(self.0);
- Self(b)
- }
- #[inline]
- fn from_little_endian(self) -> Self {
- let b = i16::from_le(self.0);
- Self(b)
- }
-}
-
-impl<'a> flatbuffers::Verifiable for DictionaryKind {
- #[inline]
- fn run_verifier(
- v: &mut flatbuffers::Verifier,
- pos: usize,
- ) -> Result<(), flatbuffers::InvalidFlatbuffer> {
- use flatbuffers::Verifiable;
- i16::run_verifier(v, pos)
- }
-}
-
-impl flatbuffers::SimpleToVerifyInSlice for DictionaryKind {}
-#[deprecated(
- since = "2.0.0",
- note = "Use associated constants instead. This will no longer be generated in 2021."
-)]
-pub const ENUM_MIN_ENDIANNESS: i16 = 0;
-#[deprecated(
- since = "2.0.0",
- note = "Use associated constants instead. This will no longer be generated in 2021."
-)]
-pub const ENUM_MAX_ENDIANNESS: i16 = 1;
-#[deprecated(
- since = "2.0.0",
- note = "Use associated constants instead. This will no longer be generated in 2021."
-)]
-#[allow(non_camel_case_types)]
-pub const ENUM_VALUES_ENDIANNESS: [Endianness; 2] = [Endianness::Little, Endianness::Big];
-
-/// ----------------------------------------------------------------------
-/// Endianness of the platform producing the data
-#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
-#[repr(transparent)]
-pub struct Endianness(pub i16);
-#[allow(non_upper_case_globals)]
-impl Endianness {
- pub const Little: Self = Self(0);
- pub const Big: Self = Self(1);
-
- pub const ENUM_MIN: i16 = 0;
- pub const ENUM_MAX: i16 = 1;
- pub const ENUM_VALUES: &'static [Self] = &[Self::Little, Self::Big];
- /// Returns the variant's name or "" if unknown.
- pub fn variant_name(self) -> Option<&'static str> {
- match self {
- Self::Little => Some("Little"),
- Self::Big => Some("Big"),
- _ => None,
- }
- }
-}
-impl std::fmt::Debug for Endianness {
- fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
- if let Some(name) = self.variant_name() {
- f.write_str(name)
- } else {
- f.write_fmt(format_args!("<UNKNOWN {:?}>", self.0))
- }
- }
-}
-impl<'a> flatbuffers::Follow<'a> for Endianness {
- type Inner = Self;
- #[inline]
- fn follow(buf: &'a [u8], loc: usize) -> Self::Inner {
- let b = flatbuffers::read_scalar_at::<i16>(buf, loc);
- Self(b)
- }
-}
-
-impl flatbuffers::Push for Endianness {
- type Output = Endianness;
- #[inline]
- fn push(&self, dst: &mut [u8], _rest: &[u8]) {
- flatbuffers::emplace_scalar::<i16>(dst, self.0);
- }
-}
-
-impl flatbuffers::EndianScalar for Endianness {
- #[inline]
- fn to_little_endian(self) -> Self {
- let b = i16::to_le(self.0);
- Self(b)
- }
- #[inline]
- fn from_little_endian(self) -> Self {
- let b = i16::from_le(self.0);
- Self(b)
- }
-}
-
-impl<'a> flatbuffers::Verifiable for Endianness {
- #[inline]
- fn run_verifier(
- v: &mut flatbuffers::Verifier,
- pos: usize,
- ) -> Result<(), flatbuffers::InvalidFlatbuffer> {
- use flatbuffers::Verifiable;
- i16::run_verifier(v, pos)
- }
-}
-
-impl flatbuffers::SimpleToVerifyInSlice for Endianness {}
-/// ----------------------------------------------------------------------
-/// A Buffer represents a single contiguous memory segment
-// struct Buffer, aligned to 8
-#[repr(transparent)]
-#[derive(Clone, Copy, PartialEq)]
-pub struct Buffer(pub [u8; 16]);
-impl std::fmt::Debug for Buffer {
- fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
- f.debug_struct("Buffer")
- .field("offset", &self.offset())
- .field("length", &self.length())
- .finish()
- }
-}
-
-impl flatbuffers::SimpleToVerifyInSlice for Buffer {}
-impl flatbuffers::SafeSliceAccess for Buffer {}
-impl<'a> flatbuffers::Follow<'a> for Buffer {
- type Inner = &'a Buffer;
- #[inline]
- fn follow(buf: &'a [u8], loc: usize) -> Self::Inner {
- <&'a Buffer>::follow(buf, loc)
- }
-}
-impl<'a> flatbuffers::Follow<'a> for &'a Buffer {
- type Inner = &'a Buffer;
- #[inline]
- fn follow(buf: &'a [u8], loc: usize) -> Self::Inner {
- flatbuffers::follow_cast_ref::<Buffer>(buf, loc)
- }
-}
-impl<'b> flatbuffers::Push for Buffer {
- type Output = Buffer;
- #[inline]
- fn push(&self, dst: &mut [u8], _rest: &[u8]) {
- let src = unsafe {
- ::std::slice::from_raw_parts(self as *const Buffer as *const u8, Self::size())
- };
- dst.copy_from_slice(src);
- }
-}
-impl<'b> flatbuffers::Push for &'b Buffer {
- type Output = Buffer;
-
- #[inline]
- fn push(&self, dst: &mut [u8], _rest: &[u8]) {
- let src = unsafe {
- ::std::slice::from_raw_parts(
- *self as *const Buffer as *const u8,
- Self::size(),
- )
- };
- dst.copy_from_slice(src);
- }
-}
-
-impl<'a> flatbuffers::Verifiable for Buffer {
- #[inline]
- fn run_verifier(
- v: &mut flatbuffers::Verifier,
- pos: usize,
- ) -> Result<(), flatbuffers::InvalidFlatbuffer> {
- use flatbuffers::Verifiable;
- v.in_buffer::<Self>(pos)
- }
-}
-impl Buffer {
- #[allow(clippy::too_many_arguments)]
- pub fn new(offset: i64, length: i64) -> Self {
- let mut s = Self([0; 16]);
- s.set_offset(offset);
- s.set_length(length);
- s
- }
-
- /// The relative offset into the shared memory page where the bytes for this
- /// buffer starts
- pub fn offset(&self) -> i64 {
- let mut mem = core::mem::MaybeUninit::<i64>::uninit();
- unsafe {
- core::ptr::copy_nonoverlapping(
- self.0[0..].as_ptr(),
- mem.as_mut_ptr() as *mut u8,
- core::mem::size_of::<i64>(),
- );
- mem.assume_init()
- }
- .from_little_endian()
- }
-
- pub fn set_offset(&mut self, x: i64) {
- let x_le = x.to_little_endian();
- unsafe {
- core::ptr::copy_nonoverlapping(
- &x_le as *const i64 as *const u8,
- self.0[0..].as_mut_ptr(),
- core::mem::size_of::<i64>(),
- );
- }
- }
-
- /// 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.
- pub fn length(&self) -> i64 {
- let mut mem = core::mem::MaybeUninit::<i64>::uninit();
- unsafe {
- core::ptr::copy_nonoverlapping(
- self.0[8..].as_ptr(),
- mem.as_mut_ptr() as *mut u8,
- core::mem::size_of::<i64>(),
- );
- mem.assume_init()
- }
- .from_little_endian()
- }
-
- pub fn set_length(&mut self, x: i64) {
- let x_le = x.to_little_endian();
- unsafe {
- core::ptr::copy_nonoverlapping(
- &x_le as *const i64 as *const u8,
- self.0[8..].as_mut_ptr(),
- core::mem::size_of::<i64>(),
- );
- }
- }
-}
-
-pub enum NullOffset {}
-#[derive(Copy, Clone, PartialEq)]
-
-/// These are stored in the flatbuffer in the Type union below
-pub struct Null<'a> {
- pub _tab: flatbuffers::Table<'a>,
-}
-
-impl<'a> flatbuffers::Follow<'a> for Null<'a> {
- type Inner = Null<'a>;
- #[inline]
- fn follow(buf: &'a [u8], loc: usize) -> Self::Inner {
- Self {
- _tab: flatbuffers::Table { buf, loc },
- }
- }
-}
-
-impl<'a> Null<'a> {
- #[inline]
- pub fn init_from_table(table: flatbuffers::Table<'a>) -> Self {
- Null { _tab: table }
- }
- #[allow(unused_mut)]
- pub fn create<'bldr: 'args, 'args: 'mut_bldr, 'mut_bldr>(
- _fbb: &'mut_bldr mut flatbuffers::FlatBufferBuilder<'bldr>,
- _args: &'args NullArgs,
- ) -> flatbuffers::WIPOffset<Null<'bldr>> {
- let mut builder = NullBuilder::new(_fbb);
- builder.finish()
- }
-}
-
-impl flatbuffers::Verifiable for Null<'_> {
- #[inline]
- fn run_verifier(
- v: &mut flatbuffers::Verifier,
- pos: usize,
- ) -> Result<(), flatbuffers::InvalidFlatbuffer> {
- use flatbuffers::Verifiable;
- v.visit_table(pos)?.finish();
- Ok(())
- }
-}
-pub struct NullArgs {}
-impl<'a> Default for NullArgs {
- #[inline]
- fn default() -> Self {
- NullArgs {}
- }
-}
-pub struct NullBuilder<'a: 'b, 'b> {
- fbb_: &'b mut flatbuffers::FlatBufferBuilder<'a>,
- start_: flatbuffers::WIPOffset<flatbuffers::TableUnfinishedWIPOffset>,
-}
-impl<'a: 'b, 'b> NullBuilder<'a, 'b> {
- #[inline]
- pub fn new(_fbb: &'b mut flatbuffers::FlatBufferBuilder<'a>) -> NullBuilder<'a, 'b> {
- let start = _fbb.start_table();
- NullBuilder {
- fbb_: _fbb,
- start_: start,
- }
- }
- #[inline]
- pub fn finish(self) -> flatbuffers::WIPOffset<Null<'a>> {
- let o = self.fbb_.end_table(self.start_);
- flatbuffers::WIPOffset::new(o.value())
- }
-}
-
-impl std::fmt::Debug for Null<'_> {
- fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
- let mut ds = f.debug_struct("Null");
- ds.finish()
- }
-}
-pub enum Struct_Offset {}
-#[derive(Copy, Clone, PartialEq)]
-
-/// 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
-pub struct Struct_<'a> {
- pub _tab: flatbuffers::Table<'a>,
-}
-
-impl<'a> flatbuffers::Follow<'a> for Struct_<'a> {
- type Inner = Struct_<'a>;
- #[inline]
- fn follow(buf: &'a [u8], loc: usize) -> Self::Inner {
- Self {
- _tab: flatbuffers::Table { buf, loc },
- }
- }
-}
-
-impl<'a> Struct_<'a> {
- #[inline]
- pub fn init_from_table(table: flatbuffers::Table<'a>) -> Self {
- Struct_ { _tab: table }
- }
- #[allow(unused_mut)]
- pub fn create<'bldr: 'args, 'args: 'mut_bldr, 'mut_bldr>(
- _fbb: &'mut_bldr mut flatbuffers::FlatBufferBuilder<'bldr>,
- _args: &'args Struct_Args,
- ) -> flatbuffers::WIPOffset<Struct_<'bldr>> {
- let mut builder = Struct_Builder::new(_fbb);
- builder.finish()
- }
-}
-
-impl flatbuffers::Verifiable for Struct_<'_> {
- #[inline]
- fn run_verifier(
- v: &mut flatbuffers::Verifier,
- pos: usize,
- ) -> Result<(), flatbuffers::InvalidFlatbuffer> {
- use flatbuffers::Verifiable;
- v.visit_table(pos)?.finish();
- Ok(())
- }
-}
-pub struct Struct_Args {}
-impl<'a> Default for Struct_Args {
- #[inline]
- fn default() -> Self {
- Struct_Args {}
- }
-}
-pub struct Struct_Builder<'a: 'b, 'b> {
- fbb_: &'b mut flatbuffers::FlatBufferBuilder<'a>,
- start_: flatbuffers::WIPOffset<flatbuffers::TableUnfinishedWIPOffset>,
-}
-impl<'a: 'b, 'b> Struct_Builder<'a, 'b> {
- #[inline]
- pub fn new(
- _fbb: &'b mut flatbuffers::FlatBufferBuilder<'a>,
- ) -> Struct_Builder<'a, 'b> {
- let start = _fbb.start_table();
- Struct_Builder {
- fbb_: _fbb,
- start_: start,
- }
- }
- #[inline]
- pub fn finish(self) -> flatbuffers::WIPOffset<Struct_<'a>> {
- let o = self.fbb_.end_table(self.start_);
- flatbuffers::WIPOffset::new(o.value())
- }
-}
-
-impl std::fmt::Debug for Struct_<'_> {
- fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
- let mut ds = f.debug_struct("Struct_");
- ds.finish()
- }
-}
-pub enum ListOffset {}
-#[derive(Copy, Clone, PartialEq)]
-
-pub struct List<'a> {
- pub _tab: flatbuffers::Table<'a>,
-}
-
-impl<'a> flatbuffers::Follow<'a> for List<'a> {
- type Inner = List<'a>;
- #[inline]
- fn follow(buf: &'a [u8], loc: usize) -> Self::Inner {
- Self {
- _tab: flatbuffers::Table { buf, loc },
- }
- }
-}
-
-impl<'a> List<'a> {
- #[inline]
- pub fn init_from_table(table: flatbuffers::Table<'a>) -> Self {
- List { _tab: table }
- }
- #[allow(unused_mut)]
- pub fn create<'bldr: 'args, 'args: 'mut_bldr, 'mut_bldr>(
- _fbb: &'mut_bldr mut flatbuffers::FlatBufferBuilder<'bldr>,
- _args: &'args ListArgs,
- ) -> flatbuffers::WIPOffset<List<'bldr>> {
- let mut builder = ListBuilder::new(_fbb);
- builder.finish()
- }
-}
-
-impl flatbuffers::Verifiable for List<'_> {
- #[inline]
- fn run_verifier(
- v: &mut flatbuffers::Verifier,
- pos: usize,
- ) -> Result<(), flatbuffers::InvalidFlatbuffer> {
- use flatbuffers::Verifiable;
- v.visit_table(pos)?.finish();
- Ok(())
- }
-}
-pub struct ListArgs {}
-impl<'a> Default for ListArgs {
- #[inline]
- fn default() -> Self {
- ListArgs {}
- }
-}
-pub struct ListBuilder<'a: 'b, 'b> {
- fbb_: &'b mut flatbuffers::FlatBufferBuilder<'a>,
- start_: flatbuffers::WIPOffset<flatbuffers::TableUnfinishedWIPOffset>,
-}
-impl<'a: 'b, 'b> ListBuilder<'a, 'b> {
- #[inline]
- pub fn new(_fbb: &'b mut flatbuffers::FlatBufferBuilder<'a>) -> ListBuilder<'a, 'b> {
- let start = _fbb.start_table();
- ListBuilder {
- fbb_: _fbb,
- start_: start,
- }
- }
- #[inline]
- pub fn finish(self) -> flatbuffers::WIPOffset<List<'a>> {
- let o = self.fbb_.end_table(self.start_);
- flatbuffers::WIPOffset::new(o.value())
- }
-}
-
-impl std::fmt::Debug for List<'_> {
- fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
- let mut ds = f.debug_struct("List");
- ds.finish()
- }
-}
-pub enum LargeListOffset {}
-#[derive(Copy, Clone, PartialEq)]
-
-/// Same as List, but with 64-bit offsets, allowing to represent
-/// extremely large data values.
-pub struct LargeList<'a> {
- pub _tab: flatbuffers::Table<'a>,
-}
-
-impl<'a> flatbuffers::Follow<'a> for LargeList<'a> {
- type Inner = LargeList<'a>;
- #[inline]
- fn follow(buf: &'a [u8], loc: usize) -> Self::Inner {
- Self {
- _tab: flatbuffers::Table { buf, loc },
- }
- }
-}
-
-impl<'a> LargeList<'a> {
- #[inline]
- pub fn init_from_table(table: flatbuffers::Table<'a>) -> Self {
- LargeList { _tab: table }
- }
- #[allow(unused_mut)]
- pub fn create<'bldr: 'args, 'args: 'mut_bldr, 'mut_bldr>(
- _fbb: &'mut_bldr mut flatbuffers::FlatBufferBuilder<'bldr>,
- _args: &'args LargeListArgs,
- ) -> flatbuffers::WIPOffset<LargeList<'bldr>> {
- let mut builder = LargeListBuilder::new(_fbb);
- builder.finish()
- }
-}
-
-impl flatbuffers::Verifiable for LargeList<'_> {
- #[inline]
- fn run_verifier(
- v: &mut flatbuffers::Verifier,
- pos: usize,
- ) -> Result<(), flatbuffers::InvalidFlatbuffer> {
- use flatbuffers::Verifiable;
- v.visit_table(pos)?.finish();
- Ok(())
- }
-}
-pub struct LargeListArgs {}
-impl<'a> Default for LargeListArgs {
- #[inline]
- fn default() -> Self {
- LargeListArgs {}
- }
-}
-pub struct LargeListBuilder<'a: 'b, 'b> {
- fbb_: &'b mut flatbuffers::FlatBufferBuilder<'a>,
- start_: flatbuffers::WIPOffset<flatbuffers::TableUnfinishedWIPOffset>,
-}
-impl<'a: 'b, 'b> LargeListBuilder<'a, 'b> {
- #[inline]
- pub fn new(
- _fbb: &'b mut flatbuffers::FlatBufferBuilder<'a>,
- ) -> LargeListBuilder<'a, 'b> {
- let start = _fbb.start_table();
- LargeListBuilder {
- fbb_: _fbb,
- start_: start,
- }
- }
- #[inline]
- pub fn finish(self) -> flatbuffers::WIPOffset<LargeList<'a>> {
- let o = self.fbb_.end_table(self.start_);
- flatbuffers::WIPOffset::new(o.value())
- }
-}
-
-impl std::fmt::Debug for LargeList<'_> {
- fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
- let mut ds = f.debug_struct("LargeList");
- ds.finish()
- }
-}
-pub enum FixedSizeListOffset {}
-#[derive(Copy, Clone, PartialEq)]
-
-pub struct FixedSizeList<'a> {
- pub _tab: flatbuffers::Table<'a>,
-}
-
-impl<'a> flatbuffers::Follow<'a> for FixedSizeList<'a> {
- type Inner = FixedSizeList<'a>;
- #[inline]
- fn follow(buf: &'a [u8], loc: usize) -> Self::Inner {
- Self {
- _tab: flatbuffers::Table { buf, loc },
- }
- }
-}
-
-impl<'a> FixedSizeList<'a> {
- #[inline]
- pub fn init_from_table(table: flatbuffers::Table<'a>) -> Self {
- FixedSizeList { _tab: table }
- }
- #[allow(unused_mut)]
- pub fn create<'bldr: 'args, 'args: 'mut_bldr, 'mut_bldr>(
- _fbb: &'mut_bldr mut flatbuffers::FlatBufferBuilder<'bldr>,
- args: &'args FixedSizeListArgs,
- ) -> flatbuffers::WIPOffset<FixedSizeList<'bldr>> {
- let mut builder = FixedSizeListBuilder::new(_fbb);
- builder.add_listSize(args.listSize);
- builder.finish()
- }
-
- pub const VT_LISTSIZE: flatbuffers::VOffsetT = 4;
-
- /// Number of list items per value
- #[inline]
- pub fn listSize(&self) -> i32 {
- self._tab
- .get::<i32>(FixedSizeList::VT_LISTSIZE, Some(0))
- .unwrap()
- }
-}
-
-impl flatbuffers::Verifiable for FixedSizeList<'_> {
- #[inline]
- fn run_verifier(
- v: &mut flatbuffers::Verifier,
- pos: usize,
- ) -> Result<(), flatbuffers::InvalidFlatbuffer> {
- use flatbuffers::Verifiable;
- v.visit_table(pos)?
- .visit_field::<i32>(&"listSize", Self::VT_LISTSIZE, false)?
- .finish();
- Ok(())
- }
-}
-pub struct FixedSizeListArgs {
- pub listSize: i32,
-}
-impl<'a> Default for FixedSizeListArgs {
- #[inline]
- fn default() -> Self {
- FixedSizeListArgs { listSize: 0 }
- }
-}
-pub struct FixedSizeListBuilder<'a: 'b, 'b> {
- fbb_: &'b mut flatbuffers::FlatBufferBuilder<'a>,
- start_: flatbuffers::WIPOffset<flatbuffers::TableUnfinishedWIPOffset>,
-}
-impl<'a: 'b, 'b> FixedSizeListBuilder<'a, 'b> {
- #[inline]
- pub fn add_listSize(&mut self, listSize: i32) {
- self.fbb_
- .push_slot::<i32>(FixedSizeList::VT_LISTSIZE, listSize, 0);
- }
- #[inline]
- pub fn new(
- _fbb: &'b mut flatbuffers::FlatBufferBuilder<'a>,
- ) -> FixedSizeListBuilder<'a, 'b> {
- let start = _fbb.start_table();
- FixedSizeListBuilder {
- fbb_: _fbb,
- start_: start,
- }
- }
- #[inline]
- pub fn finish(self) -> flatbuffers::WIPOffset<FixedSizeList<'a>> {
- let o = self.fbb_.end_table(self.start_);
- flatbuffers::WIPOffset::new(o.value())
- }
-}
-
-impl std::fmt::Debug for FixedSizeList<'_> {
- fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
- let mut ds = f.debug_struct("FixedSizeList");
- ds.field("listSize", &self.listSize());
- ds.finish()
- }
-}
-pub enum MapOffset {}
-#[derive(Copy, Clone, PartialEq)]
-
-/// A Map is a logical nested type that is represented as
-///
-/// List<entries: Struct<key: K, value: V>>
-///
-/// 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.
-pub struct Map<'a> {
- pub _tab: flatbuffers::Table<'a>,
-}
-
-impl<'a> flatbuffers::Follow<'a> for Map<'a> {
- type Inner = Map<'a>;
- #[inline]
- fn follow(buf: &'a [u8], loc: usize) -> Self::Inner {
- Self {
- _tab: flatbuffers::Table { buf, loc },
- }
- }
-}
-
-impl<'a> Map<'a> {
- #[inline]
- pub fn init_from_table(table: flatbuffers::Table<'a>) -> Self {
- Map { _tab: table }
- }
- #[allow(unused_mut)]
- pub fn create<'bldr: 'args, 'args: 'mut_bldr, 'mut_bldr>(
- _fbb: &'mut_bldr mut flatbuffers::FlatBufferBuilder<'bldr>,
- args: &'args MapArgs,
- ) -> flatbuffers::WIPOffset<Map<'bldr>> {
- let mut builder = MapBuilder::new(_fbb);
- builder.add_keysSorted(args.keysSorted);
- builder.finish()
- }
-
- pub const VT_KEYSSORTED: flatbuffers::VOffsetT = 4;
-
- /// Set to true if the keys within each value are sorted
- #[inline]
- pub fn keysSorted(&self) -> bool {
- self._tab
- .get::<bool>(Map::VT_KEYSSORTED, Some(false))
- .unwrap()
- }
-}
-
-impl flatbuffers::Verifiable for Map<'_> {
- #[inline]
- fn run_verifier(
- v: &mut flatbuffers::Verifier,
- pos: usize,
- ) -> Result<(), flatbuffers::InvalidFlatbuffer> {
- use flatbuffers::Verifiable;
- v.visit_table(pos)?
- .visit_field::<bool>(&"keysSorted", Self::VT_KEYSSORTED, false)?
- .finish();
- Ok(())
- }
-}
-pub struct MapArgs {
- pub keysSorted: bool,
-}
-impl<'a> Default for MapArgs {
- #[inline]
- fn default() -> Self {
- MapArgs { keysSorted: false }
- }
-}
-pub struct MapBuilder<'a: 'b, 'b> {
- fbb_: &'b mut flatbuffers::FlatBufferBuilder<'a>,
- start_: flatbuffers::WIPOffset<flatbuffers::TableUnfinishedWIPOffset>,
-}
-impl<'a: 'b, 'b> MapBuilder<'a, 'b> {
- #[inline]
- pub fn add_keysSorted(&mut self, keysSorted: bool) {
- self.fbb_
- .push_slot::<bool>(Map::VT_KEYSSORTED, keysSorted, false);
- }
- #[inline]
- pub fn new(_fbb: &'b mut flatbuffers::FlatBufferBuilder<'a>) -> MapBuilder<'a, 'b> {
- let start = _fbb.start_table();
- MapBuilder {
- fbb_: _fbb,
- start_: start,
- }
- }
- #[inline]
- pub fn finish(self) -> flatbuffers::WIPOffset<Map<'a>> {
- let o = self.fbb_.end_table(self.start_);
- flatbuffers::WIPOffset::new(o.value())
- }
-}
-
-impl std::fmt::Debug for Map<'_> {
- fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
- let mut ds = f.debug_struct("Map");
- ds.field("keysSorted", &self.keysSorted());
- ds.finish()
- }
-}
-pub enum UnionOffset {}
-#[derive(Copy, Clone, PartialEq)]
-
-/// 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
-pub struct Union<'a> {
- pub _tab: flatbuffers::Table<'a>,
-}
-
-impl<'a> flatbuffers::Follow<'a> for Union<'a> {
- type Inner = Union<'a>;
- #[inline]
- fn follow(buf: &'a [u8], loc: usize) -> Self::Inner {
- Self {
- _tab: flatbuffers::Table { buf, loc },
- }
- }
-}
-
-impl<'a> Union<'a> {
- #[inline]
- pub fn init_from_table(table: flatbuffers::Table<'a>) -> Self {
- Union { _tab: table }
- }
- #[allow(unused_mut)]
- pub fn create<'bldr: 'args, 'args: 'mut_bldr, 'mut_bldr>(
- _fbb: &'mut_bldr mut flatbuffers::FlatBufferBuilder<'bldr>,
- args: &'args UnionArgs<'args>,
- ) -> flatbuffers::WIPOffset<Union<'bldr>> {
- let mut builder = UnionBuilder::new(_fbb);
- if let Some(x) = args.typeIds {
- builder.add_typeIds(x);
- }
- builder.add_mode(args.mode);
- builder.finish()
- }
-
- pub const VT_MODE: flatbuffers::VOffsetT = 4;
- pub const VT_TYPEIDS: flatbuffers::VOffsetT = 6;
-
- #[inline]
- pub fn mode(&self) -> UnionMode {
- self._tab
- .get::<UnionMode>(Union::VT_MODE, Some(UnionMode::Sparse))
- .unwrap()
- }
- #[inline]
- pub fn typeIds(&self) -> Option<flatbuffers::Vector<'a, i32>> {
- self._tab
- .get::<flatbuffers::ForwardsUOffset<flatbuffers::Vector<'a, i32>>>(
- Union::VT_TYPEIDS,
- None,
- )
- }
-}
-
-impl flatbuffers::Verifiable for Union<'_> {
- #[inline]
- fn run_verifier(
- v: &mut flatbuffers::Verifier,
- pos: usize,
- ) -> Result<(), flatbuffers::InvalidFlatbuffer> {
- use flatbuffers::Verifiable;
- v.visit_table(pos)?
- .visit_field::<UnionMode>(&"mode", Self::VT_MODE, false)?
- .visit_field::<flatbuffers::ForwardsUOffset<flatbuffers::Vector<'_, i32>>>(
- &"typeIds",
- Self::VT_TYPEIDS,
- false,
- )?
- .finish();
- Ok(())
- }
-}
-pub struct UnionArgs<'a> {
- pub mode: UnionMode,
- pub typeIds: Option<flatbuffers::WIPOffset<flatbuffers::Vector<'a, i32>>>,
-}
-impl<'a> Default for UnionArgs<'a> {
- #[inline]
- fn default() -> Self {
- UnionArgs {
- mode: UnionMode::Sparse,
- typeIds: None,
- }
- }
-}
-pub struct UnionBuilder<'a: 'b, 'b> {
- fbb_: &'b mut flatbuffers::FlatBufferBuilder<'a>,
- start_: flatbuffers::WIPOffset<flatbuffers::TableUnfinishedWIPOffset>,
-}
-impl<'a: 'b, 'b> UnionBuilder<'a, 'b> {
- #[inline]
- pub fn add_mode(&mut self, mode: UnionMode) {
- self.fbb_
- .push_slot::<UnionMode>(Union::VT_MODE, mode, UnionMode::Sparse);
- }
- #[inline]
- pub fn add_typeIds(
- &mut self,
- typeIds: flatbuffers::WIPOffset<flatbuffers::Vector<'b, i32>>,
- ) {
- self.fbb_
- .push_slot_always::<flatbuffers::WIPOffset<_>>(Union::VT_TYPEIDS, typeIds);
- }
- #[inline]
- pub fn new(_fbb: &'b mut flatbuffers::FlatBufferBuilder<'a>) -> UnionBuilder<'a, 'b> {
- let start = _fbb.start_table();
- UnionBuilder {
- fbb_: _fbb,
- start_: start,
- }
- }
- #[inline]
- pub fn finish(self) -> flatbuffers::WIPOffset<Union<'a>> {
- let o = self.fbb_.end_table(self.start_);
- flatbuffers::WIPOffset::new(o.value())
- }
-}
-
-impl std::fmt::Debug for Union<'_> {
- fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
- let mut ds = f.debug_struct("Union");
- ds.field("mode", &self.mode());
- ds.field("typeIds", &self.typeIds());
- ds.finish()
- }
-}
-pub enum IntOffset {}
-#[derive(Copy, Clone, PartialEq)]
-
-pub struct Int<'a> {
- pub _tab: flatbuffers::Table<'a>,
-}
-
-impl<'a> flatbuffers::Follow<'a> for Int<'a> {
- type Inner = Int<'a>;
- #[inline]
- fn follow(buf: &'a [u8], loc: usize) -> Self::Inner {
- Self {
- _tab: flatbuffers::Table { buf, loc },
- }
- }
-}
-
-impl<'a> Int<'a> {
- #[inline]
- pub fn init_from_table(table: flatbuffers::Table<'a>) -> Self {
- Int { _tab: table }
- }
- #[allow(unused_mut)]
- pub fn create<'bldr: 'args, 'args: 'mut_bldr, 'mut_bldr>(
- _fbb: &'mut_bldr mut flatbuffers::FlatBufferBuilder<'bldr>,
- args: &'args IntArgs,
- ) -> flatbuffers::WIPOffset<Int<'bldr>> {
- let mut builder = IntBuilder::new(_fbb);
- builder.add_bitWidth(args.bitWidth);
- builder.add_is_signed(args.is_signed);
- builder.finish()
- }
-
- pub const VT_BITWIDTH: flatbuffers::VOffsetT = 4;
- pub const VT_IS_SIGNED: flatbuffers::VOffsetT = 6;
-
- #[inline]
- pub fn bitWidth(&self) -> i32 {
- self._tab.get::<i32>(Int::VT_BITWIDTH, Some(0)).unwrap()
- }
- #[inline]
- pub fn is_signed(&self) -> bool {
- self._tab
- .get::<bool>(Int::VT_IS_SIGNED, Some(false))
- .unwrap()
- }
-}
-
-impl flatbuffers::Verifiable for Int<'_> {
- #[inline]
- fn run_verifier(
- v: &mut flatbuffers::Verifier,
- pos: usize,
- ) -> Result<(), flatbuffers::InvalidFlatbuffer> {
- use flatbuffers::Verifiable;
- v.visit_table(pos)?
- .visit_field::<i32>(&"bitWidth", Self::VT_BITWIDTH, false)?
- .visit_field::<bool>(&"is_signed", Self::VT_IS_SIGNED, false)?
- .finish();
- Ok(())
- }
-}
-pub struct IntArgs {
- pub bitWidth: i32,
- pub is_signed: bool,
-}
-impl<'a> Default for IntArgs {
- #[inline]
- fn default() -> Self {
- IntArgs {
- bitWidth: 0,
- is_signed: false,
- }
- }
-}
-pub struct IntBuilder<'a: 'b, 'b> {
- fbb_: &'b mut flatbuffers::FlatBufferBuilder<'a>,
- start_: flatbuffers::WIPOffset<flatbuffers::TableUnfinishedWIPOffset>,
-}
-impl<'a: 'b, 'b> IntBuilder<'a, 'b> {
- #[inline]
- pub fn add_bitWidth(&mut self, bitWidth: i32) {
- self.fbb_.push_slot::<i32>(Int::VT_BITWIDTH, bitWidth, 0);
- }
- #[inline]
- pub fn add_is_signed(&mut self, is_signed: bool) {
- self.fbb_
- .push_slot::<bool>(Int::VT_IS_SIGNED, is_signed, false);
- }
- #[inline]
- pub fn new(_fbb: &'b mut flatbuffers::FlatBufferBuilder<'a>) -> IntBuilder<'a, 'b> {
- let start = _fbb.start_table();
- IntBuilder {
- fbb_: _fbb,
- start_: start,
- }
- }
- #[inline]
- pub fn finish(self) -> flatbuffers::WIPOffset<Int<'a>> {
- let o = self.fbb_.end_table(self.start_);
- flatbuffers::WIPOffset::new(o.value())
- }
-}
-
-impl std::fmt::Debug for Int<'_> {
- fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
- let mut ds = f.debug_struct("Int");
- ds.field("bitWidth", &self.bitWidth());
- ds.field("is_signed", &self.is_signed());
- ds.finish()
- }
-}
-pub enum FloatingPointOffset {}
-#[derive(Copy, Clone, PartialEq)]
-
-pub struct FloatingPoint<'a> {
- pub _tab: flatbuffers::Table<'a>,
-}
-
-impl<'a> flatbuffers::Follow<'a> for FloatingPoint<'a> {
- type Inner = FloatingPoint<'a>;
- #[inline]
- fn follow(buf: &'a [u8], loc: usize) -> Self::Inner {
- Self {
- _tab: flatbuffers::Table { buf, loc },
- }
- }
-}
-
-impl<'a> FloatingPoint<'a> {
- #[inline]
- pub fn init_from_table(table: flatbuffers::Table<'a>) -> Self {
- FloatingPoint { _tab: table }
- }
- #[allow(unused_mut)]
- pub fn create<'bldr: 'args, 'args: 'mut_bldr, 'mut_bldr>(
- _fbb: &'mut_bldr mut flatbuffers::FlatBufferBuilder<'bldr>,
- args: &'args FloatingPointArgs,
- ) -> flatbuffers::WIPOffset<FloatingPoint<'bldr>> {
- let mut builder = FloatingPointBuilder::new(_fbb);
- builder.add_precision(args.precision);
- builder.finish()
- }
-
- pub const VT_PRECISION: flatbuffers::VOffsetT = 4;
-
- #[inline]
- pub fn precision(&self) -> Precision {
- self._tab
- .get::<Precision>(FloatingPoint::VT_PRECISION, Some(Precision::HALF))
- .unwrap()
- }
-}
-
-impl flatbuffers::Verifiable for FloatingPoint<'_> {
- #[inline]
- fn run_verifier(
- v: &mut flatbuffers::Verifier,
- pos: usize,
- ) -> Result<(), flatbuffers::InvalidFlatbuffer> {
- use flatbuffers::Verifiable;
- v.visit_table(pos)?
- .visit_field::<Precision>(&"precision", Self::VT_PRECISION, false)?
- .finish();
- Ok(())
- }
-}
-pub struct FloatingPointArgs {
- pub precision: Precision,
-}
-impl<'a> Default for FloatingPointArgs {
- #[inline]
- fn default() -> Self {
- FloatingPointArgs {
- precision: Precision::HALF,
- }
- }
-}
-pub struct FloatingPointBuilder<'a: 'b, 'b> {
- fbb_: &'b mut flatbuffers::FlatBufferBuilder<'a>,
- start_: flatbuffers::WIPOffset<flatbuffers::TableUnfinishedWIPOffset>,
-}
-impl<'a: 'b, 'b> FloatingPointBuilder<'a, 'b> {
- #[inline]
- pub fn add_precision(&mut self, precision: Precision) {
- self.fbb_.push_slot::<Precision>(
- FloatingPoint::VT_PRECISION,
- precision,
- Precision::HALF,
- );
- }
- #[inline]
- pub fn new(
- _fbb: &'b mut flatbuffers::FlatBufferBuilder<'a>,
- ) -> FloatingPointBuilder<'a, 'b> {
- let start = _fbb.start_table();
- FloatingPointBuilder {
- fbb_: _fbb,
- start_: start,
- }
- }
- #[inline]
- pub fn finish(self) -> flatbuffers::WIPOffset<FloatingPoint<'a>> {
- let o = self.fbb_.end_table(self.start_);
- flatbuffers::WIPOffset::new(o.value())
- }
-}
-
-impl std::fmt::Debug for FloatingPoint<'_> {
- fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
- let mut ds = f.debug_struct("FloatingPoint");
- ds.field("precision", &self.precision());
- ds.finish()
- }
-}
-pub enum Utf8Offset {}
-#[derive(Copy, Clone, PartialEq)]
-
-/// Unicode with UTF-8 encoding
-pub struct Utf8<'a> {
- pub _tab: flatbuffers::Table<'a>,
-}
-
-impl<'a> flatbuffers::Follow<'a> for Utf8<'a> {
- type Inner = Utf8<'a>;
- #[inline]
- fn follow(buf: &'a [u8], loc: usize) -> Self::Inner {
- Self {
- _tab: flatbuffers::Table { buf, loc },
- }
- }
-}
-
-impl<'a> Utf8<'a> {
- #[inline]
- pub fn init_from_table(table: flatbuffers::Table<'a>) -> Self {
- Utf8 { _tab: table }
- }
- #[allow(unused_mut)]
- pub fn create<'bldr: 'args, 'args: 'mut_bldr, 'mut_bldr>(
- _fbb: &'mut_bldr mut flatbuffers::FlatBufferBuilder<'bldr>,
- _args: &'args Utf8Args,
- ) -> flatbuffers::WIPOffset<Utf8<'bldr>> {
- let mut builder = Utf8Builder::new(_fbb);
- builder.finish()
- }
-}
-
-impl flatbuffers::Verifiable for Utf8<'_> {
- #[inline]
- fn run_verifier(
- v: &mut flatbuffers::Verifier,
- pos: usize,
- ) -> Result<(), flatbuffers::InvalidFlatbuffer> {
- use flatbuffers::Verifiable;
- v.visit_table(pos)?.finish();
- Ok(())
- }
-}
-pub struct Utf8Args {}
-impl<'a> Default for Utf8Args {
- #[inline]
- fn default() -> Self {
- Utf8Args {}
- }
-}
-pub struct Utf8Builder<'a: 'b, 'b> {
- fbb_: &'b mut flatbuffers::FlatBufferBuilder<'a>,
- start_: flatbuffers::WIPOffset<flatbuffers::TableUnfinishedWIPOffset>,
-}
-impl<'a: 'b, 'b> Utf8Builder<'a, 'b> {
- #[inline]
- pub fn new(_fbb: &'b mut flatbuffers::FlatBufferBuilder<'a>) -> Utf8Builder<'a, 'b> {
- let start = _fbb.start_table();
- Utf8Builder {
- fbb_: _fbb,
- start_: start,
- }
- }
- #[inline]
- pub fn finish(self) -> flatbuffers::WIPOffset<Utf8<'a>> {
- let o = self.fbb_.end_table(self.start_);
- flatbuffers::WIPOffset::new(o.value())
- }
-}
-
-impl std::fmt::Debug for Utf8<'_> {
- fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
- let mut ds = f.debug_struct("Utf8");
- ds.finish()
- }
-}
-pub enum BinaryOffset {}
-#[derive(Copy, Clone, PartialEq)]
-
-/// Opaque binary data
-pub struct Binary<'a> {
- pub _tab: flatbuffers::Table<'a>,
-}
-
-impl<'a> flatbuffers::Follow<'a> for Binary<'a> {
- type Inner = Binary<'a>;
- #[inline]
- fn follow(buf: &'a [u8], loc: usize) -> Self::Inner {
- Self {
- _tab: flatbuffers::Table { buf, loc },
- }
- }
-}
-
-impl<'a> Binary<'a> {
- #[inline]
- pub fn init_from_table(table: flatbuffers::Table<'a>) -> Self {
- Binary { _tab: table }
- }
- #[allow(unused_mut)]
- pub fn create<'bldr: 'args, 'args: 'mut_bldr, 'mut_bldr>(
- _fbb: &'mut_bldr mut flatbuffers::FlatBufferBuilder<'bldr>,
- _args: &'args BinaryArgs,
- ) -> flatbuffers::WIPOffset<Binary<'bldr>> {
- let mut builder = BinaryBuilder::new(_fbb);
- builder.finish()
- }
-}
-
-impl flatbuffers::Verifiable for Binary<'_> {
- #[inline]
- fn run_verifier(
- v: &mut flatbuffers::Verifier,
- pos: usize,
- ) -> Result<(), flatbuffers::InvalidFlatbuffer> {
- use flatbuffers::Verifiable;
- v.visit_table(pos)?.finish();
- Ok(())
- }
-}
-pub struct BinaryArgs {}
-impl<'a> Default for BinaryArgs {
- #[inline]
- fn default() -> Self {
- BinaryArgs {}
- }
-}
-pub struct BinaryBuilder<'a: 'b, 'b> {
- fbb_: &'b mut flatbuffers::FlatBufferBuilder<'a>,
- start_: flatbuffers::WIPOffset<flatbuffers::TableUnfinishedWIPOffset>,
-}
-impl<'a: 'b, 'b> BinaryBuilder<'a, 'b> {
- #[inline]
- pub fn new(
- _fbb: &'b mut flatbuffers::FlatBufferBuilder<'a>,
- ) -> BinaryBuilder<'a, 'b> {
- let start = _fbb.start_table();
- BinaryBuilder {
- fbb_: _fbb,
- start_: start,
- }
- }
- #[inline]
- pub fn finish(self) -> flatbuffers::WIPOffset<Binary<'a>> {
- let o = self.fbb_.end_table(self.start_);
- flatbuffers::WIPOffset::new(o.value())
- }
-}
-
-impl std::fmt::Debug for Binary<'_> {
- fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
- let mut ds = f.debug_struct("Binary");
- ds.finish()
- }
-}
-pub enum LargeUtf8Offset {}
-#[derive(Copy, Clone, PartialEq)]
-
-/// Same as Utf8, but with 64-bit offsets, allowing to represent
-/// extremely large data values.
-pub struct LargeUtf8<'a> {
- pub _tab: flatbuffers::Table<'a>,
-}
-
-impl<'a> flatbuffers::Follow<'a> for LargeUtf8<'a> {
- type Inner = LargeUtf8<'a>;
- #[inline]
- fn follow(buf: &'a [u8], loc: usize) -> Self::Inner {
- Self {
- _tab: flatbuffers::Table { buf, loc },
- }
- }
-}
-
-impl<'a> LargeUtf8<'a> {
- #[inline]
- pub fn init_from_table(table: flatbuffers::Table<'a>) -> Self {
- LargeUtf8 { _tab: table }
- }
- #[allow(unused_mut)]
- pub fn create<'bldr: 'args, 'args: 'mut_bldr, 'mut_bldr>(
- _fbb: &'mut_bldr mut flatbuffers::FlatBufferBuilder<'bldr>,
- _args: &'args LargeUtf8Args,
- ) -> flatbuffers::WIPOffset<LargeUtf8<'bldr>> {
- let mut builder = LargeUtf8Builder::new(_fbb);
- builder.finish()
- }
-}
-
-impl flatbuffers::Verifiable for LargeUtf8<'_> {
- #[inline]
- fn run_verifier(
- v: &mut flatbuffers::Verifier,
- pos: usize,
- ) -> Result<(), flatbuffers::InvalidFlatbuffer> {
- use flatbuffers::Verifiable;
- v.visit_table(pos)?.finish();
- Ok(())
- }
-}
-pub struct LargeUtf8Args {}
-impl<'a> Default for LargeUtf8Args {
- #[inline]
- fn default() -> Self {
- LargeUtf8Args {}
- }
-}
-pub struct LargeUtf8Builder<'a: 'b, 'b> {
- fbb_: &'b mut flatbuffers::FlatBufferBuilder<'a>,
- start_: flatbuffers::WIPOffset<flatbuffers::TableUnfinishedWIPOffset>,
-}
-impl<'a: 'b, 'b> LargeUtf8Builder<'a, 'b> {
- #[inline]
- pub fn new(
- _fbb: &'b mut flatbuffers::FlatBufferBuilder<'a>,
- ) -> LargeUtf8Builder<'a, 'b> {
- let start = _fbb.start_table();
- LargeUtf8Builder {
- fbb_: _fbb,
- start_: start,
- }
- }
- #[inline]
- pub fn finish(self) -> flatbuffers::WIPOffset<LargeUtf8<'a>> {
- let o = self.fbb_.end_table(self.start_);
- flatbuffers::WIPOffset::new(o.value())
- }
-}
-
-impl std::fmt::Debug for LargeUtf8<'_> {
- fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
- let mut ds = f.debug_struct("LargeUtf8");
- ds.finish()
- }
-}
-pub enum LargeBinaryOffset {}
-#[derive(Copy, Clone, PartialEq)]
-
-/// Same as Binary, but with 64-bit offsets, allowing to represent
-/// extremely large data values.
-pub struct LargeBinary<'a> {
- pub _tab: flatbuffers::Table<'a>,
-}
-
-impl<'a> flatbuffers::Follow<'a> for LargeBinary<'a> {
- type Inner = LargeBinary<'a>;
- #[inline]
- fn follow(buf: &'a [u8], loc: usize) -> Self::Inner {
- Self {
- _tab: flatbuffers::Table { buf, loc },
- }
- }
-}
-
-impl<'a> LargeBinary<'a> {
- #[inline]
- pub fn init_from_table(table: flatbuffers::Table<'a>) -> Self {
- LargeBinary { _tab: table }
- }
- #[allow(unused_mut)]
- pub fn create<'bldr: 'args, 'args: 'mut_bldr, 'mut_bldr>(
- _fbb: &'mut_bldr mut flatbuffers::FlatBufferBuilder<'bldr>,
- _args: &'args LargeBinaryArgs,
- ) -> flatbuffers::WIPOffset<LargeBinary<'bldr>> {
- let mut builder = LargeBinaryBuilder::new(_fbb);
- builder.finish()
- }
-}
-
-impl flatbuffers::Verifiable for LargeBinary<'_> {
- #[inline]
- fn run_verifier(
- v: &mut flatbuffers::Verifier,
- pos: usize,
- ) -> Result<(), flatbuffers::InvalidFlatbuffer> {
- use flatbuffers::Verifiable;
- v.visit_table(pos)?.finish();
- Ok(())
- }
-}
-pub struct LargeBinaryArgs {}
-impl<'a> Default for LargeBinaryArgs {
- #[inline]
- fn default() -> Self {
- LargeBinaryArgs {}
- }
-}
-pub struct LargeBinaryBuilder<'a: 'b, 'b> {
- fbb_: &'b mut flatbuffers::FlatBufferBuilder<'a>,
- start_: flatbuffers::WIPOffset<flatbuffers::TableUnfinishedWIPOffset>,
-}
-impl<'a: 'b, 'b> LargeBinaryBuilder<'a, 'b> {
- #[inline]
- pub fn new(
- _fbb: &'b mut flatbuffers::FlatBufferBuilder<'a>,
- ) -> LargeBinaryBuilder<'a, 'b> {
- let start = _fbb.start_table();
- LargeBinaryBuilder {
- fbb_: _fbb,
- start_: start,
- }
- }
- #[inline]
- pub fn finish(self) -> flatbuffers::WIPOffset<LargeBinary<'a>> {
- let o = self.fbb_.end_table(self.start_);
- flatbuffers::WIPOffset::new(o.value())
- }
-}
-
-impl std::fmt::Debug for LargeBinary<'_> {
- fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
- let mut ds = f.debug_struct("LargeBinary");
- ds.finish()
- }
-}
-pub enum FixedSizeBinaryOffset {}
-#[derive(Copy, Clone, PartialEq)]
-
-pub struct FixedSizeBinary<'a> {
- pub _tab: flatbuffers::Table<'a>,
-}
-
-impl<'a> flatbuffers::Follow<'a> for FixedSizeBinary<'a> {
- type Inner = FixedSizeBinary<'a>;
- #[inline]
- fn follow(buf: &'a [u8], loc: usize) -> Self::Inner {
- Self {
- _tab: flatbuffers::Table { buf, loc },
- }
- }
-}
-
-impl<'a> FixedSizeBinary<'a> {
- #[inline]
- pub fn init_from_table(table: flatbuffers::Table<'a>) -> Self {
- FixedSizeBinary { _tab: table }
- }
- #[allow(unused_mut)]
- pub fn create<'bldr: 'args, 'args: 'mut_bldr, 'mut_bldr>(
- _fbb: &'mut_bldr mut flatbuffers::FlatBufferBuilder<'bldr>,
- args: &'args FixedSizeBinaryArgs,
- ) -> flatbuffers::WIPOffset<FixedSizeBinary<'bldr>> {
- let mut builder = FixedSizeBinaryBuilder::new(_fbb);
- builder.add_byteWidth(args.byteWidth);
- builder.finish()
- }
-
- pub const VT_BYTEWIDTH: flatbuffers::VOffsetT = 4;
-
- /// Number of bytes per value
- #[inline]
- pub fn byteWidth(&self) -> i32 {
- self._tab
- .get::<i32>(FixedSizeBinary::VT_BYTEWIDTH, Some(0))
- .unwrap()
- }
-}
-
-impl flatbuffers::Verifiable for FixedSizeBinary<'_> {
- #[inline]
- fn run_verifier(
- v: &mut flatbuffers::Verifier,
- pos: usize,
- ) -> Result<(), flatbuffers::InvalidFlatbuffer> {
- use flatbuffers::Verifiable;
- v.visit_table(pos)?
- .visit_field::<i32>(&"byteWidth", Self::VT_BYTEWIDTH, false)?
- .finish();
- Ok(())
- }
-}
-pub struct FixedSizeBinaryArgs {
- pub byteWidth: i32,
-}
-impl<'a> Default for FixedSizeBinaryArgs {
- #[inline]
- fn default() -> Self {
- FixedSizeBinaryArgs { byteWidth: 0 }
- }
-}
-pub struct FixedSizeBinaryBuilder<'a: 'b, 'b> {
- fbb_: &'b mut flatbuffers::FlatBufferBuilder<'a>,
- start_: flatbuffers::WIPOffset<flatbuffers::TableUnfinishedWIPOffset>,
-}
-impl<'a: 'b, 'b> FixedSizeBinaryBuilder<'a, 'b> {
- #[inline]
- pub fn add_byteWidth(&mut self, byteWidth: i32) {
- self.fbb_
- .push_slot::<i32>(FixedSizeBinary::VT_BYTEWIDTH, byteWidth, 0);
- }
- #[inline]
- pub fn new(
- _fbb: &'b mut flatbuffers::FlatBufferBuilder<'a>,
- ) -> FixedSizeBinaryBuilder<'a, 'b> {
- let start = _fbb.start_table();
- FixedSizeBinaryBuilder {
- fbb_: _fbb,
- start_: start,
- }
- }
- #[inline]
- pub fn finish(self) -> flatbuffers::WIPOffset<FixedSizeBinary<'a>> {
- let o = self.fbb_.end_table(self.start_);
- flatbuffers::WIPOffset::new(o.value())
- }
-}
-
-impl std::fmt::Debug for FixedSizeBinary<'_> {
- fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
- let mut ds = f.debug_struct("FixedSizeBinary");
- ds.field("byteWidth", &self.byteWidth());
- ds.finish()
- }
-}
-pub enum BoolOffset {}
-#[derive(Copy, Clone, PartialEq)]
-
-pub struct Bool<'a> {
- pub _tab: flatbuffers::Table<'a>,
-}
-
-impl<'a> flatbuffers::Follow<'a> for Bool<'a> {
- type Inner = Bool<'a>;
- #[inline]
- fn follow(buf: &'a [u8], loc: usize) -> Self::Inner {
- Self {
- _tab: flatbuffers::Table { buf, loc },
- }
- }
-}
-
-impl<'a> Bool<'a> {
- #[inline]
- pub fn init_from_table(table: flatbuffers::Table<'a>) -> Self {
- Bool { _tab: table }
- }
- #[allow(unused_mut)]
- pub fn create<'bldr: 'args, 'args: 'mut_bldr, 'mut_bldr>(
- _fbb: &'mut_bldr mut flatbuffers::FlatBufferBuilder<'bldr>,
- _args: &'args BoolArgs,
- ) -> flatbuffers::WIPOffset<Bool<'bldr>> {
- let mut builder = BoolBuilder::new(_fbb);
- builder.finish()
- }
-}
-
-impl flatbuffers::Verifiable for Bool<'_> {
- #[inline]
- fn run_verifier(
- v: &mut flatbuffers::Verifier,
- pos: usize,
- ) -> Result<(), flatbuffers::InvalidFlatbuffer> {
- use flatbuffers::Verifiable;
- v.visit_table(pos)?.finish();
- Ok(())
- }
-}
-pub struct BoolArgs {}
-impl<'a> Default for BoolArgs {
- #[inline]
- fn default() -> Self {
- BoolArgs {}
- }
-}
-pub struct BoolBuilder<'a: 'b, 'b> {
- fbb_: &'b mut flatbuffers::FlatBufferBuilder<'a>,
- start_: flatbuffers::WIPOffset<flatbuffers::TableUnfinishedWIPOffset>,
-}
-impl<'a: 'b, 'b> BoolBuilder<'a, 'b> {
- #[inline]
- pub fn new(_fbb: &'b mut flatbuffers::FlatBufferBuilder<'a>) -> BoolBuilder<'a, 'b> {
- let start = _fbb.start_table();
- BoolBuilder {
- fbb_: _fbb,
- start_: start,
- }
- }
- #[inline]
- pub fn finish(self) -> flatbuffers::WIPOffset<Bool<'a>> {
- let o = self.fbb_.end_table(self.start_);
- flatbuffers::WIPOffset::new(o.value())
- }
-}
-
-impl std::fmt::Debug for Bool<'_> {
- fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
- let mut ds = f.debug_struct("Bool");
- ds.finish()
- }
-}
-pub enum DecimalOffset {}
-#[derive(Copy, Clone, PartialEq)]
-
-/// 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.
-pub struct Decimal<'a> {
- pub _tab: flatbuffers::Table<'a>,
-}
-
-impl<'a> flatbuffers::Follow<'a> for Decimal<'a> {
- type Inner = Decimal<'a>;
- #[inline]
- fn follow(buf: &'a [u8], loc: usize) -> Self::Inner {
- Self {
- _tab: flatbuffers::Table { buf, loc },
- }
- }
-}
-
-impl<'a> Decimal<'a> {
- #[inline]
- pub fn init_from_table(table: flatbuffers::Table<'a>) -> Self {
- Decimal { _tab: table }
- }
- #[allow(unused_mut)]
- pub fn create<'bldr: 'args, 'args: 'mut_bldr, 'mut_bldr>(
- _fbb: &'mut_bldr mut flatbuffers::FlatBufferBuilder<'bldr>,
- args: &'args DecimalArgs,
- ) -> flatbuffers::WIPOffset<Decimal<'bldr>> {
- let mut builder = DecimalBuilder::new(_fbb);
- builder.add_bitWidth(args.bitWidth);
- builder.add_scale(args.scale);
- builder.add_precision(args.precision);
- builder.finish()
- }
-
- pub const VT_PRECISION: flatbuffers::VOffsetT = 4;
- pub const VT_SCALE: flatbuffers::VOffsetT = 6;
- pub const VT_BITWIDTH: flatbuffers::VOffsetT = 8;
-
- /// Total number of decimal digits
- #[inline]
- pub fn precision(&self) -> i32 {
- self._tab
- .get::<i32>(Decimal::VT_PRECISION, Some(0))
- .unwrap()
- }
- /// Number of digits after the decimal point "."
- #[inline]
- pub fn scale(&self) -> i32 {
- self._tab.get::<i32>(Decimal::VT_SCALE, Some(0)).unwrap()
- }
- /// Number of bits per value. The only accepted widths are 128 and 256.
- /// We use bitWidth for consistency with Int::bitWidth.
- #[inline]
- pub fn bitWidth(&self) -> i32 {
- self._tab
- .get::<i32>(Decimal::VT_BITWIDTH, Some(128))
- .unwrap()
- }
-}
-
-impl flatbuffers::Verifiable for Decimal<'_> {
- #[inline]
- fn run_verifier(
- v: &mut flatbuffers::Verifier,
- pos: usize,
- ) -> Result<(), flatbuffers::InvalidFlatbuffer> {
- use flatbuffers::Verifiable;
- v.visit_table(pos)?
- .visit_field::<i32>(&"precision", Self::VT_PRECISION, false)?
- .visit_field::<i32>(&"scale", Self::VT_SCALE, false)?
- .visit_field::<i32>(&"bitWidth", Self::VT_BITWIDTH, false)?
- .finish();
- Ok(())
- }
-}
-pub struct DecimalArgs {
- pub precision: i32,
- pub scale: i32,
- pub bitWidth: i32,
-}
-impl<'a> Default for DecimalArgs {
- #[inline]
- fn default() -> Self {
- DecimalArgs {
- precision: 0,
- scale: 0,
- bitWidth: 128,
- }
- }
-}
-pub struct DecimalBuilder<'a: 'b, 'b> {
- fbb_: &'b mut flatbuffers::FlatBufferBuilder<'a>,
- start_: flatbuffers::WIPOffset<flatbuffers::TableUnfinishedWIPOffset>,
-}
-impl<'a: 'b, 'b> DecimalBuilder<'a, 'b> {
- #[inline]
- pub fn add_precision(&mut self, precision: i32) {
- self.fbb_
- .push_slot::<i32>(Decimal::VT_PRECISION, precision, 0);
- }
- #[inline]
- pub fn add_scale(&mut self, scale: i32) {
- self.fbb_.push_slot::<i32>(Decimal::VT_SCALE, scale, 0);
- }
- #[inline]
- pub fn add_bitWidth(&mut self, bitWidth: i32) {
- self.fbb_
- .push_slot::<i32>(Decimal::VT_BITWIDTH, bitWidth, 128);
- }
- #[inline]
- pub fn new(
- _fbb: &'b mut flatbuffers::FlatBufferBuilder<'a>,
- ) -> DecimalBuilder<'a, 'b> {
- let start = _fbb.start_table();
- DecimalBuilder {
- fbb_: _fbb,
- start_: start,
- }
- }
- #[inline]
- pub fn finish(self) -> flatbuffers::WIPOffset<Decimal<'a>> {
- let o = self.fbb_.end_table(self.start_);
- flatbuffers::WIPOffset::new(o.value())
- }
-}
-
-impl std::fmt::Debug for Decimal<'_> {
- fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
- let mut ds = f.debug_struct("Decimal");
- ds.field("precision", &self.precision());
- ds.field("scale", &self.scale());
- ds.field("bitWidth", &self.bitWidth());
- ds.finish()
- }
-}
-pub enum DateOffset {}
-#[derive(Copy, Clone, PartialEq)]
-
-/// Date is either a 32-bit or 64-bit type representing 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
-pub struct Date<'a> {
- pub _tab: flatbuffers::Table<'a>,
-}
-
-impl<'a> flatbuffers::Follow<'a> for Date<'a> {
- type Inner = Date<'a>;
- #[inline]
- fn follow(buf: &'a [u8], loc: usize) -> Self::Inner {
- Self {
- _tab: flatbuffers::Table { buf, loc },
- }
- }
-}
-
-impl<'a> Date<'a> {
- #[inline]
- pub fn init_from_table(table: flatbuffers::Table<'a>) -> Self {
- Date { _tab: table }
- }
- #[allow(unused_mut)]
- pub fn create<'bldr: 'args, 'args: 'mut_bldr, 'mut_bldr>(
- _fbb: &'mut_bldr mut flatbuffers::FlatBufferBuilder<'bldr>,
- args: &'args DateArgs,
- ) -> flatbuffers::WIPOffset<Date<'bldr>> {
- let mut builder = DateBuilder::new(_fbb);
- builder.add_unit(args.unit);
- builder.finish()
- }
-
- pub const VT_UNIT: flatbuffers::VOffsetT = 4;
-
- #[inline]
- pub fn unit(&self) -> DateUnit {
- self._tab
- .get::<DateUnit>(Date::VT_UNIT, Some(DateUnit::MILLISECOND))
- .unwrap()
- }
-}
-
-impl flatbuffers::Verifiable for Date<'_> {
- #[inline]
- fn run_verifier(
- v: &mut flatbuffers::Verifier,
- pos: usize,
- ) -> Result<(), flatbuffers::InvalidFlatbuffer> {
- use flatbuffers::Verifiable;
- v.visit_table(pos)?
- .visit_field::<DateUnit>(&"unit", Self::VT_UNIT, false)?
- .finish();
- Ok(())
- }
-}
-pub struct DateArgs {
- pub unit: DateUnit,
-}
-impl<'a> Default for DateArgs {
- #[inline]
- fn default() -> Self {
- DateArgs {
- unit: DateUnit::MILLISECOND,
- }
- }
-}
-pub struct DateBuilder<'a: 'b, 'b> {
- fbb_: &'b mut flatbuffers::FlatBufferBuilder<'a>,
- start_: flatbuffers::WIPOffset<flatbuffers::TableUnfinishedWIPOffset>,
-}
-impl<'a: 'b, 'b> DateBuilder<'a, 'b> {
- #[inline]
- pub fn add_unit(&mut self, unit: DateUnit) {
- self.fbb_
- .push_slot::<DateUnit>(Date::VT_UNIT, unit, DateUnit::MILLISECOND);
- }
- #[inline]
- pub fn new(_fbb: &'b mut flatbuffers::FlatBufferBuilder<'a>) -> DateBuilder<'a, 'b> {
- let start = _fbb.start_table();
- DateBuilder {
- fbb_: _fbb,
- start_: start,
- }
- }
- #[inline]
- pub fn finish(self) -> flatbuffers::WIPOffset<Date<'a>> {
- let o = self.fbb_.end_table(self.start_);
- flatbuffers::WIPOffset::new(o.value())
- }
-}
-
-impl std::fmt::Debug for Date<'_> {
- fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
- let mut ds = f.debug_struct("Date");
- ds.field("unit", &self.unit());
- ds.finish()
- }
-}
-pub enum TimeOffset {}
-#[derive(Copy, Clone, PartialEq)]
-
-/// Time type. The physical storage type depends on the unit
-/// - SECOND and MILLISECOND: 32 bits
-/// - MICROSECOND and NANOSECOND: 64 bits
-pub struct Time<'a> {
- pub _tab: flatbuffers::Table<'a>,
-}
-
-impl<'a> flatbuffers::Follow<'a> for Time<'a> {
- type Inner = Time<'a>;
- #[inline]
- fn follow(buf: &'a [u8], loc: usize) -> Self::Inner {
- Self {
- _tab: flatbuffers::Table { buf, loc },
- }
- }
-}
-
-impl<'a> Time<'a> {
- #[inline]
- pub fn init_from_table(table: flatbuffers::Table<'a>) -> Self {
- Time { _tab: table }
- }
- #[allow(unused_mut)]
- pub fn create<'bldr: 'args, 'args: 'mut_bldr, 'mut_bldr>(
- _fbb: &'mut_bldr mut flatbuffers::FlatBufferBuilder<'bldr>,
- args: &'args TimeArgs,
- ) -> flatbuffers::WIPOffset<Time<'bldr>> {
- let mut builder = TimeBuilder::new(_fbb);
- builder.add_bitWidth(args.bitWidth);
- builder.add_unit(args.unit);
- builder.finish()
- }
-
- pub const VT_UNIT: flatbuffers::VOffsetT = 4;
- pub const VT_BITWIDTH: flatbuffers::VOffsetT = 6;
-
- #[inline]
- pub fn unit(&self) -> TimeUnit {
- self._tab
- .get::<TimeUnit>(Time::VT_UNIT, Some(TimeUnit::MILLISECOND))
- .unwrap()
- }
- #[inline]
- pub fn bitWidth(&self) -> i32 {
- self._tab.get::<i32>(Time::VT_BITWIDTH, Some(32)).unwrap()
- }
-}
-
-impl flatbuffers::Verifiable for Time<'_> {
- #[inline]
- fn run_verifier(
- v: &mut flatbuffers::Verifier,
- pos: usize,
- ) -> Result<(), flatbuffers::InvalidFlatbuffer> {
- use flatbuffers::Verifiable;
- v.visit_table(pos)?
- .visit_field::<TimeUnit>(&"unit", Self::VT_UNIT, false)?
- .visit_field::<i32>(&"bitWidth", Self::VT_BITWIDTH, false)?
- .finish();
- Ok(())
- }
-}
-pub struct TimeArgs {
- pub unit: TimeUnit,
- pub bitWidth: i32,
-}
-impl<'a> Default for TimeArgs {
- #[inline]
- fn default() -> Self {
- TimeArgs {
- unit: TimeUnit::MILLISECOND,
- bitWidth: 32,
- }
- }
-}
-pub struct TimeBuilder<'a: 'b, 'b> {
- fbb_: &'b mut flatbuffers::FlatBufferBuilder<'a>,
- start_: flatbuffers::WIPOffset<flatbuffers::TableUnfinishedWIPOffset>,
-}
-impl<'a: 'b, 'b> TimeBuilder<'a, 'b> {
- #[inline]
- pub fn add_unit(&mut self, unit: TimeUnit) {
- self.fbb_
- .push_slot::<TimeUnit>(Time::VT_UNIT, unit, TimeUnit::MILLISECOND);
- }
- #[inline]
- pub fn add_bitWidth(&mut self, bitWidth: i32) {
- self.fbb_.push_slot::<i32>(Time::VT_BITWIDTH, bitWidth, 32);
- }
- #[inline]
- pub fn new(_fbb: &'b mut flatbuffers::FlatBufferBuilder<'a>) -> TimeBuilder<'a, 'b> {
- let start = _fbb.start_table();
- TimeBuilder {
- fbb_: _fbb,
- start_: start,
- }
- }
- #[inline]
- pub fn finish(self) -> flatbuffers::WIPOffset<Time<'a>> {
- let o = self.fbb_.end_table(self.start_);
- flatbuffers::WIPOffset::new(o.value())
- }
-}
-
-impl std::fmt::Debug for Time<'_> {
- fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
- let mut ds = f.debug_struct("Time");
- ds.field("unit", &self.unit());
- ds.field("bitWidth", &self.bitWidth());
- ds.finish()
- }
-}
-pub enum TimestampOffset {}
-#[derive(Copy, Clone, PartialEq)]
-
-/// Time elapsed from the Unix epoch, 00:00:00.000 on 1 January 1970, excluding
-/// leap seconds, as a 64-bit integer. Note that UNIX time does not include
-/// leap seconds.
-///
-/// The Timestamp metadata supports both "time zone naive" and "time zone
-/// aware" timestamps. Read about the timezone attribute for more detail
-pub struct Timestamp<'a> {
- pub _tab: flatbuffers::Table<'a>,
-}
-
-impl<'a> flatbuffers::Follow<'a> for Timestamp<'a> {
- type Inner = Timestamp<'a>;
- #[inline]
- fn follow(buf: &'a [u8], loc: usize) -> Self::Inner {
- Self {
- _tab: flatbuffers::Table { buf, loc },
- }
- }
-}
-
-impl<'a> Timestamp<'a> {
- #[inline]
- pub fn init_from_table(table: flatbuffers::Table<'a>) -> Self {
- Timestamp { _tab: table }
- }
- #[allow(unused_mut)]
- pub fn create<'bldr: 'args, 'args: 'mut_bldr, 'mut_bldr>(
- _fbb: &'mut_bldr mut flatbuffers::FlatBufferBuilder<'bldr>,
- args: &'args TimestampArgs<'args>,
- ) -> flatbuffers::WIPOffset<Timestamp<'bldr>> {
- let mut builder = TimestampBuilder::new(_fbb);
- if let Some(x) = args.timezone {
- builder.add_timezone(x);
- }
- builder.add_unit(args.unit);
- builder.finish()
- }
-
- pub const VT_UNIT: flatbuffers::VOffsetT = 4;
- pub const VT_TIMEZONE: flatbuffers::VOffsetT = 6;
-
- #[inline]
- pub fn unit(&self) -> TimeUnit {
- self._tab
- .get::<TimeUnit>(Timestamp::VT_UNIT, Some(TimeUnit::SECOND))
- .unwrap()
- }
- /// The time zone is a string indicating the name of a time zone, one of:
- ///
- /// * As used in the Olson time zone database (the "tz database" or
- /// "tzdata"), such as "America/New_York"
- /// * An absolute time zone 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:
- ///
- /// * If the time zone is null or equal to an empty string, the data is "time
- /// zone naive" and shall be displayed *as is* to the user, not localized
- /// to the locale of the user. This data can be though of as UTC but
- /// without having "UTC" as the time zone, it is not considered to be
- /// localized to any time zone
- ///
- /// * If the time zone is set to a valid value, values can be displayed as
- /// "localized" to that time zone, even though the underlying 64-bit
- /// integers are identical to the same data stored in UTC. Converting
- /// between time zones is a metadata-only operation and does not change the
- /// underlying values
- #[inline]
- pub fn timezone(&self) -> Option<&'a str> {
- self._tab
- .get::<flatbuffers::ForwardsUOffset<&str>>(Timestamp::VT_TIMEZONE, None)
- }
-}
-
-impl flatbuffers::Verifiable for Timestamp<'_> {
- #[inline]
- fn run_verifier(
- v: &mut flatbuffers::Verifier,
- pos: usize,
- ) -> Result<(), flatbuffers::InvalidFlatbuffer> {
- use flatbuffers::Verifiable;
- v.visit_table(pos)?
- .visit_field::<TimeUnit>(&"unit", Self::VT_UNIT, false)?
- .visit_field::<flatbuffers::ForwardsUOffset<&str>>(
- &"timezone",
- Self::VT_TIMEZONE,
- false,
- )?
- .finish();
- Ok(())
- }
-}
-pub struct TimestampArgs<'a> {
- pub unit: TimeUnit,
- pub timezone: Option<flatbuffers::WIPOffset<&'a str>>,
-}
-impl<'a> Default for TimestampArgs<'a> {
- #[inline]
- fn default() -> Self {
- TimestampArgs {
- unit: TimeUnit::SECOND,
- timezone: None,
- }
- }
-}
-pub struct TimestampBuilder<'a: 'b, 'b> {
- fbb_: &'b mut flatbuffers::FlatBufferBuilder<'a>,
- start_: flatbuffers::WIPOffset<flatbuffers::TableUnfinishedWIPOffset>,
-}
-impl<'a: 'b, 'b> TimestampBuilder<'a, 'b> {
- #[inline]
- pub fn add_unit(&mut self, unit: TimeUnit) {
- self.fbb_
- .push_slot::<TimeUnit>(Timestamp::VT_UNIT, unit, TimeUnit::SECOND);
- }
- #[inline]
- pub fn add_timezone(&mut self, timezone: flatbuffers::WIPOffset<&'b str>) {
- self.fbb_.push_slot_always::<flatbuffers::WIPOffset<_>>(
- Timestamp::VT_TIMEZONE,
- timezone,
- );
- }
- #[inline]
- pub fn new(
- _fbb: &'b mut flatbuffers::FlatBufferBuilder<'a>,
- ) -> TimestampBuilder<'a, 'b> {
- let start = _fbb.start_table();
- TimestampBuilder {
- fbb_: _fbb,
- start_: start,
- }
- }
- #[inline]
- pub fn finish(self) -> flatbuffers::WIPOffset<Timestamp<'a>> {
- let o = self.fbb_.end_table(self.start_);
- flatbuffers::WIPOffset::new(o.value())
- }
-}
-
-impl std::fmt::Debug for Timestamp<'_> {
- fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
- let mut ds = f.debug_struct("Timestamp");
- ds.field("unit", &self.unit());
- ds.field("timezone", &self.timezone());
- ds.finish()
- }
-}
-pub enum IntervalOffset {}
-#[derive(Copy, Clone, PartialEq)]
-
-pub struct Interval<'a> {
- pub _tab: flatbuffers::Table<'a>,
-}
-
-impl<'a> flatbuffers::Follow<'a> for Interval<'a> {
- type Inner = Interval<'a>;
- #[inline]
- fn follow(buf: &'a [u8], loc: usize) -> Self::Inner {
- Self {
- _tab: flatbuffers::Table { buf, loc },
- }
- }
-}
-
-impl<'a> Interval<'a> {
- #[inline]
- pub fn init_from_table(table: flatbuffers::Table<'a>) -> Self {
- Interval { _tab: table }
- }
- #[allow(unused_mut)]
- pub fn create<'bldr: 'args, 'args: 'mut_bldr, 'mut_bldr>(
- _fbb: &'mut_bldr mut flatbuffers::FlatBufferBuilder<'bldr>,
- args: &'args IntervalArgs,
- ) -> flatbuffers::WIPOffset<Interval<'bldr>> {
- let mut builder = IntervalBuilder::new(_fbb);
- builder.add_unit(args.unit);
- builder.finish()
- }
-
- pub const VT_UNIT: flatbuffers::VOffsetT = 4;
-
- #[inline]
- pub fn unit(&self) -> IntervalUnit {
- self._tab
- .get::<IntervalUnit>(Interval::VT_UNIT, Some(IntervalUnit::YEAR_MONTH))
- .unwrap()
- }
-}
-
-impl flatbuffers::Verifiable for Interval<'_> {
- #[inline]
- fn run_verifier(
- v: &mut flatbuffers::Verifier,
- pos: usize,
- ) -> Result<(), flatbuffers::InvalidFlatbuffer> {
- use flatbuffers::Verifiable;
- v.visit_table(pos)?
- .visit_field::<IntervalUnit>(&"unit", Self::VT_UNIT, false)?
- .finish();
- Ok(())
- }
-}
-pub struct IntervalArgs {
- pub unit: IntervalUnit,
-}
-impl<'a> Default for IntervalArgs {
- #[inline]
- fn default() -> Self {
- IntervalArgs {
- unit: IntervalUnit::YEAR_MONTH,
- }
- }
-}
-pub struct IntervalBuilder<'a: 'b, 'b> {
- fbb_: &'b mut flatbuffers::FlatBufferBuilder<'a>,
- start_: flatbuffers::WIPOffset<flatbuffers::TableUnfinishedWIPOffset>,
-}
-impl<'a: 'b, 'b> IntervalBuilder<'a, 'b> {
- #[inline]
- pub fn add_unit(&mut self, unit: IntervalUnit) {
- self.fbb_.push_slot::<IntervalUnit>(
- Interval::VT_UNIT,
- unit,
- IntervalUnit::YEAR_MONTH,
- );
- }
- #[inline]
- pub fn new(
- _fbb: &'b mut flatbuffers::FlatBufferBuilder<'a>,
- ) -> IntervalBuilder<'a, 'b> {
- let start = _fbb.start_table();
- IntervalBuilder {
- fbb_: _fbb,
- start_: start,
- }
- }
- #[inline]
- pub fn finish(self) -> flatbuffers::WIPOffset<Interval<'a>> {
- let o = self.fbb_.end_table(self.start_);
- flatbuffers::WIPOffset::new(o.value())
- }
-}
-
-impl std::fmt::Debug for Interval<'_> {
- fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
- let mut ds = f.debug_struct("Interval");
- ds.field("unit", &self.unit());
- ds.finish()
- }
-}
-pub enum DurationOffset {}
-#[derive(Copy, Clone, PartialEq)]
-
-pub struct Duration<'a> {
- pub _tab: flatbuffers::Table<'a>,
-}
-
-impl<'a> flatbuffers::Follow<'a> for Duration<'a> {
- type Inner = Duration<'a>;
- #[inline]
- fn follow(buf: &'a [u8], loc: usize) -> Self::Inner {
- Self {
- _tab: flatbuffers::Table { buf, loc },
- }
- }
-}
-
-impl<'a> Duration<'a> {
- #[inline]
- pub fn init_from_table(table: flatbuffers::Table<'a>) -> Self {
- Duration { _tab: table }
- }
- #[allow(unused_mut)]
- pub fn create<'bldr: 'args, 'args: 'mut_bldr, 'mut_bldr>(
- _fbb: &'mut_bldr mut flatbuffers::FlatBufferBuilder<'bldr>,
- args: &'args DurationArgs,
- ) -> flatbuffers::WIPOffset<Duration<'bldr>> {
- let mut builder = DurationBuilder::new(_fbb);
- builder.add_unit(args.unit);
- builder.finish()
- }
-
- pub const VT_UNIT: flatbuffers::VOffsetT = 4;
-
- #[inline]
- pub fn unit(&self) -> TimeUnit {
- self._tab
- .get::<TimeUnit>(Duration::VT_UNIT, Some(TimeUnit::MILLISECOND))
- .unwrap()
- }
-}
-
-impl flatbuffers::Verifiable for Duration<'_> {
- #[inline]
- fn run_verifier(
- v: &mut flatbuffers::Verifier,
- pos: usize,
- ) -> Result<(), flatbuffers::InvalidFlatbuffer> {
- use flatbuffers::Verifiable;
- v.visit_table(pos)?
- .visit_field::<TimeUnit>(&"unit", Self::VT_UNIT, false)?
- .finish();
- Ok(())
- }
-}
-pub struct DurationArgs {
- pub unit: TimeUnit,
-}
-impl<'a> Default for DurationArgs {
- #[inline]
- fn default() -> Self {
- DurationArgs {
- unit: TimeUnit::MILLISECOND,
- }
- }
-}
-pub struct DurationBuilder<'a: 'b, 'b> {
- fbb_: &'b mut flatbuffers::FlatBufferBuilder<'a>,
- start_: flatbuffers::WIPOffset<flatbuffers::TableUnfinishedWIPOffset>,
-}
-impl<'a: 'b, 'b> DurationBuilder<'a, 'b> {
- #[inline]
- pub fn add_unit(&mut self, unit: TimeUnit) {
- self.fbb_
- .push_slot::<TimeUnit>(Duration::VT_UNIT, unit, TimeUnit::MILLISECOND);
- }
- #[inline]
- pub fn new(
- _fbb: &'b mut flatbuffers::FlatBufferBuilder<'a>,
- ) -> DurationBuilder<'a, 'b> {
- let start = _fbb.start_table();
- DurationBuilder {
- fbb_: _fbb,
- start_: start,
- }
- }
- #[inline]
- pub fn finish(self) -> flatbuffers::WIPOffset<Duration<'a>> {
- let o = self.fbb_.end_table(self.start_);
- flatbuffers::WIPOffset::new(o.value())
- }
-}
-
-impl std::fmt::Debug for Duration<'_> {
- fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
- let mut ds = f.debug_struct("Duration");
- ds.field("unit", &self.unit());
- ds.finish()
- }
-}
-pub enum KeyValueOffset {}
-#[derive(Copy, Clone, PartialEq)]
-
-/// ----------------------------------------------------------------------
-/// user defined key value pairs to add custom metadata to arrow
-/// key namespacing is the responsibility of the user
-pub struct KeyValue<'a> {
- pub _tab: flatbuffers::Table<'a>,
-}
-
-impl<'a> flatbuffers::Follow<'a> for KeyValue<'a> {
- type Inner = KeyValue<'a>;
- #[inline]
- fn follow(buf: &'a [u8], loc: usize) -> Self::Inner {
- Self {
- _tab: flatbuffers::Table { buf, loc },
- }
- }
-}
-
-impl<'a> KeyValue<'a> {
- #[inline]
- pub fn init_from_table(table: flatbuffers::Table<'a>) -> Self {
- KeyValue { _tab: table }
- }
- #[allow(unused_mut)]
- pub fn create<'bldr: 'args, 'args: 'mut_bldr, 'mut_bldr>(
- _fbb: &'mut_bldr mut flatbuffers::FlatBufferBuilder<'bldr>,
- args: &'args KeyValueArgs<'args>,
- ) -> flatbuffers::WIPOffset<KeyValue<'bldr>> {
- let mut builder = KeyValueBuilder::new(_fbb);
- if let Some(x) = args.value {
- builder.add_value(x);
- }
- if let Some(x) = args.key {
- builder.add_key(x);
- }
- builder.finish()
- }
-
- pub const VT_KEY: flatbuffers::VOffsetT = 4;
- pub const VT_VALUE: flatbuffers::VOffsetT = 6;
-
- #[inline]
- pub fn key(&self) -> Option<&'a str> {
- self._tab
- .get::<flatbuffers::ForwardsUOffset<&str>>(KeyValue::VT_KEY, None)
- }
- #[inline]
- pub fn value(&self) -> Option<&'a str> {
- self._tab
- .get::<flatbuffers::ForwardsUOffset<&str>>(KeyValue::VT_VALUE, None)
- }
-}
-
-impl flatbuffers::Verifiable for KeyValue<'_> {
- #[inline]
- fn run_verifier(
- v: &mut flatbuffers::Verifier,
- pos: usize,
- ) -> Result<(), flatbuffers::InvalidFlatbuffer> {
- use flatbuffers::Verifiable;
- v.visit_table(pos)?
- .visit_field::<flatbuffers::ForwardsUOffset<&str>>(
- &"key",
- Self::VT_KEY,
- false,
- )?
- .visit_field::<flatbuffers::ForwardsUOffset<&str>>(
- &"value",
- Self::VT_VALUE,
- false,
- )?
- .finish();
- Ok(())
- }
-}
-pub struct KeyValueArgs<'a> {
- pub key: Option<flatbuffers::WIPOffset<&'a str>>,
- pub value: Option<flatbuffers::WIPOffset<&'a str>>,
-}
-impl<'a> Default for KeyValueArgs<'a> {
- #[inline]
- fn default() -> Self {
- KeyValueArgs {
- key: None,
- value: None,
- }
- }
-}
-pub struct KeyValueBuilder<'a: 'b, 'b> {
- fbb_: &'b mut flatbuffers::FlatBufferBuilder<'a>,
- start_: flatbuffers::WIPOffset<flatbuffers::TableUnfinishedWIPOffset>,
-}
-impl<'a: 'b, 'b> KeyValueBuilder<'a, 'b> {
- #[inline]
- pub fn add_key(&mut self, key: flatbuffers::WIPOffset<&'b str>) {
- self.fbb_
- .push_slot_always::<flatbuffers::WIPOffset<_>>(KeyValue::VT_KEY, key);
- }
- #[inline]
- pub fn add_value(&mut self, value: flatbuffers::WIPOffset<&'b str>) {
- self.fbb_
- .push_slot_always::<flatbuffers::WIPOffset<_>>(KeyValue::VT_VALUE, value);
- }
- #[inline]
- pub fn new(
- _fbb: &'b mut flatbuffers::FlatBufferBuilder<'a>,
- ) -> KeyValueBuilder<'a, 'b> {
- let start = _fbb.start_table();
- KeyValueBuilder {
- fbb_: _fbb,
- start_: start,
- }
- }
- #[inline]
- pub fn finish(self) -> flatbuffers::WIPOffset<KeyValue<'a>> {
- let o = self.fbb_.end_table(self.start_);
- flatbuffers::WIPOffset::new(o.value())
- }
-}
-
-impl std::fmt::Debug for KeyValue<'_> {
- fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
- let mut ds = f.debug_struct("KeyValue");
- ds.field("key", &self.key());
- ds.field("value", &self.value());
- ds.finish()
- }
-}
-pub enum DictionaryEncodingOffset {}
-#[derive(Copy, Clone, PartialEq)]
-
-pub struct DictionaryEncoding<'a> {
- pub _tab: flatbuffers::Table<'a>,
-}
-
-impl<'a> flatbuffers::Follow<'a> for DictionaryEncoding<'a> {
- type Inner = DictionaryEncoding<'a>;
- #[inline]
- fn follow(buf: &'a [u8], loc: usize) -> Self::Inner {
- Self {
- _tab: flatbuffers::Table { buf, loc },
- }
- }
-}
-
-impl<'a> DictionaryEncoding<'a> {
- #[inline]
- pub fn init_from_table(table: flatbuffers::Table<'a>) -> Self {
- DictionaryEncoding { _tab: table }
- }
- #[allow(unused_mut)]
- pub fn create<'bldr: 'args, 'args: 'mut_bldr, 'mut_bldr>(
- _fbb: &'mut_bldr mut flatbuffers::FlatBufferBuilder<'bldr>,
- args: &'args DictionaryEncodingArgs<'args>,
- ) -> flatbuffers::WIPOffset<DictionaryEncoding<'bldr>> {
- let mut builder = DictionaryEncodingBuilder::new(_fbb);
- builder.add_id(args.id);
- if let Some(x) = args.indexType {
- builder.add_indexType(x);
- }
- builder.add_dictionaryKind(args.dictionaryKind);
- builder.add_isOrdered(args.isOrdered);
- builder.finish()
- }
-
- pub const VT_ID: flatbuffers::VOffsetT = 4;
- pub const VT_INDEXTYPE: flatbuffers::VOffsetT = 6;
- pub const VT_ISORDERED: flatbuffers::VOffsetT = 8;
- pub const VT_DICTIONARYKIND: flatbuffers::VOffsetT = 10;
-
- /// 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
- #[inline]
- pub fn id(&self) -> i64 {
- self._tab
- .get::<i64>(DictionaryEncoding::VT_ID, Some(0))
- .unwrap()
- }
- /// 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.
- #[inline]
- pub fn indexType(&self) -> Option<Int<'a>> {
- self._tab.get::<flatbuffers::ForwardsUOffset<Int>>(
- DictionaryEncoding::VT_INDEXTYPE,
- None,
- )
- }
- /// 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
- #[inline]
- pub fn isOrdered(&self) -> bool {
- self._tab
- .get::<bool>(DictionaryEncoding::VT_ISORDERED, Some(false))
- .unwrap()
- }
- #[inline]
- pub fn dictionaryKind(&self) -> DictionaryKind {
- self._tab
- .get::<DictionaryKind>(
- DictionaryEncoding::VT_DICTIONARYKIND,
- Some(DictionaryKind::DenseArray),
- )
- .unwrap()
- }
-}
-
-impl flatbuffers::Verifiable for DictionaryEncoding<'_> {
- #[inline]
- fn run_verifier(
- v: &mut flatbuffers::Verifier,
- pos: usize,
- ) -> Result<(), flatbuffers::InvalidFlatbuffer> {
- use flatbuffers::Verifiable;
- v.visit_table(pos)?
- .visit_field::<i64>(&"id", Self::VT_ID, false)?
- .visit_field::<flatbuffers::ForwardsUOffset<Int>>(
- &"indexType",
- Self::VT_INDEXTYPE,
- false,
- )?
- .visit_field::<bool>(&"isOrdered", Self::VT_ISORDERED, false)?
- .visit_field::<DictionaryKind>(
- &"dictionaryKind",
- Self::VT_DICTIONARYKIND,
- false,
- )?
- .finish();
- Ok(())
- }
-}
-pub struct DictionaryEncodingArgs<'a> {
- pub id: i64,
- pub indexType: Option<flatbuffers::WIPOffset<Int<'a>>>,
- pub isOrdered: bool,
- pub dictionaryKind: DictionaryKind,
-}
-impl<'a> Default for DictionaryEncodingArgs<'a> {
- #[inline]
- fn default() -> Self {
- DictionaryEncodingArgs {
- id: 0,
- indexType: None,
- isOrdered: false,
- dictionaryKind: DictionaryKind::DenseArray,
- }
- }
-}
-pub struct DictionaryEncodingBuilder<'a: 'b, 'b> {
- fbb_: &'b mut flatbuffers::FlatBufferBuilder<'a>,
- start_: flatbuffers::WIPOffset<flatbuffers::TableUnfinishedWIPOffset>,
-}
-impl<'a: 'b, 'b> DictionaryEncodingBuilder<'a, 'b> {
- #[inline]
- pub fn add_id(&mut self, id: i64) {
- self.fbb_.push_slot::<i64>(DictionaryEncoding::VT_ID, id, 0);
- }
- #[inline]
- pub fn add_indexType(&mut self, indexType: flatbuffers::WIPOffset<Int<'b>>) {
- self.fbb_.push_slot_always::<flatbuffers::WIPOffset<Int>>(
- DictionaryEncoding::VT_INDEXTYPE,
- indexType,
- );
- }
- #[inline]
- pub fn add_isOrdered(&mut self, isOrdered: bool) {
- self.fbb_
- .push_slot::<bool>(DictionaryEncoding::VT_ISORDERED, isOrdered, false);
- }
- #[inline]
- pub fn add_dictionaryKind(&mut self, dictionaryKind: DictionaryKind) {
- self.fbb_.push_slot::<DictionaryKind>(
- DictionaryEncoding::VT_DICTIONARYKIND,
- dictionaryKind,
- DictionaryKind::DenseArray,
- );
- }
- #[inline]
- pub fn new(
- _fbb: &'b mut flatbuffers::FlatBufferBuilder<'a>,
- ) -> DictionaryEncodingBuilder<'a, 'b> {
- let start = _fbb.start_table();
- DictionaryEncodingBuilder {
- fbb_: _fbb,
- start_: start,
- }
- }
- #[inline]
- pub fn finish(self) -> flatbuffers::WIPOffset<DictionaryEncoding<'a>> {
- let o = self.fbb_.end_table(self.start_);
- flatbuffers::WIPOffset::new(o.value())
- }
-}
-
-impl std::fmt::Debug for DictionaryEncoding<'_> {
- fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
- let mut ds = f.debug_struct("DictionaryEncoding");
- ds.field("id", &self.id());
- ds.field("indexType", &self.indexType());
- ds.field("isOrdered", &self.isOrdered());
- ds.field("dictionaryKind", &self.dictionaryKind());
- ds.finish()
- }
-}
-pub enum FieldOffset {}
-#[derive(Copy, Clone, PartialEq)]
-
-/// ----------------------------------------------------------------------
-/// A field represents a named column in a record / row batch or child of a
-/// nested type.
-pub struct Field<'a> {
- pub _tab: flatbuffers::Table<'a>,
-}
-
-impl<'a> flatbuffers::Follow<'a> for Field<'a> {
- type Inner = Field<'a>;
- #[inline]
- fn follow(buf: &'a [u8], loc: usize) -> Self::Inner {
- Self {
- _tab: flatbuffers::Table { buf, loc },
- }
- }
-}
-
-impl<'a> Field<'a> {
- #[inline]
- pub fn init_from_table(table: flatbuffers::Table<'a>) -> Self {
- Field { _tab: table }
- }
- #[allow(unused_mut)]
- pub fn create<'bldr: 'args, 'args: 'mut_bldr, 'mut_bldr>(
- _fbb: &'mut_bldr mut flatbuffers::FlatBufferBuilder<'bldr>,
- args: &'args FieldArgs<'args>,
- ) -> flatbuffers::WIPOffset<Field<'bldr>> {
- let mut builder = FieldBuilder::new(_fbb);
- if let Some(x) = args.custom_metadata {
- builder.add_custom_metadata(x);
- }
- if let Some(x) = args.children {
- builder.add_children(x);
- }
- if let Some(x) = args.dictionary {
- builder.add_dictionary(x);
- }
- if let Some(x) = args.type_ {
- builder.add_type_(x);
- }
- if let Some(x) = args.name {
- builder.add_name(x);
- }
- builder.add_type_type(args.type_type);
- builder.add_nullable(args.nullable);
- builder.finish()
- }
-
- pub const VT_NAME: flatbuffers::VOffsetT = 4;
- pub const VT_NULLABLE: flatbuffers::VOffsetT = 6;
- pub const VT_TYPE_TYPE: flatbuffers::VOffsetT = 8;
- pub const VT_TYPE_: flatbuffers::VOffsetT = 10;
- pub const VT_DICTIONARY: flatbuffers::VOffsetT = 12;
- pub const VT_CHILDREN: flatbuffers::VOffsetT = 14;
- pub const VT_CUSTOM_METADATA: flatbuffers::VOffsetT = 16;
-
- /// Name is not required, in i.e. a List
- #[inline]
- pub fn name(&self) -> Option<&'a str> {
- self._tab
- .get::<flatbuffers::ForwardsUOffset<&str>>(Field::VT_NAME, None)
- }
- /// Whether or not this field can contain nulls. Should be true in general.
- #[inline]
- pub fn nullable(&self) -> bool {
- self._tab
- .get::<bool>(Field::VT_NULLABLE, Some(false))
- .unwrap()
- }
- #[inline]
- pub fn type_type(&self) -> Type {
- self._tab
- .get::<Type>(Field::VT_TYPE_TYPE, Some(Type::NONE))
- .unwrap()
- }
- /// This is the type of the decoded value if the field is dictionary encoded.
- #[inline]
- pub fn type_(&self) -> Option<flatbuffers::Table<'a>> {
- self._tab
- .get::<flatbuffers::ForwardsUOffset<flatbuffers::Table<'a>>>(
- Field::VT_TYPE_,
- None,
- )
- }
- /// Present only if the field is dictionary encoded.
- #[inline]
- pub fn dictionary(&self) -> Option<DictionaryEncoding<'a>> {
- self._tab
- .get::<flatbuffers::ForwardsUOffset<DictionaryEncoding>>(
- Field::VT_DICTIONARY,
- None,
- )
- }
- /// children apply only to nested data types like Struct, List and Union. For
- /// primitive types children will have length 0.
- #[inline]
- pub fn children(
- &self,
- ) -> Option<flatbuffers::Vector<'a, flatbuffers::ForwardsUOffset<Field<'a>>>> {
- self._tab.get::<flatbuffers::ForwardsUOffset<
- flatbuffers::Vector<'a, flatbuffers::ForwardsUOffset<Field>>,
- >>(Field::VT_CHILDREN, None)
- }
- /// User-defined metadata
- #[inline]
- pub fn custom_metadata(
- &self,
- ) -> Option<flatbuffers::Vector<'a, flatbuffers::ForwardsUOffset<KeyValue<'a>>>> {
- self._tab.get::<flatbuffers::ForwardsUOffset<
- flatbuffers::Vector<'a, flatbuffers::ForwardsUOffset<KeyValue>>,
- >>(Field::VT_CUSTOM_METADATA, None)
- }
- #[inline]
- #[allow(non_snake_case)]
- pub fn type_as_null(&self) -> Option<Null<'a>> {
- if self.type_type() == Type::Null {
- self.type_().map(Null::init_from_table)
- } else {
- None
- }
- }
-
- #[inline]
- #[allow(non_snake_case)]
- pub fn type_as_int(&self) -> Option<Int<'a>> {
- if self.type_type() == Type::Int {
- self.type_().map(Int::init_from_table)
- } else {
- None
- }
- }
-
- #[inline]
- #[allow(non_snake_case)]
- pub fn type_as_floating_point(&self) -> Option<FloatingPoint<'a>> {
- if self.type_type() == Type::FloatingPoint {
- self.type_().map(FloatingPoint::init_from_table)
- } else {
- None
- }
- }
-
- #[inline]
- #[allow(non_snake_case)]
- pub fn type_as_binary(&self) -> Option<Binary<'a>> {
- if self.type_type() == Type::Binary {
- self.type_().map(Binary::init_from_table)
- } else {
- None
- }
- }
-
- #[inline]
- #[allow(non_snake_case)]
- pub fn type_as_utf_8(&self) -> Option<Utf8<'a>> {
- if self.type_type() == Type::Utf8 {
- self.type_().map(Utf8::init_from_table)
- } else {
- None
- }
- }
-
- #[inline]
- #[allow(non_snake_case)]
- pub fn type_as_bool(&self) -> Option<Bool<'a>> {
- if self.type_type() == Type::Bool {
- self.type_().map(Bool::init_from_table)
- } else {
- None
- }
- }
-
- #[inline]
- #[allow(non_snake_case)]
- pub fn type_as_decimal(&self) -> Option<Decimal<'a>> {
- if self.type_type() == Type::Decimal {
- self.type_().map(Decimal::init_from_table)
- } else {
- None
- }
- }
-
- #[inline]
- #[allow(non_snake_case)]
- pub fn type_as_date(&self) -> Option<Date<'a>> {
- if self.type_type() == Type::Date {
- self.type_().map(Date::init_from_table)
- } else {
- None
- }
- }
-
- #[inline]
- #[allow(non_snake_case)]
- pub fn type_as_time(&self) -> Option<Time<'a>> {
- if self.type_type() == Type::Time {
- self.type_().map(Time::init_from_table)
- } else {
- None
- }
- }
-
- #[inline]
- #[allow(non_snake_case)]
- pub fn type_as_timestamp(&self) -> Option<Timestamp<'a>> {
- if self.type_type() == Type::Timestamp {
- self.type_().map(Timestamp::init_from_table)
- } else {
- None
- }
- }
-
- #[inline]
- #[allow(non_snake_case)]
- pub fn type_as_interval(&self) -> Option<Interval<'a>> {
- if self.type_type() == Type::Interval {
- self.type_().map(Interval::init_from_table)
- } else {
- None
- }
- }
-
- #[inline]
- #[allow(non_snake_case)]
- pub fn type_as_list(&self) -> Option<List<'a>> {
- if self.type_type() == Type::List {
- self.type_().map(List::init_from_table)
- } else {
- None
- }
- }
-
- #[inline]
- #[allow(non_snake_case)]
- pub fn type_as_struct_(&self) -> Option<Struct_<'a>> {
- if self.type_type() == Type::Struct_ {
- self.type_().map(Struct_::init_from_table)
- } else {
- None
- }
- }
-
- #[inline]
- #[allow(non_snake_case)]
- pub fn type_as_union(&self) -> Option<Union<'a>> {
- if self.type_type() == Type::Union {
- self.type_().map(Union::init_from_table)
- } else {
- None
- }
- }
-
- #[inline]
- #[allow(non_snake_case)]
- pub fn type_as_fixed_size_binary(&self) -> Option<FixedSizeBinary<'a>> {
- if self.type_type() == Type::FixedSizeBinary {
- self.type_().map(FixedSizeBinary::init_from_table)
- } else {
- None
- }
- }
-
- #[inline]
- #[allow(non_snake_case)]
- pub fn type_as_fixed_size_list(&self) -> Option<FixedSizeList<'a>> {
- if self.type_type() == Type::FixedSizeList {
- self.type_().map(FixedSizeList::init_from_table)
- } else {
- None
- }
- }
-
- #[inline]
- #[allow(non_snake_case)]
- pub fn type_as_map(&self) -> Option<Map<'a>> {
- if self.type_type() == Type::Map {
- self.type_().map(Map::init_from_table)
- } else {
- None
- }
- }
-
- #[inline]
- #[allow(non_snake_case)]
- pub fn type_as_duration(&self) -> Option<Duration<'a>> {
- if self.type_type() == Type::Duration {
- self.type_().map(Duration::init_from_table)
- } else {
- None
- }
- }
-
- #[inline]
- #[allow(non_snake_case)]
- pub fn type_as_large_binary(&self) -> Option<LargeBinary<'a>> {
- if self.type_type() == Type::LargeBinary {
- self.type_().map(LargeBinary::init_from_table)
- } else {
- None
- }
- }
-
- #[inline]
- #[allow(non_snake_case)]
- pub fn type_as_large_utf_8(&self) -> Option<LargeUtf8<'a>> {
- if self.type_type() == Type::LargeUtf8 {
- self.type_().map(LargeUtf8::init_from_table)
- } else {
- None
- }
- }
-
- #[inline]
- #[allow(non_snake_case)]
- pub fn type_as_large_list(&self) -> Option<LargeList<'a>> {
- if self.type_type() == Type::LargeList {
- self.type_().map(LargeList::init_from_table)
- } else {
- None
- }
- }
-}
-
-impl flatbuffers::Verifiable for Field<'_> {
- #[inline]
- fn run_verifier(
- v: &mut flatbuffers::Verifier,
- pos: usize,
- ) -> Result<(), flatbuffers::InvalidFlatbuffer> {
- use flatbuffers::Verifiable;
- v.visit_table(pos)?
- .visit_field::<flatbuffers::ForwardsUOffset<&str>>(&"name", Self::VT_NAME, false)?
- .visit_field::<bool>(&"nullable", Self::VT_NULLABLE, false)?
- .visit_union::<Type, _>(&"type_type", Self::VT_TYPE_TYPE, &"type_", Self::VT_TYPE_, false, |key, v, pos| {
- match key {
- Type::Null => v.verify_union_variant::<flatbuffers::ForwardsUOffset<Null>>("Type::Null", pos),
- Type::Int => v.verify_union_variant::<flatbuffers::ForwardsUOffset<Int>>("Type::Int", pos),
- Type::FloatingPoint => v.verify_union_variant::<flatbuffers::ForwardsUOffset<FloatingPoint>>("Type::FloatingPoint", pos),
- Type::Binary => v.verify_union_variant::<flatbuffers::ForwardsUOffset<Binary>>("Type::Binary", pos),
- Type::Utf8 => v.verify_union_variant::<flatbuffers::ForwardsUOffset<Utf8>>("Type::Utf8", pos),
- Type::Bool => v.verify_union_variant::<flatbuffers::ForwardsUOffset<Bool>>("Type::Bool", pos),
- Type::Decimal => v.verify_union_variant::<flatbuffers::ForwardsUOffset<Decimal>>("Type::Decimal", pos),
- Type::Date => v.verify_union_variant::<flatbuffers::ForwardsUOffset<Date>>("Type::Date", pos),
- Type::Time => v.verify_union_variant::<flatbuffers::ForwardsUOffset<Time>>("Type::Time", pos),
- Type::Timestamp => v.verify_union_variant::<flatbuffers::ForwardsUOffset<Timestamp>>("Type::Timestamp", pos),
- Type::Interval => v.verify_union_variant::<flatbuffers::ForwardsUOffset<Interval>>("Type::Interval", pos),
- Type::List => v.verify_union_variant::<flatbuffers::ForwardsUOffset<List>>("Type::List", pos),
- Type::Struct_ => v.verify_union_variant::<flatbuffers::ForwardsUOffset<Struct_>>("Type::Struct_", pos),
- Type::Union => v.verify_union_variant::<flatbuffers::ForwardsUOffset<Union>>("Type::Union", pos),
- Type::FixedSizeBinary => v.verify_union_variant::<flatbuffers::ForwardsUOffset<FixedSizeBinary>>("Type::FixedSizeBinary", pos),
- Type::FixedSizeList => v.verify_union_variant::<flatbuffers::ForwardsUOffset<FixedSizeList>>("Type::FixedSizeList", pos),
- Type::Map => v.verify_union_variant::<flatbuffers::ForwardsUOffset<Map>>("Type::Map", pos),
- Type::Duration => v.verify_union_variant::<flatbuffers::ForwardsUOffset<Duration>>("Type::Duration", pos),
- Type::LargeBinary => v.verify_union_variant::<flatbuffers::ForwardsUOffset<LargeBinary>>("Type::LargeBinary", pos),
- Type::LargeUtf8 => v.verify_union_variant::<flatbuffers::ForwardsUOffset<LargeUtf8>>("Type::LargeUtf8", pos),
- Type::LargeList => v.verify_union_variant::<flatbuffers::ForwardsUOffset<LargeList>>("Type::LargeList", pos),
- _ => Ok(()),
- }
- })?
- .visit_field::<flatbuffers::ForwardsUOffset<DictionaryEncoding>>(&"dictionary", Self::VT_DICTIONARY, false)?
- .visit_field::<flatbuffers::ForwardsUOffset<flatbuffers::Vector<'_, flatbuffers::ForwardsUOffset<Field>>>>(&"children", Self::VT_CHILDREN, false)?
- .visit_field::<flatbuffers::ForwardsUOffset<flatbuffers::Vector<'_, flatbuffers::ForwardsUOffset<KeyValue>>>>(&"custom_metadata", Self::VT_CUSTOM_METADATA, false)?
- .finish();
- Ok(())
- }
-}
-pub struct FieldArgs<'a> {
- pub name: Option<flatbuffers::WIPOffset<&'a str>>,
- pub nullable: bool,
- pub type_type: Type,
- pub type_: Option<flatbuffers::WIPOffset<flatbuffers::UnionWIPOffset>>,
- pub dictionary: Option<flatbuffers::WIPOffset<DictionaryEncoding<'a>>>,
- pub children: Option<
- flatbuffers::WIPOffset<
- flatbuffers::Vector<'a, flatbuffers::ForwardsUOffset<Field<'a>>>,
- >,
- >,
- pub custom_metadata: Option<
- flatbuffers::WIPOffset<
- flatbuffers::Vector<'a, flatbuffers::ForwardsUOffset<KeyValue<'a>>>,
- >,
- >,
-}
-impl<'a> Default for FieldArgs<'a> {
- #[inline]
- fn default() -> Self {
- FieldArgs {
- name: None,
- nullable: false,
- type_type: Type::NONE,
- type_: None,
- dictionary: None,
- children: None,
- custom_metadata: None,
- }
- }
-}
-pub struct FieldBuilder<'a: 'b, 'b> {
- fbb_: &'b mut flatbuffers::FlatBufferBuilder<'a>,
- start_: flatbuffers::WIPOffset<flatbuffers::TableUnfinishedWIPOffset>,
-}
-impl<'a: 'b, 'b> FieldBuilder<'a, 'b> {
- #[inline]
- pub fn add_name(&mut self, name: flatbuffers::WIPOffset<&'b str>) {
- self.fbb_
- .push_slot_always::<flatbuffers::WIPOffset<_>>(Field::VT_NAME, name);
- }
- #[inline]
- pub fn add_nullable(&mut self, nullable: bool) {
- self.fbb_
- .push_slot::<bool>(Field::VT_NULLABLE, nullable, false);
- }
- #[inline]
- pub fn add_type_type(&mut self, type_type: Type) {
- self.fbb_
- .push_slot::<Type>(Field::VT_TYPE_TYPE, type_type, Type::NONE);
- }
- #[inline]
- pub fn add_type_(
- &mut self,
- type_: flatbuffers::WIPOffset<flatbuffers::UnionWIPOffset>,
- ) {
- self.fbb_
- .push_slot_always::<flatbuffers::WIPOffset<_>>(Field::VT_TYPE_, type_);
- }
- #[inline]
- pub fn add_dictionary(
- &mut self,
- dictionary: flatbuffers::WIPOffset<DictionaryEncoding<'b>>,
- ) {
- self.fbb_
- .push_slot_always::<flatbuffers::WIPOffset<DictionaryEncoding>>(
- Field::VT_DICTIONARY,
- dictionary,
- );
- }
- #[inline]
- pub fn add_children(
- &mut self,
- children: flatbuffers::WIPOffset<
- flatbuffers::Vector<'b, flatbuffers::ForwardsUOffset<Field<'b>>>,
- >,
- ) {
- self.fbb_
- .push_slot_always::<flatbuffers::WIPOffset<_>>(Field::VT_CHILDREN, children);
- }
- #[inline]
- pub fn add_custom_metadata(
- &mut self,
- custom_metadata: flatbuffers::WIPOffset<
- flatbuffers::Vector<'b, flatbuffers::ForwardsUOffset<KeyValue<'b>>>,
- >,
- ) {
- self.fbb_.push_slot_always::<flatbuffers::WIPOffset<_>>(
- Field::VT_CUSTOM_METADATA,
- custom_metadata,
- );
- }
- #[inline]
- pub fn new(_fbb: &'b mut flatbuffers::FlatBufferBuilder<'a>) -> FieldBuilder<'a, 'b> {
- let start = _fbb.start_table();
- FieldBuilder {
- fbb_: _fbb,
- start_: start,
- }
- }
- #[inline]
- pub fn finish(self) -> flatbuffers::WIPOffset<Field<'a>> {
- let o = self.fbb_.end_table(self.start_);
- flatbuffers::WIPOffset::new(o.value())
- }
-}
-
-impl std::fmt::Debug for Field<'_> {
- fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
- let mut ds = f.debug_struct("Field");
- ds.field("name", &self.name());
- ds.field("nullable", &self.nullable());
- ds.field("type_type", &self.type_type());
- match self.type_type() {
- Type::Null => {
- if let Some(x) = self.type_as_null() {
- ds.field("type_", &x)
- } else {
- ds.field(
- "type_",
- &"InvalidFlatbuffer: Union discriminant does not match value.",
- )
- }
- }
- Type::Int => {
- if let Some(x) = self.type_as_int() {
- ds.field("type_", &x)
- } else {
- ds.field(
- "type_",
- &"InvalidFlatbuffer: Union discriminant does not match value.",
- )
- }
- }
- Type::FloatingPoint => {
- if let Some(x) = self.type_as_floating_point() {
- ds.field("type_", &x)
- } else {
- ds.field(
- "type_",
- &"InvalidFlatbuffer: Union discriminant does not match value.",
- )
- }
- }
- Type::Binary => {
- if let Some(x) = self.type_as_binary() {
- ds.field("type_", &x)
- } else {
- ds.field(
- "type_",
- &"InvalidFlatbuffer: Union discriminant does not match value.",
- )
- }
- }
- Type::Utf8 => {
- if let Some(x) = self.type_as_utf_8() {
- ds.field("type_", &x)
- } else {
- ds.field(
- "type_",
- &"InvalidFlatbuffer: Union discriminant does not match value.",
- )
- }
- }
- Type::Bool => {
- if let Some(x) = self.type_as_bool() {
- ds.field("type_", &x)
- } else {
- ds.field(
- "type_",
- &"InvalidFlatbuffer: Union discriminant does not match value.",
- )
- }
- }
- Type::Decimal => {
- if let Some(x) = self.type_as_decimal() {
- ds.field("type_", &x)
- } else {
- ds.field(
- "type_",
- &"InvalidFlatbuffer: Union discriminant does not match value.",
- )
- }
- }
- Type::Date => {
- if let Some(x) = self.type_as_date() {
- ds.field("type_", &x)
- } else {
- ds.field(
- "type_",
- &"InvalidFlatbuffer: Union discriminant does not match value.",
- )
- }
- }
- Type::Time => {
- if let Some(x) = self.type_as_time() {
- ds.field("type_", &x)
- } else {
- ds.field(
- "type_",
- &"InvalidFlatbuffer: Union discriminant does not match value.",
- )
- }
- }
- Type::Timestamp => {
- if let Some(x) = self.type_as_timestamp() {
- ds.field("type_", &x)
- } else {
- ds.field(
- "type_",
- &"InvalidFlatbuffer: Union discriminant does not match value.",
- )
- }
- }
- Type::Interval => {
- if let Some(x) = self.type_as_interval() {
- ds.field("type_", &x)
- } else {
- ds.field(
- "type_",
- &"InvalidFlatbuffer: Union discriminant does not match value.",
- )
- }
- }
- Type::List => {
- if let Some(x) = self.type_as_list() {
- ds.field("type_", &x)
- } else {
- ds.field(
- "type_",
- &"InvalidFlatbuffer: Union discriminant does not match value.",
- )
- }
- }
- Type::Struct_ => {
- if let Some(x) = self.type_as_struct_() {
- ds.field("type_", &x)
- } else {
- ds.field(
- "type_",
- &"InvalidFlatbuffer: Union discriminant does not match value.",
- )
- }
- }
- Type::Union => {
- if let Some(x) = self.type_as_union() {
- ds.field("type_", &x)
- } else {
- ds.field(
- "type_",
- &"InvalidFlatbuffer: Union discriminant does not match value.",
- )
- }
- }
- Type::FixedSizeBinary => {
- if let Some(x) = self.type_as_fixed_size_binary() {
- ds.field("type_", &x)
- } else {
- ds.field(
- "type_",
- &"InvalidFlatbuffer: Union discriminant does not match value.",
- )
- }
- }
- Type::FixedSizeList => {
- if let Some(x) = self.type_as_fixed_size_list() {
- ds.field("type_", &x)
- } else {
- ds.field(
- "type_",
- &"InvalidFlatbuffer: Union discriminant does not match value.",
- )
- }
- }
- Type::Map => {
- if let Some(x) = self.type_as_map() {
- ds.field("type_", &x)
- } else {
- ds.field(
- "type_",
- &"InvalidFlatbuffer: Union discriminant does not match value.",
- )
- }
- }
- Type::Duration => {
- if let Some(x) = self.type_as_duration() {
- ds.field("type_", &x)
- } else {
- ds.field(
- "type_",
- &"InvalidFlatbuffer: Union discriminant does not match value.",
- )
- }
- }
- Type::LargeBinary => {
- if let Some(x) = self.type_as_large_binary() {
- ds.field("type_", &x)
- } else {
- ds.field(
- "type_",
- &"InvalidFlatbuffer: Union discriminant does not match value.",
- )
- }
- }
- Type::LargeUtf8 => {
- if let Some(x) = self.type_as_large_utf_8() {
- ds.field("type_", &x)
- } else {
- ds.field(
- "type_",
- &"InvalidFlatbuffer: Union discriminant does not match value.",
- )
- }
- }
- Type::LargeList => {
- if let Some(x) = self.type_as_large_list() {
- ds.field("type_", &x)
- } else {
- ds.field(
- "type_",
- &"InvalidFlatbuffer: Union discriminant does not match value.",
- )
- }
- }
- _ => {
- let x: Option<()> = None;
- ds.field("type_", &x)
- }
- };
- ds.field("dictionary", &self.dictionary());
- ds.field("children", &self.children());
- ds.field("custom_metadata", &self.custom_metadata());
- ds.finish()
- }
-}
-pub enum SchemaOffset {}
-#[derive(Copy, Clone, PartialEq)]
-
-/// ----------------------------------------------------------------------
-/// A Schema describes the columns in a row batch
-pub struct Schema<'a> {
- pub _tab: flatbuffers::Table<'a>,
-}
-
-impl<'a> flatbuffers::Follow<'a> for Schema<'a> {
- type Inner = Schema<'a>;
- #[inline]
- fn follow(buf: &'a [u8], loc: usize) -> Self::Inner {
- Self {
- _tab: flatbuffers::Table { buf, loc },
- }
- }
-}
-
-impl<'a> Schema<'a> {
- #[inline]
- pub fn init_from_table(table: flatbuffers::Table<'a>) -> Self {
- Schema { _tab: table }
- }
- #[allow(unused_mut)]
- pub fn create<'bldr: 'args, 'args: 'mut_bldr, 'mut_bldr>(
- _fbb: &'mut_bldr mut flatbuffers::FlatBufferBuilder<'bldr>,
- args: &'args SchemaArgs<'args>,
- ) -> flatbuffers::WIPOffset<Schema<'bldr>> {
- let mut builder = SchemaBuilder::new(_fbb);
- if let Some(x) = args.features {
- builder.add_features(x);
- }
- if let Some(x) = args.custom_metadata {
- builder.add_custom_metadata(x);
- }
- if let Some(x) = args.fields {
- builder.add_fields(x);
- }
- builder.add_endianness(args.endianness);
- builder.finish()
- }
-
- pub const VT_ENDIANNESS: flatbuffers::VOffsetT = 4;
- pub const VT_FIELDS: flatbuffers::VOffsetT = 6;
- pub const VT_CUSTOM_METADATA: flatbuffers::VOffsetT = 8;
- pub const VT_FEATURES: flatbuffers::VOffsetT = 10;
-
- /// 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
- #[inline]
- pub fn endianness(&self) -> Endianness {
- self._tab
- .get::<Endianness>(Schema::VT_ENDIANNESS, Some(Endianness::Little))
- .unwrap()
- }
- #[inline]
- pub fn fields(
- &self,
- ) -> Option<flatbuffers::Vector<'a, flatbuffers::ForwardsUOffset<Field<'a>>>> {
- self._tab.get::<flatbuffers::ForwardsUOffset<
- flatbuffers::Vector<'a, flatbuffers::ForwardsUOffset<Field>>,
- >>(Schema::VT_FIELDS, None)
- }
- #[inline]
- pub fn custom_metadata(
- &self,
- ) -> Option<flatbuffers::Vector<'a, flatbuffers::ForwardsUOffset<KeyValue<'a>>>> {
- self._tab.get::<flatbuffers::ForwardsUOffset<
- flatbuffers::Vector<'a, flatbuffers::ForwardsUOffset<KeyValue>>,
- >>(Schema::VT_CUSTOM_METADATA, None)
- }
- /// Features used in the stream/file.
- #[inline]
- pub fn features(&self) -> Option<flatbuffers::Vector<'a, Feature>> {
- self._tab
- .get::<flatbuffers::ForwardsUOffset<flatbuffers::Vector<'a, Feature>>>(
- Schema::VT_FEATURES,
- None,
- )
- }
-}
-
-impl flatbuffers::Verifiable for Schema<'_> {
- #[inline]
- fn run_verifier(
- v: &mut flatbuffers::Verifier,
- pos: usize,
- ) -> Result<(), flatbuffers::InvalidFlatbuffer> {
- use flatbuffers::Verifiable;
- v.visit_table(pos)?
- .visit_field::<Endianness>(&"endianness", Self::VT_ENDIANNESS, false)?
- .visit_field::<flatbuffers::ForwardsUOffset<flatbuffers::Vector<'_, flatbuffers::ForwardsUOffset<Field>>>>(&"fields", Self::VT_FIELDS, false)?
- .visit_field::<flatbuffers::ForwardsUOffset<flatbuffers::Vector<'_, flatbuffers::ForwardsUOffset<KeyValue>>>>(&"custom_metadata", Self::VT_CUSTOM_METADATA, false)?
- .visit_field::<flatbuffers::ForwardsUOffset<flatbuffers::Vector<'_, Feature>>>(&"features", Self::VT_FEATURES, false)?
- .finish();
- Ok(())
- }
-}
-pub struct SchemaArgs<'a> {
- pub endianness: Endianness,
- pub fields: Option<
- flatbuffers::WIPOffset<
- flatbuffers::Vector<'a, flatbuffers::ForwardsUOffset<Field<'a>>>,
- >,
- >,
- pub custom_metadata: Option<
- flatbuffers::WIPOffset<
- flatbuffers::Vector<'a, flatbuffers::ForwardsUOffset<KeyValue<'a>>>,
- >,
- >,
- pub features: Option<flatbuffers::WIPOffset<flatbuffers::Vector<'a, Feature>>>,
-}
-impl<'a> Default for SchemaArgs<'a> {
- #[inline]
- fn default() -> Self {
- SchemaArgs {
- endianness: Endianness::Little,
- fields: None,
- custom_metadata: None,
- features: None,
- }
- }
-}
-pub struct SchemaBuilder<'a: 'b, 'b> {
- fbb_: &'b mut flatbuffers::FlatBufferBuilder<'a>,
- start_: flatbuffers::WIPOffset<flatbuffers::TableUnfinishedWIPOffset>,
-}
-impl<'a: 'b, 'b> SchemaBuilder<'a, 'b> {
- #[inline]
- pub fn add_endianness(&mut self, endianness: Endianness) {
- self.fbb_.push_slot::<Endianness>(
- Schema::VT_ENDIANNESS,
- endianness,
- Endianness::Little,
- );
- }
- #[inline]
- pub fn add_fields(
- &mut self,
- fields: flatbuffers::WIPOffset<
- flatbuffers::Vector<'b, flatbuffers::ForwardsUOffset<Field<'b>>>,
- >,
- ) {
- self.fbb_
- .push_slot_always::<flatbuffers::WIPOffset<_>>(Schema::VT_FIELDS, fields);
- }
- #[inline]
- pub fn add_custom_metadata(
- &mut self,
- custom_metadata: flatbuffers::WIPOffset<
- flatbuffers::Vector<'b, flatbuffers::ForwardsUOffset<KeyValue<'b>>>,
- >,
- ) {
- self.fbb_.push_slot_always::<flatbuffers::WIPOffset<_>>(
- Schema::VT_CUSTOM_METADATA,
- custom_metadata,
- );
- }
- #[inline]
- pub fn add_features(
- &mut self,
- features: flatbuffers::WIPOffset<flatbuffers::Vector<'b, Feature>>,
- ) {
- self.fbb_
- .push_slot_always::<flatbuffers::WIPOffset<_>>(Schema::VT_FEATURES, features);
- }
- #[inline]
- pub fn new(
- _fbb: &'b mut flatbuffers::FlatBufferBuilder<'a>,
- ) -> SchemaBuilder<'a, 'b> {
- let start = _fbb.start_table();
- SchemaBuilder {
- fbb_: _fbb,
- start_: start,
- }
- }
- #[inline]
- pub fn finish(self) -> flatbuffers::WIPOffset<Schema<'a>> {
- let o = self.fbb_.end_table(self.start_);
- flatbuffers::WIPOffset::new(o.value())
- }
-}
-
-impl std::fmt::Debug for Schema<'_> {
- fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
- let mut ds = f.debug_struct("Schema");
- ds.field("endianness", &self.endianness());
- ds.field("fields", &self.fields());
- ds.field("custom_metadata", &self.custom_metadata());
- ds.field("features", &self.features());
- ds.finish()
- }
-}
-#[inline]
-#[deprecated(since = "2.0.0", note = "Deprecated in favor of `root_as...` methods.")]
-pub fn get_root_as_schema<'a>(buf: &'a [u8]) -> Schema<'a> {
- unsafe { flatbuffers::root_unchecked::<Schema<'a>>(buf) }
-}
-
-#[inline]
-#[deprecated(since = "2.0.0", note = "Deprecated in favor of `root_as...` methods.")]
-pub fn get_size_prefixed_root_as_schema<'a>(buf: &'a [u8]) -> Schema<'a> {
- unsafe { flatbuffers::size_prefixed_root_unchecked::<Schema<'a>>(buf) }
-}
-
-#[inline]
-/// Verifies that a buffer of bytes contains a `Schema`
-/// and returns it.
-/// Note that verification is still experimental and may not
-/// catch every error, or be maximally performant. For the
-/// previous, unchecked, behavior use
-/// `root_as_schema_unchecked`.
-pub fn root_as_schema(buf: &[u8]) -> Result<Schema, flatbuffers::InvalidFlatbuffer> {
- flatbuffers::root::<Schema>(buf)
-}
-#[inline]
-/// Verifies that a buffer of bytes contains a size prefixed
-/// `Schema` and returns it.
-/// Note that verification is still experimental and may not
-/// catch every error, or be maximally performant. For the
-/// previous, unchecked, behavior use
-/// `size_prefixed_root_as_schema_unchecked`.
-pub fn size_prefixed_root_as_schema(
- buf: &[u8],
-) -> Result<Schema, flatbuffers::InvalidFlatbuffer> {
- flatbuffers::size_prefixed_root::<Schema>(buf)
-}
-#[inline]
-/// Verifies, with the given options, that a buffer of bytes
-/// contains a `Schema` and returns it.
-/// Note that verification is still experimental and may not
-/// catch every error, or be maximally performant. For the
-/// previous, unchecked, behavior use
-/// `root_as_schema_unchecked`.
-pub fn root_as_schema_with_opts<'b, 'o>(
- opts: &'o flatbuffers::VerifierOptions,
- buf: &'b [u8],
-) -> Result<Schema<'b>, flatbuffers::InvalidFlatbuffer> {
- flatbuffers::root_with_opts::<Schema<'b>>(opts, buf)
-}
-#[inline]
-/// Verifies, with the given verifier options, that a buffer of
-/// bytes contains a size prefixed `Schema` and returns
-/// it. Note that verification is still experimental and may not
-/// catch every error, or be maximally performant. For the
-/// previous, unchecked, behavior use
-/// `root_as_schema_unchecked`.
-pub fn size_prefixed_root_as_schema_with_opts<'b, 'o>(
- opts: &'o flatbuffers::VerifierOptions,
- buf: &'b [u8],
-) -> Result<Schema<'b>, flatbuffers::InvalidFlatbuffer> {
- flatbuffers::size_prefixed_root_with_opts::<Schema<'b>>(opts, buf)
-}
-#[inline]
-/// Assumes, without verification, that a buffer of bytes contains a Schema and returns it.
-/// # Safety
-/// Callers must trust the given bytes do indeed contain a valid `Schema`.
-pub unsafe fn root_as_schema_unchecked(buf: &[u8]) -> Schema {
- flatbuffers::root_unchecked::<Schema>(buf)
-}
-#[inline]
-/// Assumes, without verification, that a buffer of bytes contains a size prefixed Schema and returns it.
-/// # Safety
-/// Callers must trust the given bytes do indeed contain a valid size prefixed `Schema`.
-pub unsafe fn size_prefixed_root_as_schema_unchecked(buf: &[u8]) -> Schema {
- flatbuffers::size_prefixed_root_unchecked::<Schema>(buf)
-}
-#[inline]
-pub fn finish_schema_buffer<'a, 'b>(
- fbb: &'b mut flatbuffers::FlatBufferBuilder<'a>,
- root: flatbuffers::WIPOffset<Schema<'a>>,
-) {
- fbb.finish(root, None);
-}
-
-#[inline]
-pub fn finish_size_prefixed_schema_buffer<'a, 'b>(
- fbb: &'b mut flatbuffers::FlatBufferBuilder<'a>,
- root: flatbuffers::WIPOffset<Schema<'a>>,
-) {
- fbb.finish_size_prefixed(root, None);
-}
diff --git a/arrow/src/ipc/gen/SparseTensor.rs b/arrow/src/ipc/gen/SparseTensor.rs
deleted file mode 100644
index 5d12d4e..0000000
--- a/arrow/src/ipc/gen/SparseTensor.rs
+++ /dev/null
@@ -1,1902 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-#![allow(dead_code)]
-#![allow(unused_imports)]
-
-use crate::ipc::gen::Schema::*;
-use crate::ipc::gen::Tensor::*;
-use flatbuffers::EndianScalar;
-use std::{cmp::Ordering, mem};
-// automatically generated by the FlatBuffers compiler, do not modify
-
-#[deprecated(
- since = "2.0.0",
- note = "Use associated constants instead. This will no longer be generated in 2021."
-)]
-pub const ENUM_MIN_SPARSE_MATRIX_COMPRESSED_AXIS: i16 = 0;
-#[deprecated(
- since = "2.0.0",
- note = "Use associated constants instead. This will no longer be generated in 2021."
-)]
-pub const ENUM_MAX_SPARSE_MATRIX_COMPRESSED_AXIS: i16 = 1;
-#[deprecated(
- since = "2.0.0",
- note = "Use associated constants instead. This will no longer be generated in 2021."
-)]
-#[allow(non_camel_case_types)]
-pub const ENUM_VALUES_SPARSE_MATRIX_COMPRESSED_AXIS: [SparseMatrixCompressedAxis; 2] = [
- SparseMatrixCompressedAxis::Row,
- SparseMatrixCompressedAxis::Column,
-];
-
-#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
-#[repr(transparent)]
-pub struct SparseMatrixCompressedAxis(pub i16);
-#[allow(non_upper_case_globals)]
-impl SparseMatrixCompressedAxis {
- pub const Row: Self = Self(0);
- pub const Column: Self = Self(1);
-
- pub const ENUM_MIN: i16 = 0;
- pub const ENUM_MAX: i16 = 1;
- pub const ENUM_VALUES: &'static [Self] = &[Self::Row, Self::Column];
- /// Returns the variant's name or "" if unknown.
- pub fn variant_name(self) -> Option<&'static str> {
- match self {
- Self::Row => Some("Row"),
- Self::Column => Some("Column"),
- _ => None,
- }
- }
-}
-impl std::fmt::Debug for SparseMatrixCompressedAxis {
- fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
- if let Some(name) = self.variant_name() {
- f.write_str(name)
- } else {
- f.write_fmt(format_args!("<UNKNOWN {:?}>", self.0))
- }
- }
-}
-impl<'a> flatbuffers::Follow<'a> for SparseMatrixCompressedAxis {
- type Inner = Self;
- #[inline]
- fn follow(buf: &'a [u8], loc: usize) -> Self::Inner {
- let b = flatbuffers::read_scalar_at::<i16>(buf, loc);
- Self(b)
- }
-}
-
-impl flatbuffers::Push for SparseMatrixCompressedAxis {
- type Output = SparseMatrixCompressedAxis;
- #[inline]
- fn push(&self, dst: &mut [u8], _rest: &[u8]) {
- flatbuffers::emplace_scalar::<i16>(dst, self.0);
- }
-}
-
-impl flatbuffers::EndianScalar for SparseMatrixCompressedAxis {
- #[inline]
- fn to_little_endian(self) -> Self {
- let b = i16::to_le(self.0);
- Self(b)
- }
- #[inline]
- fn from_little_endian(self) -> Self {
- let b = i16::from_le(self.0);
- Self(b)
- }
-}
-
-impl<'a> flatbuffers::Verifiable for SparseMatrixCompressedAxis {
- #[inline]
- fn run_verifier(
- v: &mut flatbuffers::Verifier,
- pos: usize,
- ) -> Result<(), flatbuffers::InvalidFlatbuffer> {
- use flatbuffers::Verifiable;
- i16::run_verifier(v, pos)
- }
-}
-
-impl flatbuffers::SimpleToVerifyInSlice for SparseMatrixCompressedAxis {}
-#[deprecated(
- since = "2.0.0",
- note = "Use associated constants instead. This will no longer be generated in 2021."
-)]
-pub const ENUM_MIN_SPARSE_TENSOR_INDEX: u8 = 0;
-#[deprecated(
- since = "2.0.0",
- note = "Use associated constants instead. This will no longer be generated in 2021."
-)]
-pub const ENUM_MAX_SPARSE_TENSOR_INDEX: u8 = 3;
-#[deprecated(
- since = "2.0.0",
- note = "Use associated constants instead. This will no longer be generated in 2021."
-)]
-#[allow(non_camel_case_types)]
-pub const ENUM_VALUES_SPARSE_TENSOR_INDEX: [SparseTensorIndex; 4] = [
- SparseTensorIndex::NONE,
- SparseTensorIndex::SparseTensorIndexCOO,
- SparseTensorIndex::SparseMatrixIndexCSX,
- SparseTensorIndex::SparseTensorIndexCSF,
-];
-
-#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
-#[repr(transparent)]
-pub struct SparseTensorIndex(pub u8);
-#[allow(non_upper_case_globals)]
-impl SparseTensorIndex {
- pub const NONE: Self = Self(0);
- pub const SparseTensorIndexCOO: Self = Self(1);
- pub const SparseMatrixIndexCSX: Self = Self(2);
- pub const SparseTensorIndexCSF: Self = Self(3);
-
- pub const ENUM_MIN: u8 = 0;
- pub const ENUM_MAX: u8 = 3;
- pub const ENUM_VALUES: &'static [Self] = &[
- Self::NONE,
- Self::SparseTensorIndexCOO,
- Self::SparseMatrixIndexCSX,
- Self::SparseTensorIndexCSF,
- ];
- /// Returns the variant's name or "" if unknown.
- pub fn variant_name(self) -> Option<&'static str> {
- match self {
- Self::NONE => Some("NONE"),
- Self::SparseTensorIndexCOO => Some("SparseTensorIndexCOO"),
- Self::SparseMatrixIndexCSX => Some("SparseMatrixIndexCSX"),
- Self::SparseTensorIndexCSF => Some("SparseTensorIndexCSF"),
- _ => None,
- }
- }
-}
-impl std::fmt::Debug for SparseTensorIndex {
- fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
- if let Some(name) = self.variant_name() {
- f.write_str(name)
- } else {
- f.write_fmt(format_args!("<UNKNOWN {:?}>", self.0))
- }
- }
-}
-pub struct SparseTensorIndexUnionTableOffset {}
-impl<'a> flatbuffers::Follow<'a> for SparseTensorIndex {
- type Inner = Self;
- #[inline]
- fn follow(buf: &'a [u8], loc: usize) -> Self::Inner {
- let b = flatbuffers::read_scalar_at::<u8>(buf, loc);
- Self(b)
- }
-}
-
-impl flatbuffers::Push for SparseTensorIndex {
- type Output = SparseTensorIndex;
- #[inline]
- fn push(&self, dst: &mut [u8], _rest: &[u8]) {
- flatbuffers::emplace_scalar::<u8>(dst, self.0);
- }
-}
-
-impl flatbuffers::EndianScalar for SparseTensorIndex {
- #[inline]
- fn to_little_endian(self) -> Self {
- let b = u8::to_le(self.0);
- Self(b)
- }
- #[inline]
- fn from_little_endian(self) -> Self {
- let b = u8::from_le(self.0);
- Self(b)
- }
-}
-
-impl<'a> flatbuffers::Verifiable for SparseTensorIndex {
- #[inline]
- fn run_verifier(
- v: &mut flatbuffers::Verifier,
- pos: usize,
- ) -> Result<(), flatbuffers::InvalidFlatbuffer> {
- use flatbuffers::Verifiable;
- u8::run_verifier(v, pos)
- }
-}
-
-impl flatbuffers::SimpleToVerifyInSlice for SparseTensorIndex {}
-pub enum SparseTensorIndexCOOOffset {}
-#[derive(Copy, Clone, PartialEq)]
-
-/// ----------------------------------------------------------------------
-/// 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.
-pub struct SparseTensorIndexCOO<'a> {
- pub _tab: flatbuffers::Table<'a>,
-}
-
-impl<'a> flatbuffers::Follow<'a> for SparseTensorIndexCOO<'a> {
- type Inner = SparseTensorIndexCOO<'a>;
- #[inline]
- fn follow(buf: &'a [u8], loc: usize) -> Self::Inner {
- Self {
- _tab: flatbuffers::Table { buf, loc },
- }
- }
-}
-
-impl<'a> SparseTensorIndexCOO<'a> {
- #[inline]
- pub fn init_from_table(table: flatbuffers::Table<'a>) -> Self {
- SparseTensorIndexCOO { _tab: table }
- }
- #[allow(unused_mut)]
- pub fn create<'bldr: 'args, 'args: 'mut_bldr, 'mut_bldr>(
- _fbb: &'mut_bldr mut flatbuffers::FlatBufferBuilder<'bldr>,
- args: &'args SparseTensorIndexCOOArgs<'args>,
- ) -> flatbuffers::WIPOffset<SparseTensorIndexCOO<'bldr>> {
- let mut builder = SparseTensorIndexCOOBuilder::new(_fbb);
- if let Some(x) = args.indicesBuffer {
- builder.add_indicesBuffer(x);
- }
- if let Some(x) = args.indicesStrides {
- builder.add_indicesStrides(x);
- }
- if let Some(x) = args.indicesType {
- builder.add_indicesType(x);
- }
- builder.add_isCanonical(args.isCanonical);
- builder.finish()
- }
-
- pub const VT_INDICESTYPE: flatbuffers::VOffsetT = 4;
- pub const VT_INDICESSTRIDES: flatbuffers::VOffsetT = 6;
- pub const VT_INDICESBUFFER: flatbuffers::VOffsetT = 8;
- pub const VT_ISCANONICAL: flatbuffers::VOffsetT = 10;
-
- /// The type of values in indicesBuffer
- #[inline]
- pub fn indicesType(&self) -> Int<'a> {
- self._tab
- .get::<flatbuffers::ForwardsUOffset<Int>>(
- SparseTensorIndexCOO::VT_INDICESTYPE,
- None,
- )
- .unwrap()
- }
- /// Non-negative byte offsets to advance one value cell along each dimension
- /// If omitted, default to row-major order (C-like).
- #[inline]
- pub fn indicesStrides(&self) -> Option<flatbuffers::Vector<'a, i64>> {
- self._tab
- .get::<flatbuffers::ForwardsUOffset<flatbuffers::Vector<'a, i64>>>(
- SparseTensorIndexCOO::VT_INDICESSTRIDES,
- None,
- )
- }
- /// The location and size of the indices matrix's data
- #[inline]
- pub fn indicesBuffer(&self) -> &'a Buffer {
- self._tab
- .get::<Buffer>(SparseTensorIndexCOO::VT_INDICESBUFFER, None)
- .unwrap()
- }
- /// 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).
- #[inline]
- pub fn isCanonical(&self) -> bool {
- self._tab
- .get::<bool>(SparseTensorIndexCOO::VT_ISCANONICAL, Some(false))
- .unwrap()
- }
-}
-
-impl flatbuffers::Verifiable for SparseTensorIndexCOO<'_> {
- #[inline]
- fn run_verifier(
- v: &mut flatbuffers::Verifier,
- pos: usize,
- ) -> Result<(), flatbuffers::InvalidFlatbuffer> {
- use flatbuffers::Verifiable;
- v.visit_table(pos)?
- .visit_field::<flatbuffers::ForwardsUOffset<Int>>(
- &"indicesType",
- Self::VT_INDICESTYPE,
- true,
- )?
- .visit_field::<flatbuffers::ForwardsUOffset<flatbuffers::Vector<'_, i64>>>(
- &"indicesStrides",
- Self::VT_INDICESSTRIDES,
- false,
- )?
- .visit_field::<Buffer>(&"indicesBuffer", Self::VT_INDICESBUFFER, true)?
- .visit_field::<bool>(&"isCanonical", Self::VT_ISCANONICAL, false)?
- .finish();
- Ok(())
- }
-}
-pub struct SparseTensorIndexCOOArgs<'a> {
- pub indicesType: Option<flatbuffers::WIPOffset<Int<'a>>>,
- pub indicesStrides: Option<flatbuffers::WIPOffset<flatbuffers::Vector<'a, i64>>>,
- pub indicesBuffer: Option<&'a Buffer>,
- pub isCanonical: bool,
-}
-impl<'a> Default for SparseTensorIndexCOOArgs<'a> {
- #[inline]
- fn default() -> Self {
- SparseTensorIndexCOOArgs {
- indicesType: None, // required field
- indicesStrides: None,
- indicesBuffer: None, // required field
- isCanonical: false,
- }
- }
-}
-pub struct SparseTensorIndexCOOBuilder<'a: 'b, 'b> {
- fbb_: &'b mut flatbuffers::FlatBufferBuilder<'a>,
- start_: flatbuffers::WIPOffset<flatbuffers::TableUnfinishedWIPOffset>,
-}
-impl<'a: 'b, 'b> SparseTensorIndexCOOBuilder<'a, 'b> {
- #[inline]
- pub fn add_indicesType(&mut self, indicesType: flatbuffers::WIPOffset<Int<'b>>) {
- self.fbb_.push_slot_always::<flatbuffers::WIPOffset<Int>>(
- SparseTensorIndexCOO::VT_INDICESTYPE,
- indicesType,
- );
- }
- #[inline]
- pub fn add_indicesStrides(
- &mut self,
- indicesStrides: flatbuffers::WIPOffset<flatbuffers::Vector<'b, i64>>,
- ) {
- self.fbb_.push_slot_always::<flatbuffers::WIPOffset<_>>(
- SparseTensorIndexCOO::VT_INDICESSTRIDES,
- indicesStrides,
- );
- }
- #[inline]
- pub fn add_indicesBuffer(&mut self, indicesBuffer: &Buffer) {
- self.fbb_.push_slot_always::<&Buffer>(
- SparseTensorIndexCOO::VT_INDICESBUFFER,
- indicesBuffer,
- );
- }
- #[inline]
- pub fn add_isCanonical(&mut self, isCanonical: bool) {
- self.fbb_.push_slot::<bool>(
- SparseTensorIndexCOO::VT_ISCANONICAL,
- isCanonical,
- false,
- );
- }
- #[inline]
- pub fn new(
- _fbb: &'b mut flatbuffers::FlatBufferBuilder<'a>,
- ) -> SparseTensorIndexCOOBuilder<'a, 'b> {
- let start = _fbb.start_table();
- SparseTensorIndexCOOBuilder {
- fbb_: _fbb,
- start_: start,
- }
- }
- #[inline]
- pub fn finish(self) -> flatbuffers::WIPOffset<SparseTensorIndexCOO<'a>> {
- let o = self.fbb_.end_table(self.start_);
- self.fbb_
- .required(o, SparseTensorIndexCOO::VT_INDICESTYPE, "indicesType");
- self.fbb_
- .required(o, SparseTensorIndexCOO::VT_INDICESBUFFER, "indicesBuffer");
- flatbuffers::WIPOffset::new(o.value())
- }
-}
-
-impl std::fmt::Debug for SparseTensorIndexCOO<'_> {
- fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
- let mut ds = f.debug_struct("SparseTensorIndexCOO");
- ds.field("indicesType", &self.indicesType());
- ds.field("indicesStrides", &self.indicesStrides());
- ds.field("indicesBuffer", &self.indicesBuffer());
- ds.field("isCanonical", &self.isCanonical());
- ds.finish()
- }
-}
-pub enum SparseMatrixIndexCSXOffset {}
-#[derive(Copy, Clone, PartialEq)]
-
-/// Compressed Sparse format, that is matrix-specific.
-pub struct SparseMatrixIndexCSX<'a> {
- pub _tab: flatbuffers::Table<'a>,
-}
-
-impl<'a> flatbuffers::Follow<'a> for SparseMatrixIndexCSX<'a> {
- type Inner = SparseMatrixIndexCSX<'a>;
- #[inline]
- fn follow(buf: &'a [u8], loc: usize) -> Self::Inner {
- Self {
- _tab: flatbuffers::Table { buf, loc },
- }
- }
-}
-
-impl<'a> SparseMatrixIndexCSX<'a> {
- #[inline]
- pub fn init_from_table(table: flatbuffers::Table<'a>) -> Self {
- SparseMatrixIndexCSX { _tab: table }
- }
- #[allow(unused_mut)]
- pub fn create<'bldr: 'args, 'args: 'mut_bldr, 'mut_bldr>(
- _fbb: &'mut_bldr mut flatbuffers::FlatBufferBuilder<'bldr>,
- args: &'args SparseMatrixIndexCSXArgs<'args>,
- ) -> flatbuffers::WIPOffset<SparseMatrixIndexCSX<'bldr>> {
- let mut builder = SparseMatrixIndexCSXBuilder::new(_fbb);
- if let Some(x) = args.indicesBuffer {
- builder.add_indicesBuffer(x);
- }
- if let Some(x) = args.indicesType {
- builder.add_indicesType(x);
- }
- if let Some(x) = args.indptrBuffer {
- builder.add_indptrBuffer(x);
- }
- if let Some(x) = args.indptrType {
- builder.add_indptrType(x);
- }
- builder.add_compressedAxis(args.compressedAxis);
- builder.finish()
- }
-
- pub const VT_COMPRESSEDAXIS: flatbuffers::VOffsetT = 4;
- pub const VT_INDPTRTYPE: flatbuffers::VOffsetT = 6;
- pub const VT_INDPTRBUFFER: flatbuffers::VOffsetT = 8;
- pub const VT_INDICESTYPE: flatbuffers::VOffsetT = 10;
- pub const VT_INDICESBUFFER: flatbuffers::VOffsetT = 12;
-
- /// Which axis, row or column, is compressed
- #[inline]
- pub fn compressedAxis(&self) -> SparseMatrixCompressedAxis {
- self._tab
- .get::<SparseMatrixCompressedAxis>(
- SparseMatrixIndexCSX::VT_COMPRESSEDAXIS,
- Some(SparseMatrixCompressedAxis::Row),
- )
- .unwrap()
- }
- /// The type of values in indptrBuffer
- #[inline]
- pub fn indptrType(&self) -> Int<'a> {
- self._tab
- .get::<flatbuffers::ForwardsUOffset<Int>>(
- SparseMatrixIndexCSX::VT_INDPTRTYPE,
- None,
- )
- .unwrap()
- }
- /// 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].
- /// ```
- #[inline]
- pub fn indptrBuffer(&self) -> &'a Buffer {
- self._tab
- .get::<Buffer>(SparseMatrixIndexCSX::VT_INDPTRBUFFER, None)
- .unwrap()
- }
- /// The type of values in indicesBuffer
- #[inline]
- pub fn indicesType(&self) -> Int<'a> {
- self._tab
- .get::<flatbuffers::ForwardsUOffset<Int>>(
- SparseMatrixIndexCSX::VT_INDICESTYPE,
- None,
- )
- .unwrap()
- }
- /// 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.
- #[inline]
- pub fn indicesBuffer(&self) -> &'a Buffer {
- self._tab
- .get::<Buffer>(SparseMatrixIndexCSX::VT_INDICESBUFFER, None)
- .unwrap()
- }
-}
-
-impl flatbuffers::Verifiable for SparseMatrixIndexCSX<'_> {
- #[inline]
- fn run_verifier(
- v: &mut flatbuffers::Verifier,
- pos: usize,
- ) -> Result<(), flatbuffers::InvalidFlatbuffer> {
- use flatbuffers::Verifiable;
- v.visit_table(pos)?
- .visit_field::<SparseMatrixCompressedAxis>(
- &"compressedAxis",
- Self::VT_COMPRESSEDAXIS,
- false,
- )?
- .visit_field::<flatbuffers::ForwardsUOffset<Int>>(
- &"indptrType",
- Self::VT_INDPTRTYPE,
- true,
- )?
- .visit_field::<Buffer>(&"indptrBuffer", Self::VT_INDPTRBUFFER, true)?
- .visit_field::<flatbuffers::ForwardsUOffset<Int>>(
- &"indicesType",
- Self::VT_INDICESTYPE,
- true,
- )?
- .visit_field::<Buffer>(&"indicesBuffer", Self::VT_INDICESBUFFER, true)?
- .finish();
- Ok(())
- }
-}
-pub struct SparseMatrixIndexCSXArgs<'a> {
- pub compressedAxis: SparseMatrixCompressedAxis,
- pub indptrType: Option<flatbuffers::WIPOffset<Int<'a>>>,
- pub indptrBuffer: Option<&'a Buffer>,
- pub indicesType: Option<flatbuffers::WIPOffset<Int<'a>>>,
- pub indicesBuffer: Option<&'a Buffer>,
-}
-impl<'a> Default for SparseMatrixIndexCSXArgs<'a> {
- #[inline]
- fn default() -> Self {
- SparseMatrixIndexCSXArgs {
- compressedAxis: SparseMatrixCompressedAxis::Row,
- indptrType: None, // required field
- indptrBuffer: None, // required field
- indicesType: None, // required field
- indicesBuffer: None, // required field
- }
- }
-}
-pub struct SparseMatrixIndexCSXBuilder<'a: 'b, 'b> {
- fbb_: &'b mut flatbuffers::FlatBufferBuilder<'a>,
- start_: flatbuffers::WIPOffset<flatbuffers::TableUnfinishedWIPOffset>,
-}
-impl<'a: 'b, 'b> SparseMatrixIndexCSXBuilder<'a, 'b> {
- #[inline]
- pub fn add_compressedAxis(&mut self, compressedAxis: SparseMatrixCompressedAxis) {
- self.fbb_.push_slot::<SparseMatrixCompressedAxis>(
- SparseMatrixIndexCSX::VT_COMPRESSEDAXIS,
- compressedAxis,
- SparseMatrixCompressedAxis::Row,
- );
- }
- #[inline]
- pub fn add_indptrType(&mut self, indptrType: flatbuffers::WIPOffset<Int<'b>>) {
- self.fbb_.push_slot_always::<flatbuffers::WIPOffset<Int>>(
- SparseMatrixIndexCSX::VT_INDPTRTYPE,
- indptrType,
- );
- }
- #[inline]
- pub fn add_indptrBuffer(&mut self, indptrBuffer: &Buffer) {
- self.fbb_.push_slot_always::<&Buffer>(
- SparseMatrixIndexCSX::VT_INDPTRBUFFER,
- indptrBuffer,
- );
- }
- #[inline]
- pub fn add_indicesType(&mut self, indicesType: flatbuffers::WIPOffset<Int<'b>>) {
- self.fbb_.push_slot_always::<flatbuffers::WIPOffset<Int>>(
- SparseMatrixIndexCSX::VT_INDICESTYPE,
- indicesType,
- );
- }
- #[inline]
- pub fn add_indicesBuffer(&mut self, indicesBuffer: &Buffer) {
- self.fbb_.push_slot_always::<&Buffer>(
- SparseMatrixIndexCSX::VT_INDICESBUFFER,
- indicesBuffer,
- );
- }
- #[inline]
- pub fn new(
- _fbb: &'b mut flatbuffers::FlatBufferBuilder<'a>,
- ) -> SparseMatrixIndexCSXBuilder<'a, 'b> {
- let start = _fbb.start_table();
- SparseMatrixIndexCSXBuilder {
- fbb_: _fbb,
- start_: start,
- }
- }
- #[inline]
- pub fn finish(self) -> flatbuffers::WIPOffset<SparseMatrixIndexCSX<'a>> {
- let o = self.fbb_.end_table(self.start_);
- self.fbb_
- .required(o, SparseMatrixIndexCSX::VT_INDPTRTYPE, "indptrType");
- self.fbb_
- .required(o, SparseMatrixIndexCSX::VT_INDPTRBUFFER, "indptrBuffer");
- self.fbb_
- .required(o, SparseMatrixIndexCSX::VT_INDICESTYPE, "indicesType");
- self.fbb_
- .required(o, SparseMatrixIndexCSX::VT_INDICESBUFFER, "indicesBuffer");
- flatbuffers::WIPOffset::new(o.value())
- }
-}
-
-impl std::fmt::Debug for SparseMatrixIndexCSX<'_> {
- fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
- let mut ds = f.debug_struct("SparseMatrixIndexCSX");
- ds.field("compressedAxis", &self.compressedAxis());
- ds.field("indptrType", &self.indptrType());
- ds.field("indptrBuffer", &self.indptrBuffer());
- ds.field("indicesType", &self.indicesType());
- ds.field("indicesBuffer", &self.indicesBuffer());
- ds.finish()
- }
-}
-pub enum SparseTensorIndexCSFOffset {}
-#[derive(Copy, Clone, PartialEq)]
-
-/// Compressed Sparse Fiber (CSF) sparse tensor index.
-pub struct SparseTensorIndexCSF<'a> {
- pub _tab: flatbuffers::Table<'a>,
-}
-
-impl<'a> flatbuffers::Follow<'a> for SparseTensorIndexCSF<'a> {
- type Inner = SparseTensorIndexCSF<'a>;
- #[inline]
- fn follow(buf: &'a [u8], loc: usize) -> Self::Inner {
- Self {
- _tab: flatbuffers::Table { buf, loc },
- }
- }
-}
-
-impl<'a> SparseTensorIndexCSF<'a> {
- #[inline]
- pub fn init_from_table(table: flatbuffers::Table<'a>) -> Self {
- SparseTensorIndexCSF { _tab: table }
- }
- #[allow(unused_mut)]
- pub fn create<'bldr: 'args, 'args: 'mut_bldr, 'mut_bldr>(
- _fbb: &'mut_bldr mut flatbuffers::FlatBufferBuilder<'bldr>,
- args: &'args SparseTensorIndexCSFArgs<'args>,
- ) -> flatbuffers::WIPOffset<SparseTensorIndexCSF<'bldr>> {
- let mut builder = SparseTensorIndexCSFBuilder::new(_fbb);
- if let Some(x) = args.axisOrder {
- builder.add_axisOrder(x);
- }
- if let Some(x) = args.indicesBuffers {
- builder.add_indicesBuffers(x);
- }
- if let Some(x) = args.indicesType {
- builder.add_indicesType(x);
- }
- if let Some(x) = args.indptrBuffers {
- builder.add_indptrBuffers(x);
- }
- if let Some(x) = args.indptrType {
- builder.add_indptrType(x);
- }
- builder.finish()
- }
-
- pub const VT_INDPTRTYPE: flatbuffers::VOffsetT = 4;
- pub const VT_INDPTRBUFFERS: flatbuffers::VOffsetT = 6;
- pub const VT_INDICESTYPE: flatbuffers::VOffsetT = 8;
- pub const VT_INDICESBUFFERS: flatbuffers::VOffsetT = 10;
- pub const VT_AXISORDER: flatbuffers::VOffsetT = 12;
-
- /// 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
- #[inline]
- pub fn indptrType(&self) -> Int<'a> {
- self._tab
- .get::<flatbuffers::ForwardsUOffset<Int>>(
- SparseTensorIndexCSF::VT_INDPTRTYPE,
- None,
- )
- .unwrap()
- }
- /// 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]
- /// ].
- /// ```
- #[inline]
- pub fn indptrBuffers(&self) -> &'a [Buffer] {
- self._tab
- .get::<flatbuffers::ForwardsUOffset<flatbuffers::Vector<'a, Buffer>>>(
- SparseTensorIndexCSF::VT_INDPTRBUFFERS,
- None,
- )
- .map(|v| v.safe_slice())
- .unwrap()
- }
- /// The type of values in indicesBuffers
- #[inline]
- pub fn indicesType(&self) -> Int<'a> {
- self._tab
- .get::<flatbuffers::ForwardsUOffset<Int>>(
- SparseTensorIndexCSF::VT_INDICESTYPE,
- None,
- )
- .unwrap()
- }
- /// 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]
- /// ].
- /// ```
- #[inline]
- pub fn indicesBuffers(&self) -> &'a [Buffer] {
- self._tab
- .get::<flatbuffers::ForwardsUOffset<flatbuffers::Vector<'a, Buffer>>>(
- SparseTensorIndexCSF::VT_INDICESBUFFERS,
- None,
- )
- .map(|v| v.safe_slice())
- .unwrap()
- }
- /// 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].
- /// ```
- #[inline]
- pub fn axisOrder(&self) -> flatbuffers::Vector<'a, i32> {
- self._tab
- .get::<flatbuffers::ForwardsUOffset<flatbuffers::Vector<'a, i32>>>(
- SparseTensorIndexCSF::VT_AXISORDER,
- None,
- )
- .unwrap()
- }
-}
-
-impl flatbuffers::Verifiable for SparseTensorIndexCSF<'_> {
- #[inline]
- fn run_verifier(
- v: &mut flatbuffers::Verifier,
- pos: usize,
- ) -> Result<(), flatbuffers::InvalidFlatbuffer> {
- use flatbuffers::Verifiable;
- v.visit_table(pos)?
- .visit_field::<flatbuffers::ForwardsUOffset<Int>>(
- &"indptrType",
- Self::VT_INDPTRTYPE,
- true,
- )?
- .visit_field::<flatbuffers::ForwardsUOffset<flatbuffers::Vector<'_, Buffer>>>(
- &"indptrBuffers",
- Self::VT_INDPTRBUFFERS,
- true,
- )?
- .visit_field::<flatbuffers::ForwardsUOffset<Int>>(
- &"indicesType",
- Self::VT_INDICESTYPE,
- true,
- )?
- .visit_field::<flatbuffers::ForwardsUOffset<flatbuffers::Vector<'_, Buffer>>>(
- &"indicesBuffers",
- Self::VT_INDICESBUFFERS,
- true,
- )?
- .visit_field::<flatbuffers::ForwardsUOffset<flatbuffers::Vector<'_, i32>>>(
- &"axisOrder",
- Self::VT_AXISORDER,
- true,
- )?
- .finish();
- Ok(())
- }
-}
-pub struct SparseTensorIndexCSFArgs<'a> {
- pub indptrType: Option<flatbuffers::WIPOffset<Int<'a>>>,
- pub indptrBuffers: Option<flatbuffers::WIPOffset<flatbuffers::Vector<'a, Buffer>>>,
- pub indicesType: Option<flatbuffers::WIPOffset<Int<'a>>>,
- pub indicesBuffers: Option<flatbuffers::WIPOffset<flatbuffers::Vector<'a, Buffer>>>,
- pub axisOrder: Option<flatbuffers::WIPOffset<flatbuffers::Vector<'a, i32>>>,
-}
-impl<'a> Default for SparseTensorIndexCSFArgs<'a> {
- #[inline]
- fn default() -> Self {
- SparseTensorIndexCSFArgs {
- indptrType: None, // required field
- indptrBuffers: None, // required field
- indicesType: None, // required field
- indicesBuffers: None, // required field
- axisOrder: None, // required field
- }
- }
-}
-pub struct SparseTensorIndexCSFBuilder<'a: 'b, 'b> {
- fbb_: &'b mut flatbuffers::FlatBufferBuilder<'a>,
- start_: flatbuffers::WIPOffset<flatbuffers::TableUnfinishedWIPOffset>,
-}
-impl<'a: 'b, 'b> SparseTensorIndexCSFBuilder<'a, 'b> {
- #[inline]
- pub fn add_indptrType(&mut self, indptrType: flatbuffers::WIPOffset<Int<'b>>) {
- self.fbb_.push_slot_always::<flatbuffers::WIPOffset<Int>>(
- SparseTensorIndexCSF::VT_INDPTRTYPE,
- indptrType,
- );
- }
- #[inline]
- pub fn add_indptrBuffers(
- &mut self,
- indptrBuffers: flatbuffers::WIPOffset<flatbuffers::Vector<'b, Buffer>>,
- ) {
- self.fbb_.push_slot_always::<flatbuffers::WIPOffset<_>>(
- SparseTensorIndexCSF::VT_INDPTRBUFFERS,
- indptrBuffers,
- );
- }
- #[inline]
- pub fn add_indicesType(&mut self, indicesType: flatbuffers::WIPOffset<Int<'b>>) {
- self.fbb_.push_slot_always::<flatbuffers::WIPOffset<Int>>(
- SparseTensorIndexCSF::VT_INDICESTYPE,
- indicesType,
- );
- }
- #[inline]
- pub fn add_indicesBuffers(
- &mut self,
- indicesBuffers: flatbuffers::WIPOffset<flatbuffers::Vector<'b, Buffer>>,
- ) {
- self.fbb_.push_slot_always::<flatbuffers::WIPOffset<_>>(
- SparseTensorIndexCSF::VT_INDICESBUFFERS,
- indicesBuffers,
- );
- }
- #[inline]
- pub fn add_axisOrder(
- &mut self,
- axisOrder: flatbuffers::WIPOffset<flatbuffers::Vector<'b, i32>>,
- ) {
- self.fbb_.push_slot_always::<flatbuffers::WIPOffset<_>>(
- SparseTensorIndexCSF::VT_AXISORDER,
- axisOrder,
- );
- }
- #[inline]
- pub fn new(
- _fbb: &'b mut flatbuffers::FlatBufferBuilder<'a>,
- ) -> SparseTensorIndexCSFBuilder<'a, 'b> {
- let start = _fbb.start_table();
- SparseTensorIndexCSFBuilder {
- fbb_: _fbb,
- start_: start,
- }
- }
- #[inline]
- pub fn finish(self) -> flatbuffers::WIPOffset<SparseTensorIndexCSF<'a>> {
- let o = self.fbb_.end_table(self.start_);
- self.fbb_
- .required(o, SparseTensorIndexCSF::VT_INDPTRTYPE, "indptrType");
- self.fbb_
- .required(o, SparseTensorIndexCSF::VT_INDPTRBUFFERS, "indptrBuffers");
- self.fbb_
- .required(o, SparseTensorIndexCSF::VT_INDICESTYPE, "indicesType");
- self.fbb_
- .required(o, SparseTensorIndexCSF::VT_INDICESBUFFERS, "indicesBuffers");
- self.fbb_
- .required(o, SparseTensorIndexCSF::VT_AXISORDER, "axisOrder");
- flatbuffers::WIPOffset::new(o.value())
- }
-}
-
-impl std::fmt::Debug for SparseTensorIndexCSF<'_> {
- fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
- let mut ds = f.debug_struct("SparseTensorIndexCSF");
- ds.field("indptrType", &self.indptrType());
- ds.field("indptrBuffers", &self.indptrBuffers());
- ds.field("indicesType", &self.indicesType());
- ds.field("indicesBuffers", &self.indicesBuffers());
- ds.field("axisOrder", &self.axisOrder());
- ds.finish()
- }
-}
-pub enum SparseTensorOffset {}
-#[derive(Copy, Clone, PartialEq)]
-
-pub struct SparseTensor<'a> {
- pub _tab: flatbuffers::Table<'a>,
-}
-
-impl<'a> flatbuffers::Follow<'a> for SparseTensor<'a> {
- type Inner = SparseTensor<'a>;
- #[inline]
- fn follow(buf: &'a [u8], loc: usize) -> Self::Inner {
- Self {
- _tab: flatbuffers::Table { buf, loc },
- }
- }
-}
-
-impl<'a> SparseTensor<'a> {
- #[inline]
- pub fn init_from_table(table: flatbuffers::Table<'a>) -> Self {
- SparseTensor { _tab: table }
- }
- #[allow(unused_mut)]
- pub fn create<'bldr: 'args, 'args: 'mut_bldr, 'mut_bldr>(
- _fbb: &'mut_bldr mut flatbuffers::FlatBufferBuilder<'bldr>,
- args: &'args SparseTensorArgs<'args>,
- ) -> flatbuffers::WIPOffset<SparseTensor<'bldr>> {
- let mut builder = SparseTensorBuilder::new(_fbb);
- builder.add_non_zero_length(args.non_zero_length);
- if let Some(x) = args.data {
- builder.add_data(x);
- }
- if let Some(x) = args.sparseIndex {
- builder.add_sparseIndex(x);
- }
- if let Some(x) = args.shape {
- builder.add_shape(x);
- }
- if let Some(x) = args.type_ {
- builder.add_type_(x);
- }
- builder.add_sparseIndex_type(args.sparseIndex_type);
- builder.add_type_type(args.type_type);
- builder.finish()
- }
-
- pub const VT_TYPE_TYPE: flatbuffers::VOffsetT = 4;
- pub const VT_TYPE_: flatbuffers::VOffsetT = 6;
- pub const VT_SHAPE: flatbuffers::VOffsetT = 8;
- pub const VT_NON_ZERO_LENGTH: flatbuffers::VOffsetT = 10;
- pub const VT_SPARSEINDEX_TYPE: flatbuffers::VOffsetT = 12;
- pub const VT_SPARSEINDEX: flatbuffers::VOffsetT = 14;
- pub const VT_DATA: flatbuffers::VOffsetT = 16;
-
- #[inline]
- pub fn type_type(&self) -> Type {
- self._tab
- .get::<Type>(SparseTensor::VT_TYPE_TYPE, Some(Type::NONE))
- .unwrap()
- }
- /// The type of data contained in a value cell.
- /// Currently only fixed-width value types are supported,
- /// no strings or nested types.
- #[inline]
- pub fn type_(&self) -> flatbuffers::Table<'a> {
- self._tab
- .get::<flatbuffers::ForwardsUOffset<flatbuffers::Table<'a>>>(
- SparseTensor::VT_TYPE_,
- None,
- )
- .unwrap()
- }
- /// The dimensions of the tensor, optionally named.
- #[inline]
- pub fn shape(
- &self,
- ) -> flatbuffers::Vector<'a, flatbuffers::ForwardsUOffset<TensorDim<'a>>> {
- self._tab
- .get::<flatbuffers::ForwardsUOffset<
- flatbuffers::Vector<'a, flatbuffers::ForwardsUOffset<TensorDim>>,
- >>(SparseTensor::VT_SHAPE, None)
- .unwrap()
- }
- /// The number of non-zero values in a sparse tensor.
- #[inline]
- pub fn non_zero_length(&self) -> i64 {
- self._tab
- .get::<i64>(SparseTensor::VT_NON_ZERO_LENGTH, Some(0))
- .unwrap()
- }
- #[inline]
- pub fn sparseIndex_type(&self) -> SparseTensorIndex {
- self._tab
- .get::<SparseTensorIndex>(
- SparseTensor::VT_SPARSEINDEX_TYPE,
- Some(SparseTensorIndex::NONE),
- )
- .unwrap()
- }
- /// Sparse tensor index
- #[inline]
- pub fn sparseIndex(&self) -> flatbuffers::Table<'a> {
- self._tab
- .get::<flatbuffers::ForwardsUOffset<flatbuffers::Table<'a>>>(
- SparseTensor::VT_SPARSEINDEX,
- None,
- )
- .unwrap()
- }
- /// The location and size of the tensor's data
- #[inline]
- pub fn data(&self) -> &'a Buffer {
- self._tab
- .get::<Buffer>(SparseTensor::VT_DATA, None)
- .unwrap()
- }
- #[inline]
- #[allow(non_snake_case)]
- pub fn type_as_null(&self) -> Option<Null<'a>> {
- if self.type_type() == Type::Null {
- let u = self.type_();
- Some(Null::init_from_table(u))
- } else {
- None
- }
- }
-
- #[inline]
- #[allow(non_snake_case)]
- pub fn type_as_int(&self) -> Option<Int<'a>> {
- if self.type_type() == Type::Int {
- let u = self.type_();
- Some(Int::init_from_table(u))
- } else {
- None
- }
- }
-
- #[inline]
- #[allow(non_snake_case)]
- pub fn type_as_floating_point(&self) -> Option<FloatingPoint<'a>> {
- if self.type_type() == Type::FloatingPoint {
- let u = self.type_();
- Some(FloatingPoint::init_from_table(u))
- } else {
- None
- }
- }
-
- #[inline]
- #[allow(non_snake_case)]
- pub fn type_as_binary(&self) -> Option<Binary<'a>> {
- if self.type_type() == Type::Binary {
- let u = self.type_();
- Some(Binary::init_from_table(u))
- } else {
- None
- }
- }
-
- #[inline]
- #[allow(non_snake_case)]
- pub fn type_as_utf_8(&self) -> Option<Utf8<'a>> {
- if self.type_type() == Type::Utf8 {
- let u = self.type_();
- Some(Utf8::init_from_table(u))
- } else {
- None
- }
- }
-
- #[inline]
- #[allow(non_snake_case)]
- pub fn type_as_bool(&self) -> Option<Bool<'a>> {
- if self.type_type() == Type::Bool {
- let u = self.type_();
- Some(Bool::init_from_table(u))
- } else {
- None
- }
- }
-
- #[inline]
- #[allow(non_snake_case)]
- pub fn type_as_decimal(&self) -> Option<Decimal<'a>> {
- if self.type_type() == Type::Decimal {
- let u = self.type_();
- Some(Decimal::init_from_table(u))
- } else {
- None
- }
- }
-
- #[inline]
- #[allow(non_snake_case)]
- pub fn type_as_date(&self) -> Option<Date<'a>> {
- if self.type_type() == Type::Date {
- let u = self.type_();
- Some(Date::init_from_table(u))
- } else {
- None
- }
- }
-
- #[inline]
- #[allow(non_snake_case)]
- pub fn type_as_time(&self) -> Option<Time<'a>> {
- if self.type_type() == Type::Time {
- let u = self.type_();
- Some(Time::init_from_table(u))
- } else {
- None
- }
- }
-
- #[inline]
- #[allow(non_snake_case)]
- pub fn type_as_timestamp(&self) -> Option<Timestamp<'a>> {
- if self.type_type() == Type::Timestamp {
- let u = self.type_();
- Some(Timestamp::init_from_table(u))
- } else {
- None
- }
- }
-
- #[inline]
- #[allow(non_snake_case)]
- pub fn type_as_interval(&self) -> Option<Interval<'a>> {
- if self.type_type() == Type::Interval {
- let u = self.type_();
- Some(Interval::init_from_table(u))
- } else {
- None
- }
- }
-
- #[inline]
- #[allow(non_snake_case)]
- pub fn type_as_list(&self) -> Option<List<'a>> {
- if self.type_type() == Type::List {
- let u = self.type_();
- Some(List::init_from_table(u))
- } else {
- None
- }
- }
-
- #[inline]
- #[allow(non_snake_case)]
- pub fn type_as_struct_(&self) -> Option<Struct_<'a>> {
- if self.type_type() == Type::Struct_ {
- let u = self.type_();
- Some(Struct_::init_from_table(u))
- } else {
- None
- }
- }
-
- #[inline]
- #[allow(non_snake_case)]
- pub fn type_as_union(&self) -> Option<Union<'a>> {
- if self.type_type() == Type::Union {
- let u = self.type_();
- Some(Union::init_from_table(u))
- } else {
- None
- }
- }
-
- #[inline]
- #[allow(non_snake_case)]
- pub fn type_as_fixed_size_binary(&self) -> Option<FixedSizeBinary<'a>> {
- if self.type_type() == Type::FixedSizeBinary {
- let u = self.type_();
- Some(FixedSizeBinary::init_from_table(u))
- } else {
- None
- }
- }
-
- #[inline]
- #[allow(non_snake_case)]
- pub fn type_as_fixed_size_list(&self) -> Option<FixedSizeList<'a>> {
- if self.type_type() == Type::FixedSizeList {
- let u = self.type_();
- Some(FixedSizeList::init_from_table(u))
- } else {
- None
- }
- }
-
- #[inline]
- #[allow(non_snake_case)]
- pub fn type_as_map(&self) -> Option<Map<'a>> {
- if self.type_type() == Type::Map {
- let u = self.type_();
- Some(Map::init_from_table(u))
- } else {
- None
- }
- }
-
- #[inline]
- #[allow(non_snake_case)]
- pub fn type_as_duration(&self) -> Option<Duration<'a>> {
- if self.type_type() == Type::Duration {
- let u = self.type_();
- Some(Duration::init_from_table(u))
- } else {
- None
- }
- }
-
- #[inline]
- #[allow(non_snake_case)]
- pub fn type_as_large_binary(&self) -> Option<LargeBinary<'a>> {
- if self.type_type() == Type::LargeBinary {
- let u = self.type_();
- Some(LargeBinary::init_from_table(u))
- } else {
- None
- }
- }
-
- #[inline]
- #[allow(non_snake_case)]
- pub fn type_as_large_utf_8(&self) -> Option<LargeUtf8<'a>> {
- if self.type_type() == Type::LargeUtf8 {
- let u = self.type_();
- Some(LargeUtf8::init_from_table(u))
- } else {
- None
- }
- }
-
- #[inline]
- #[allow(non_snake_case)]
- pub fn type_as_large_list(&self) -> Option<LargeList<'a>> {
- if self.type_type() == Type::LargeList {
- let u = self.type_();
- Some(LargeList::init_from_table(u))
- } else {
- None
- }
- }
-
- #[inline]
- #[allow(non_snake_case)]
- pub fn sparseIndex_as_sparse_tensor_index_coo(
- &self,
- ) -> Option<SparseTensorIndexCOO<'a>> {
- if self.sparseIndex_type() == SparseTensorIndex::SparseTensorIndexCOO {
- let u = self.sparseIndex();
- Some(SparseTensorIndexCOO::init_from_table(u))
- } else {
- None
- }
- }
-
- #[inline]
- #[allow(non_snake_case)]
- pub fn sparseIndex_as_sparse_matrix_index_csx(
- &self,
- ) -> Option<SparseMatrixIndexCSX<'a>> {
- if self.sparseIndex_type() == SparseTensorIndex::SparseMatrixIndexCSX {
- let u = self.sparseIndex();
- Some(SparseMatrixIndexCSX::init_from_table(u))
- } else {
- None
- }
- }
-
- #[inline]
- #[allow(non_snake_case)]
- pub fn sparseIndex_as_sparse_tensor_index_csf(
- &self,
- ) -> Option<SparseTensorIndexCSF<'a>> {
- if self.sparseIndex_type() == SparseTensorIndex::SparseTensorIndexCSF {
- let u = self.sparseIndex();
- Some(SparseTensorIndexCSF::init_from_table(u))
- } else {
- None
- }
- }
-}
-
-impl flatbuffers::Verifiable for SparseTensor<'_> {
- #[inline]
- fn run_verifier(
- v: &mut flatbuffers::Verifier,
- pos: usize,
- ) -> Result<(), flatbuffers::InvalidFlatbuffer> {
- use flatbuffers::Verifiable;
- v.visit_table(pos)?
- .visit_union::<Type, _>(&"type_type", Self::VT_TYPE_TYPE, &"type_", Self::VT_TYPE_, true, |key, v, pos| {
- match key {
- Type::Null => v.verify_union_variant::<flatbuffers::ForwardsUOffset<Null>>("Type::Null", pos),
- Type::Int => v.verify_union_variant::<flatbuffers::ForwardsUOffset<Int>>("Type::Int", pos),
- Type::FloatingPoint => v.verify_union_variant::<flatbuffers::ForwardsUOffset<FloatingPoint>>("Type::FloatingPoint", pos),
- Type::Binary => v.verify_union_variant::<flatbuffers::ForwardsUOffset<Binary>>("Type::Binary", pos),
- Type::Utf8 => v.verify_union_variant::<flatbuffers::ForwardsUOffset<Utf8>>("Type::Utf8", pos),
- Type::Bool => v.verify_union_variant::<flatbuffers::ForwardsUOffset<Bool>>("Type::Bool", pos),
- Type::Decimal => v.verify_union_variant::<flatbuffers::ForwardsUOffset<Decimal>>("Type::Decimal", pos),
- Type::Date => v.verify_union_variant::<flatbuffers::ForwardsUOffset<Date>>("Type::Date", pos),
- Type::Time => v.verify_union_variant::<flatbuffers::ForwardsUOffset<Time>>("Type::Time", pos),
- Type::Timestamp => v.verify_union_variant::<flatbuffers::ForwardsUOffset<Timestamp>>("Type::Timestamp", pos),
- Type::Interval => v.verify_union_variant::<flatbuffers::ForwardsUOffset<Interval>>("Type::Interval", pos),
- Type::List => v.verify_union_variant::<flatbuffers::ForwardsUOffset<List>>("Type::List", pos),
- Type::Struct_ => v.verify_union_variant::<flatbuffers::ForwardsUOffset<Struct_>>("Type::Struct_", pos),
- Type::Union => v.verify_union_variant::<flatbuffers::ForwardsUOffset<Union>>("Type::Union", pos),
- Type::FixedSizeBinary => v.verify_union_variant::<flatbuffers::ForwardsUOffset<FixedSizeBinary>>("Type::FixedSizeBinary", pos),
- Type::FixedSizeList => v.verify_union_variant::<flatbuffers::ForwardsUOffset<FixedSizeList>>("Type::FixedSizeList", pos),
- Type::Map => v.verify_union_variant::<flatbuffers::ForwardsUOffset<Map>>("Type::Map", pos),
- Type::Duration => v.verify_union_variant::<flatbuffers::ForwardsUOffset<Duration>>("Type::Duration", pos),
- Type::LargeBinary => v.verify_union_variant::<flatbuffers::ForwardsUOffset<LargeBinary>>("Type::LargeBinary", pos),
- Type::LargeUtf8 => v.verify_union_variant::<flatbuffers::ForwardsUOffset<LargeUtf8>>("Type::LargeUtf8", pos),
- Type::LargeList => v.verify_union_variant::<flatbuffers::ForwardsUOffset<LargeList>>("Type::LargeList", pos),
- _ => Ok(()),
- }
- })?
- .visit_field::<flatbuffers::ForwardsUOffset<flatbuffers::Vector<'_, flatbuffers::ForwardsUOffset<TensorDim>>>>(&"shape", Self::VT_SHAPE, true)?
- .visit_field::<i64>(&"non_zero_length", Self::VT_NON_ZERO_LENGTH, false)?
- .visit_union::<SparseTensorIndex, _>(&"sparseIndex_type", Self::VT_SPARSEINDEX_TYPE, &"sparseIndex", Self::VT_SPARSEINDEX, true, |key, v, pos| {
- match key {
- SparseTensorIndex::SparseTensorIndexCOO => v.verify_union_variant::<flatbuffers::ForwardsUOffset<SparseTensorIndexCOO>>("SparseTensorIndex::SparseTensorIndexCOO", pos),
- SparseTensorIndex::SparseMatrixIndexCSX => v.verify_union_variant::<flatbuffers::ForwardsUOffset<SparseMatrixIndexCSX>>("SparseTensorIndex::SparseMatrixIndexCSX", pos),
- SparseTensorIndex::SparseTensorIndexCSF => v.verify_union_variant::<flatbuffers::ForwardsUOffset<SparseTensorIndexCSF>>("SparseTensorIndex::SparseTensorIndexCSF", pos),
- _ => Ok(()),
- }
- })?
- .visit_field::<Buffer>(&"data", Self::VT_DATA, true)?
- .finish();
- Ok(())
- }
-}
-pub struct SparseTensorArgs<'a> {
- pub type_type: Type,
- pub type_: Option<flatbuffers::WIPOffset<flatbuffers::UnionWIPOffset>>,
- pub shape: Option<
- flatbuffers::WIPOffset<
- flatbuffers::Vector<'a, flatbuffers::ForwardsUOffset<TensorDim<'a>>>,
- >,
- >,
- pub non_zero_length: i64,
- pub sparseIndex_type: SparseTensorIndex,
- pub sparseIndex: Option<flatbuffers::WIPOffset<flatbuffers::UnionWIPOffset>>,
- pub data: Option<&'a Buffer>,
-}
-impl<'a> Default for SparseTensorArgs<'a> {
- #[inline]
- fn default() -> Self {
- SparseTensorArgs {
- type_type: Type::NONE,
- type_: None, // required field
- shape: None, // required field
- non_zero_length: 0,
- sparseIndex_type: SparseTensorIndex::NONE,
- sparseIndex: None, // required field
- data: None, // required field
- }
- }
-}
-pub struct SparseTensorBuilder<'a: 'b, 'b> {
- fbb_: &'b mut flatbuffers::FlatBufferBuilder<'a>,
- start_: flatbuffers::WIPOffset<flatbuffers::TableUnfinishedWIPOffset>,
-}
-impl<'a: 'b, 'b> SparseTensorBuilder<'a, 'b> {
- #[inline]
- pub fn add_type_type(&mut self, type_type: Type) {
- self.fbb_
- .push_slot::<Type>(SparseTensor::VT_TYPE_TYPE, type_type, Type::NONE);
- }
- #[inline]
- pub fn add_type_(
- &mut self,
- type_: flatbuffers::WIPOffset<flatbuffers::UnionWIPOffset>,
- ) {
- self.fbb_
- .push_slot_always::<flatbuffers::WIPOffset<_>>(SparseTensor::VT_TYPE_, type_);
- }
- #[inline]
- pub fn add_shape(
- &mut self,
- shape: flatbuffers::WIPOffset<
- flatbuffers::Vector<'b, flatbuffers::ForwardsUOffset<TensorDim<'b>>>,
- >,
- ) {
- self.fbb_
- .push_slot_always::<flatbuffers::WIPOffset<_>>(SparseTensor::VT_SHAPE, shape);
- }
- #[inline]
- pub fn add_non_zero_length(&mut self, non_zero_length: i64) {
- self.fbb_
- .push_slot::<i64>(SparseTensor::VT_NON_ZERO_LENGTH, non_zero_length, 0);
- }
- #[inline]
- pub fn add_sparseIndex_type(&mut self, sparseIndex_type: SparseTensorIndex) {
- self.fbb_.push_slot::<SparseTensorIndex>(
- SparseTensor::VT_SPARSEINDEX_TYPE,
- sparseIndex_type,
- SparseTensorIndex::NONE,
- );
- }
- #[inline]
- pub fn add_sparseIndex(
- &mut self,
- sparseIndex: flatbuffers::WIPOffset<flatbuffers::UnionWIPOffset>,
- ) {
- self.fbb_.push_slot_always::<flatbuffers::WIPOffset<_>>(
- SparseTensor::VT_SPARSEINDEX,
- sparseIndex,
- );
- }
- #[inline]
- pub fn add_data(&mut self, data: &Buffer) {
- self.fbb_
- .push_slot_always::<&Buffer>(SparseTensor::VT_DATA, data);
- }
- #[inline]
- pub fn new(
- _fbb: &'b mut flatbuffers::FlatBufferBuilder<'a>,
- ) -> SparseTensorBuilder<'a, 'b> {
- let start = _fbb.start_table();
- SparseTensorBuilder {
- fbb_: _fbb,
- start_: start,
- }
- }
- #[inline]
- pub fn finish(self) -> flatbuffers::WIPOffset<SparseTensor<'a>> {
- let o = self.fbb_.end_table(self.start_);
- self.fbb_.required(o, SparseTensor::VT_TYPE_, "type_");
- self.fbb_.required(o, SparseTensor::VT_SHAPE, "shape");
- self.fbb_
- .required(o, SparseTensor::VT_SPARSEINDEX, "sparseIndex");
- self.fbb_.required(o, SparseTensor::VT_DATA, "data");
- flatbuffers::WIPOffset::new(o.value())
- }
-}
-
-impl std::fmt::Debug for SparseTensor<'_> {
- fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
- let mut ds = f.debug_struct("SparseTensor");
- ds.field("type_type", &self.type_type());
- match self.type_type() {
- Type::Null => {
- if let Some(x) = self.type_as_null() {
- ds.field("type_", &x)
- } else {
- ds.field(
- "type_",
- &"InvalidFlatbuffer: Union discriminant does not match value.",
- )
- }
- }
- Type::Int => {
- if let Some(x) = self.type_as_int() {
- ds.field("type_", &x)
- } else {
- ds.field(
- "type_",
- &"InvalidFlatbuffer: Union discriminant does not match value.",
- )
- }
- }
- Type::FloatingPoint => {
- if let Some(x) = self.type_as_floating_point() {
- ds.field("type_", &x)
- } else {
- ds.field(
- "type_",
- &"InvalidFlatbuffer: Union discriminant does not match value.",
- )
- }
- }
- Type::Binary => {
- if let Some(x) = self.type_as_binary() {
- ds.field("type_", &x)
- } else {
- ds.field(
- "type_",
- &"InvalidFlatbuffer: Union discriminant does not match value.",
- )
- }
- }
- Type::Utf8 => {
- if let Some(x) = self.type_as_utf_8() {
- ds.field("type_", &x)
- } else {
- ds.field(
- "type_",
- &"InvalidFlatbuffer: Union discriminant does not match value.",
- )
- }
- }
- Type::Bool => {
- if let Some(x) = self.type_as_bool() {
- ds.field("type_", &x)
- } else {
- ds.field(
- "type_",
- &"InvalidFlatbuffer: Union discriminant does not match value.",
- )
- }
- }
- Type::Decimal => {
- if let Some(x) = self.type_as_decimal() {
- ds.field("type_", &x)
- } else {
- ds.field(
- "type_",
- &"InvalidFlatbuffer: Union discriminant does not match value.",
- )
- }
- }
- Type::Date => {
- if let Some(x) = self.type_as_date() {
- ds.field("type_", &x)
- } else {
- ds.field(
- "type_",
- &"InvalidFlatbuffer: Union discriminant does not match value.",
- )
- }
- }
- Type::Time => {
- if let Some(x) = self.type_as_time() {
- ds.field("type_", &x)
- } else {
- ds.field(
- "type_",
- &"InvalidFlatbuffer: Union discriminant does not match value.",
- )
- }
- }
- Type::Timestamp => {
- if let Some(x) = self.type_as_timestamp() {
- ds.field("type_", &x)
- } else {
- ds.field(
- "type_",
- &"InvalidFlatbuffer: Union discriminant does not match value.",
- )
- }
- }
- Type::Interval => {
- if let Some(x) = self.type_as_interval() {
- ds.field("type_", &x)
- } else {
- ds.field(
- "type_",
- &"InvalidFlatbuffer: Union discriminant does not match value.",
- )
- }
- }
- Type::List => {
- if let Some(x) = self.type_as_list() {
- ds.field("type_", &x)
- } else {
- ds.field(
- "type_",
- &"InvalidFlatbuffer: Union discriminant does not match value.",
- )
- }
- }
- Type::Struct_ => {
- if let Some(x) = self.type_as_struct_() {
- ds.field("type_", &x)
- } else {
- ds.field(
- "type_",
- &"InvalidFlatbuffer: Union discriminant does not match value.",
- )
- }
- }
- Type::Union => {
- if let Some(x) = self.type_as_union() {
- ds.field("type_", &x)
- } else {
- ds.field(
- "type_",
- &"InvalidFlatbuffer: Union discriminant does not match value.",
- )
- }
- }
- Type::FixedSizeBinary => {
- if let Some(x) = self.type_as_fixed_size_binary() {
- ds.field("type_", &x)
- } else {
- ds.field(
- "type_",
- &"InvalidFlatbuffer: Union discriminant does not match value.",
- )
- }
- }
- Type::FixedSizeList => {
- if let Some(x) = self.type_as_fixed_size_list() {
- ds.field("type_", &x)
- } else {
- ds.field(
- "type_",
- &"InvalidFlatbuffer: Union discriminant does not match value.",
- )
- }
- }
- Type::Map => {
- if let Some(x) = self.type_as_map() {
- ds.field("type_", &x)
- } else {
- ds.field(
- "type_",
- &"InvalidFlatbuffer: Union discriminant does not match value.",
- )
- }
- }
- Type::Duration => {
- if let Some(x) = self.type_as_duration() {
- ds.field("type_", &x)
- } else {
- ds.field(
- "type_",
- &"InvalidFlatbuffer: Union discriminant does not match value.",
- )
- }
- }
- Type::LargeBinary => {
- if let Some(x) = self.type_as_large_binary() {
- ds.field("type_", &x)
- } else {
- ds.field(
- "type_",
- &"InvalidFlatbuffer: Union discriminant does not match value.",
- )
- }
- }
- Type::LargeUtf8 => {
- if let Some(x) = self.type_as_large_utf_8() {
- ds.field("type_", &x)
- } else {
- ds.field(
- "type_",
- &"InvalidFlatbuffer: Union discriminant does not match value.",
- )
- }
- }
- Type::LargeList => {
- if let Some(x) = self.type_as_large_list() {
- ds.field("type_", &x)
- } else {
- ds.field(
- "type_",
- &"InvalidFlatbuffer: Union discriminant does not match value.",
- )
- }
- }
- _ => {
- let x: Option<()> = None;
- ds.field("type_", &x)
- }
- };
- ds.field("shape", &self.shape());
- ds.field("non_zero_length", &self.non_zero_length());
- ds.field("sparseIndex_type", &self.sparseIndex_type());
- match self.sparseIndex_type() {
- SparseTensorIndex::SparseTensorIndexCOO => {
- if let Some(x) = self.sparseIndex_as_sparse_tensor_index_coo() {
- ds.field("sparseIndex", &x)
- } else {
- ds.field(
- "sparseIndex",
- &"InvalidFlatbuffer: Union discriminant does not match value.",
- )
- }
- }
- SparseTensorIndex::SparseMatrixIndexCSX => {
- if let Some(x) = self.sparseIndex_as_sparse_matrix_index_csx() {
- ds.field("sparseIndex", &x)
- } else {
- ds.field(
- "sparseIndex",
- &"InvalidFlatbuffer: Union discriminant does not match value.",
- )
- }
- }
- SparseTensorIndex::SparseTensorIndexCSF => {
- if let Some(x) = self.sparseIndex_as_sparse_tensor_index_csf() {
- ds.field("sparseIndex", &x)
- } else {
- ds.field(
- "sparseIndex",
- &"InvalidFlatbuffer: Union discriminant does not match value.",
- )
- }
- }
- _ => {
- let x: Option<()> = None;
- ds.field("sparseIndex", &x)
- }
- };
- ds.field("data", &self.data());
- ds.finish()
- }
-}
-#[inline]
-#[deprecated(since = "2.0.0", note = "Deprecated in favor of `root_as...` methods.")]
-pub fn get_root_as_sparse_tensor<'a>(buf: &'a [u8]) -> SparseTensor<'a> {
- unsafe { flatbuffers::root_unchecked::<SparseTensor<'a>>(buf) }
-}
-
-#[inline]
-#[deprecated(since = "2.0.0", note = "Deprecated in favor of `root_as...` methods.")]
-pub fn get_size_prefixed_root_as_sparse_tensor<'a>(buf: &'a [u8]) -> SparseTensor<'a> {
- unsafe { flatbuffers::size_prefixed_root_unchecked::<SparseTensor<'a>>(buf) }
-}
-
-#[inline]
-/// Verifies that a buffer of bytes contains a `SparseTensor`
-/// and returns it.
-/// Note that verification is still experimental and may not
-/// catch every error, or be maximally performant. For the
-/// previous, unchecked, behavior use
-/// `root_as_sparse_tensor_unchecked`.
-pub fn root_as_sparse_tensor(
- buf: &[u8],
-) -> Result<SparseTensor, flatbuffers::InvalidFlatbuffer> {
- flatbuffers::root::<SparseTensor>(buf)
-}
-#[inline]
-/// Verifies that a buffer of bytes contains a size prefixed
-/// `SparseTensor` and returns it.
-/// Note that verification is still experimental and may not
-/// catch every error, or be maximally performant. For the
-/// previous, unchecked, behavior use
-/// `size_prefixed_root_as_sparse_tensor_unchecked`.
-pub fn size_prefixed_root_as_sparse_tensor(
- buf: &[u8],
-) -> Result<SparseTensor, flatbuffers::InvalidFlatbuffer> {
- flatbuffers::size_prefixed_root::<SparseTensor>(buf)
-}
-#[inline]
-/// Verifies, with the given options, that a buffer of bytes
-/// contains a `SparseTensor` and returns it.
-/// Note that verification is still experimental and may not
-/// catch every error, or be maximally performant. For the
-/// previous, unchecked, behavior use
-/// `root_as_sparse_tensor_unchecked`.
-pub fn root_as_sparse_tensor_with_opts<'b, 'o>(
- opts: &'o flatbuffers::VerifierOptions,
- buf: &'b [u8],
-) -> Result<SparseTensor<'b>, flatbuffers::InvalidFlatbuffer> {
- flatbuffers::root_with_opts::<SparseTensor<'b>>(opts, buf)
-}
-#[inline]
-/// Verifies, with the given verifier options, that a buffer of
-/// bytes contains a size prefixed `SparseTensor` and returns
-/// it. Note that verification is still experimental and may not
-/// catch every error, or be maximally performant. For the
-/// previous, unchecked, behavior use
-/// `root_as_sparse_tensor_unchecked`.
-pub fn size_prefixed_root_as_sparse_tensor_with_opts<'b, 'o>(
- opts: &'o flatbuffers::VerifierOptions,
- buf: &'b [u8],
-) -> Result<SparseTensor<'b>, flatbuffers::InvalidFlatbuffer> {
- flatbuffers::size_prefixed_root_with_opts::<SparseTensor<'b>>(opts, buf)
-}
-#[inline]
-/// Assumes, without verification, that a buffer of bytes contains a SparseTensor and returns it.
-/// # Safety
-/// Callers must trust the given bytes do indeed contain a valid `SparseTensor`.
-pub unsafe fn root_as_sparse_tensor_unchecked(buf: &[u8]) -> SparseTensor {
- flatbuffers::root_unchecked::<SparseTensor>(buf)
-}
-#[inline]
-/// Assumes, without verification, that a buffer of bytes contains a size prefixed SparseTensor and returns it.
-/// # Safety
-/// Callers must trust the given bytes do indeed contain a valid size prefixed `SparseTensor`.
-pub unsafe fn size_prefixed_root_as_sparse_tensor_unchecked(buf: &[u8]) -> SparseTensor {
- flatbuffers::size_prefixed_root_unchecked::<SparseTensor>(buf)
-}
-#[inline]
-pub fn finish_sparse_tensor_buffer<'a, 'b>(
- fbb: &'b mut flatbuffers::FlatBufferBuilder<'a>,
- root: flatbuffers::WIPOffset<SparseTensor<'a>>,
-) {
- fbb.finish(root, None);
-}
-
-#[inline]
-pub fn finish_size_prefixed_sparse_tensor_buffer<'a, 'b>(
- fbb: &'b mut flatbuffers::FlatBufferBuilder<'a>,
- root: flatbuffers::WIPOffset<SparseTensor<'a>>,
-) {
- fbb.finish_size_prefixed(root, None);
-}
diff --git a/arrow/src/ipc/gen/Tensor.rs b/arrow/src/ipc/gen/Tensor.rs
deleted file mode 100644
index 120636e..0000000
--- a/arrow/src/ipc/gen/Tensor.rs
+++ /dev/null
@@ -1,913 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-#![allow(dead_code)]
-#![allow(unused_imports)]
-
-use crate::ipc::gen::Schema::*;
-use flatbuffers::EndianScalar;
-use std::{cmp::Ordering, mem};
-// automatically generated by the FlatBuffers compiler, do not modify
-
-pub enum TensorDimOffset {}
-#[derive(Copy, Clone, PartialEq)]
-
-/// ----------------------------------------------------------------------
-/// Data structures for dense tensors
-/// Shape data for a single axis in a tensor
-pub struct TensorDim<'a> {
- pub _tab: flatbuffers::Table<'a>,
-}
-
-impl<'a> flatbuffers::Follow<'a> for TensorDim<'a> {
- type Inner = TensorDim<'a>;
- #[inline]
- fn follow(buf: &'a [u8], loc: usize) -> Self::Inner {
- Self {
- _tab: flatbuffers::Table { buf, loc },
- }
- }
-}
-
-impl<'a> TensorDim<'a> {
- #[inline]
- pub fn init_from_table(table: flatbuffers::Table<'a>) -> Self {
- TensorDim { _tab: table }
- }
- #[allow(unused_mut)]
- pub fn create<'bldr: 'args, 'args: 'mut_bldr, 'mut_bldr>(
- _fbb: &'mut_bldr mut flatbuffers::FlatBufferBuilder<'bldr>,
- args: &'args TensorDimArgs<'args>,
- ) -> flatbuffers::WIPOffset<TensorDim<'bldr>> {
- let mut builder = TensorDimBuilder::new(_fbb);
- builder.add_size_(args.size_);
- if let Some(x) = args.name {
- builder.add_name(x);
- }
- builder.finish()
- }
-
- pub const VT_SIZE_: flatbuffers::VOffsetT = 4;
- pub const VT_NAME: flatbuffers::VOffsetT = 6;
-
- /// Length of dimension
- #[inline]
- pub fn size_(&self) -> i64 {
- self._tab.get::<i64>(TensorDim::VT_SIZE_, Some(0)).unwrap()
- }
- /// Name of the dimension, optional
- #[inline]
- pub fn name(&self) -> Option<&'a str> {
- self._tab
- .get::<flatbuffers::ForwardsUOffset<&str>>(TensorDim::VT_NAME, None)
- }
-}
-
-impl flatbuffers::Verifiable for TensorDim<'_> {
- #[inline]
- fn run_verifier(
- v: &mut flatbuffers::Verifier,
- pos: usize,
- ) -> Result<(), flatbuffers::InvalidFlatbuffer> {
- use flatbuffers::Verifiable;
- v.visit_table(pos)?
- .visit_field::<i64>(&"size_", Self::VT_SIZE_, false)?
- .visit_field::<flatbuffers::ForwardsUOffset<&str>>(
- &"name",
- Self::VT_NAME,
- false,
- )?
- .finish();
- Ok(())
- }
-}
-pub struct TensorDimArgs<'a> {
- pub size_: i64,
- pub name: Option<flatbuffers::WIPOffset<&'a str>>,
-}
-impl<'a> Default for TensorDimArgs<'a> {
- #[inline]
- fn default() -> Self {
- TensorDimArgs {
- size_: 0,
- name: None,
- }
- }
-}
-pub struct TensorDimBuilder<'a: 'b, 'b> {
- fbb_: &'b mut flatbuffers::FlatBufferBuilder<'a>,
- start_: flatbuffers::WIPOffset<flatbuffers::TableUnfinishedWIPOffset>,
-}
-impl<'a: 'b, 'b> TensorDimBuilder<'a, 'b> {
- #[inline]
- pub fn add_size_(&mut self, size_: i64) {
- self.fbb_.push_slot::<i64>(TensorDim::VT_SIZE_, size_, 0);
- }
- #[inline]
- pub fn add_name(&mut self, name: flatbuffers::WIPOffset<&'b str>) {
- self.fbb_
- .push_slot_always::<flatbuffers::WIPOffset<_>>(TensorDim::VT_NAME, name);
- }
- #[inline]
- pub fn new(
- _fbb: &'b mut flatbuffers::FlatBufferBuilder<'a>,
- ) -> TensorDimBuilder<'a, 'b> {
- let start = _fbb.start_table();
- TensorDimBuilder {
- fbb_: _fbb,
- start_: start,
- }
- }
- #[inline]
- pub fn finish(self) -> flatbuffers::WIPOffset<TensorDim<'a>> {
- let o = self.fbb_.end_table(self.start_);
- flatbuffers::WIPOffset::new(o.value())
- }
-}
-
-impl std::fmt::Debug for TensorDim<'_> {
- fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
- let mut ds = f.debug_struct("TensorDim");
- ds.field("size_", &self.size_());
- ds.field("name", &self.name());
- ds.finish()
- }
-}
-pub enum TensorOffset {}
-#[derive(Copy, Clone, PartialEq)]
-
-pub struct Tensor<'a> {
- pub _tab: flatbuffers::Table<'a>,
-}
-
-impl<'a> flatbuffers::Follow<'a> for Tensor<'a> {
- type Inner = Tensor<'a>;
- #[inline]
- fn follow(buf: &'a [u8], loc: usize) -> Self::Inner {
- Self {
- _tab: flatbuffers::Table { buf, loc },
- }
- }
-}
-
-impl<'a> Tensor<'a> {
- #[inline]
- pub fn init_from_table(table: flatbuffers::Table<'a>) -> Self {
- Tensor { _tab: table }
- }
- #[allow(unused_mut)]
- pub fn create<'bldr: 'args, 'args: 'mut_bldr, 'mut_bldr>(
- _fbb: &'mut_bldr mut flatbuffers::FlatBufferBuilder<'bldr>,
- args: &'args TensorArgs<'args>,
- ) -> flatbuffers::WIPOffset<Tensor<'bldr>> {
- let mut builder = TensorBuilder::new(_fbb);
- if let Some(x) = args.data {
- builder.add_data(x);
- }
- if let Some(x) = args.strides {
- builder.add_strides(x);
- }
- if let Some(x) = args.shape {
- builder.add_shape(x);
- }
- if let Some(x) = args.type_ {
- builder.add_type_(x);
- }
- builder.add_type_type(args.type_type);
- builder.finish()
- }
-
- pub const VT_TYPE_TYPE: flatbuffers::VOffsetT = 4;
- pub const VT_TYPE_: flatbuffers::VOffsetT = 6;
- pub const VT_SHAPE: flatbuffers::VOffsetT = 8;
- pub const VT_STRIDES: flatbuffers::VOffsetT = 10;
- pub const VT_DATA: flatbuffers::VOffsetT = 12;
-
- #[inline]
- pub fn type_type(&self) -> Type {
- self._tab
- .get::<Type>(Tensor::VT_TYPE_TYPE, Some(Type::NONE))
- .unwrap()
- }
- /// The type of data contained in a value cell. Currently only fixed-width
- /// value types are supported, no strings or nested types
- #[inline]
- pub fn type_(&self) -> flatbuffers::Table<'a> {
- self._tab
- .get::<flatbuffers::ForwardsUOffset<flatbuffers::Table<'a>>>(
- Tensor::VT_TYPE_,
- None,
- )
- .unwrap()
- }
- /// The dimensions of the tensor, optionally named
- #[inline]
- pub fn shape(
- &self,
- ) -> flatbuffers::Vector<'a, flatbuffers::ForwardsUOffset<TensorDim<'a>>> {
- self._tab
- .get::<flatbuffers::ForwardsUOffset<
- flatbuffers::Vector<'a, flatbuffers::ForwardsUOffset<TensorDim>>,
- >>(Tensor::VT_SHAPE, None)
- .unwrap()
- }
- /// Non-negative byte offsets to advance one value cell along each dimension
- /// If omitted, default to row-major order (C-like).
- #[inline]
- pub fn strides(&self) -> Option<flatbuffers::Vector<'a, i64>> {
- self._tab
- .get::<flatbuffers::ForwardsUOffset<flatbuffers::Vector<'a, i64>>>(
- Tensor::VT_STRIDES,
- None,
- )
- }
- /// The location and size of the tensor's data
- #[inline]
- pub fn data(&self) -> &'a Buffer {
- self._tab.get::<Buffer>(Tensor::VT_DATA, None).unwrap()
- }
- #[inline]
- #[allow(non_snake_case)]
- pub fn type_as_null(&self) -> Option<Null<'a>> {
- if self.type_type() == Type::Null {
- let u = self.type_();
- Some(Null::init_from_table(u))
- } else {
- None
- }
- }
-
- #[inline]
- #[allow(non_snake_case)]
- pub fn type_as_int(&self) -> Option<Int<'a>> {
- if self.type_type() == Type::Int {
- let u = self.type_();
- Some(Int::init_from_table(u))
- } else {
- None
- }
- }
-
- #[inline]
- #[allow(non_snake_case)]
- pub fn type_as_floating_point(&self) -> Option<FloatingPoint<'a>> {
- if self.type_type() == Type::FloatingPoint {
- let u = self.type_();
- Some(FloatingPoint::init_from_table(u))
- } else {
- None
- }
- }
-
- #[inline]
- #[allow(non_snake_case)]
- pub fn type_as_binary(&self) -> Option<Binary<'a>> {
- if self.type_type() == Type::Binary {
- let u = self.type_();
- Some(Binary::init_from_table(u))
- } else {
- None
- }
- }
-
- #[inline]
- #[allow(non_snake_case)]
- pub fn type_as_utf_8(&self) -> Option<Utf8<'a>> {
- if self.type_type() == Type::Utf8 {
- let u = self.type_();
- Some(Utf8::init_from_table(u))
- } else {
- None
- }
- }
-
- #[inline]
- #[allow(non_snake_case)]
- pub fn type_as_bool(&self) -> Option<Bool<'a>> {
- if self.type_type() == Type::Bool {
- let u = self.type_();
- Some(Bool::init_from_table(u))
- } else {
- None
- }
- }
-
- #[inline]
- #[allow(non_snake_case)]
- pub fn type_as_decimal(&self) -> Option<Decimal<'a>> {
- if self.type_type() == Type::Decimal {
- let u = self.type_();
- Some(Decimal::init_from_table(u))
- } else {
- None
- }
- }
-
- #[inline]
- #[allow(non_snake_case)]
- pub fn type_as_date(&self) -> Option<Date<'a>> {
- if self.type_type() == Type::Date {
- let u = self.type_();
- Some(Date::init_from_table(u))
- } else {
- None
- }
- }
-
- #[inline]
- #[allow(non_snake_case)]
- pub fn type_as_time(&self) -> Option<Time<'a>> {
- if self.type_type() == Type::Time {
- let u = self.type_();
- Some(Time::init_from_table(u))
- } else {
- None
- }
- }
-
- #[inline]
- #[allow(non_snake_case)]
- pub fn type_as_timestamp(&self) -> Option<Timestamp<'a>> {
- if self.type_type() == Type::Timestamp {
- let u = self.type_();
- Some(Timestamp::init_from_table(u))
- } else {
- None
- }
- }
-
- #[inline]
- #[allow(non_snake_case)]
- pub fn type_as_interval(&self) -> Option<Interval<'a>> {
- if self.type_type() == Type::Interval {
- let u = self.type_();
- Some(Interval::init_from_table(u))
- } else {
- None
- }
- }
-
- #[inline]
- #[allow(non_snake_case)]
- pub fn type_as_list(&self) -> Option<List<'a>> {
- if self.type_type() == Type::List {
- let u = self.type_();
- Some(List::init_from_table(u))
- } else {
- None
- }
- }
-
- #[inline]
- #[allow(non_snake_case)]
- pub fn type_as_struct_(&self) -> Option<Struct_<'a>> {
- if self.type_type() == Type::Struct_ {
- let u = self.type_();
- Some(Struct_::init_from_table(u))
- } else {
- None
- }
- }
-
- #[inline]
- #[allow(non_snake_case)]
- pub fn type_as_union(&self) -> Option<Union<'a>> {
- if self.type_type() == Type::Union {
- let u = self.type_();
- Some(Union::init_from_table(u))
- } else {
- None
- }
- }
-
- #[inline]
- #[allow(non_snake_case)]
- pub fn type_as_fixed_size_binary(&self) -> Option<FixedSizeBinary<'a>> {
- if self.type_type() == Type::FixedSizeBinary {
- let u = self.type_();
- Some(FixedSizeBinary::init_from_table(u))
- } else {
- None
- }
- }
-
- #[inline]
- #[allow(non_snake_case)]
- pub fn type_as_fixed_size_list(&self) -> Option<FixedSizeList<'a>> {
- if self.type_type() == Type::FixedSizeList {
- let u = self.type_();
- Some(FixedSizeList::init_from_table(u))
- } else {
- None
- }
- }
-
- #[inline]
- #[allow(non_snake_case)]
- pub fn type_as_map(&self) -> Option<Map<'a>> {
- if self.type_type() == Type::Map {
- let u = self.type_();
- Some(Map::init_from_table(u))
- } else {
- None
- }
- }
-
- #[inline]
- #[allow(non_snake_case)]
- pub fn type_as_duration(&self) -> Option<Duration<'a>> {
- if self.type_type() == Type::Duration {
- let u = self.type_();
- Some(Duration::init_from_table(u))
- } else {
- None
- }
- }
-
- #[inline]
- #[allow(non_snake_case)]
- pub fn type_as_large_binary(&self) -> Option<LargeBinary<'a>> {
- if self.type_type() == Type::LargeBinary {
- let u = self.type_();
- Some(LargeBinary::init_from_table(u))
- } else {
- None
- }
- }
-
- #[inline]
- #[allow(non_snake_case)]
- pub fn type_as_large_utf_8(&self) -> Option<LargeUtf8<'a>> {
- if self.type_type() == Type::LargeUtf8 {
- let u = self.type_();
- Some(LargeUtf8::init_from_table(u))
- } else {
- None
- }
- }
-
- #[inline]
- #[allow(non_snake_case)]
- pub fn type_as_large_list(&self) -> Option<LargeList<'a>> {
- if self.type_type() == Type::LargeList {
- let u = self.type_();
- Some(LargeList::init_from_table(u))
- } else {
- None
- }
- }
-}
-
-impl flatbuffers::Verifiable for Tensor<'_> {
- #[inline]
- fn run_verifier(
- v: &mut flatbuffers::Verifier,
- pos: usize,
- ) -> Result<(), flatbuffers::InvalidFlatbuffer> {
- use flatbuffers::Verifiable;
- v.visit_table(pos)?
- .visit_union::<Type, _>(&"type_type", Self::VT_TYPE_TYPE, &"type_", Self::VT_TYPE_, true, |key, v, pos| {
- match key {
- Type::Null => v.verify_union_variant::<flatbuffers::ForwardsUOffset<Null>>("Type::Null", pos),
- Type::Int => v.verify_union_variant::<flatbuffers::ForwardsUOffset<Int>>("Type::Int", pos),
- Type::FloatingPoint => v.verify_union_variant::<flatbuffers::ForwardsUOffset<FloatingPoint>>("Type::FloatingPoint", pos),
- Type::Binary => v.verify_union_variant::<flatbuffers::ForwardsUOffset<Binary>>("Type::Binary", pos),
- Type::Utf8 => v.verify_union_variant::<flatbuffers::ForwardsUOffset<Utf8>>("Type::Utf8", pos),
- Type::Bool => v.verify_union_variant::<flatbuffers::ForwardsUOffset<Bool>>("Type::Bool", pos),
- Type::Decimal => v.verify_union_variant::<flatbuffers::ForwardsUOffset<Decimal>>("Type::Decimal", pos),
- Type::Date => v.verify_union_variant::<flatbuffers::ForwardsUOffset<Date>>("Type::Date", pos),
- Type::Time => v.verify_union_variant::<flatbuffers::ForwardsUOffset<Time>>("Type::Time", pos),
- Type::Timestamp => v.verify_union_variant::<flatbuffers::ForwardsUOffset<Timestamp>>("Type::Timestamp", pos),
- Type::Interval => v.verify_union_variant::<flatbuffers::ForwardsUOffset<Interval>>("Type::Interval", pos),
- Type::List => v.verify_union_variant::<flatbuffers::ForwardsUOffset<List>>("Type::List", pos),
- Type::Struct_ => v.verify_union_variant::<flatbuffers::ForwardsUOffset<Struct_>>("Type::Struct_", pos),
- Type::Union => v.verify_union_variant::<flatbuffers::ForwardsUOffset<Union>>("Type::Union", pos),
- Type::FixedSizeBinary => v.verify_union_variant::<flatbuffers::ForwardsUOffset<FixedSizeBinary>>("Type::FixedSizeBinary", pos),
- Type::FixedSizeList => v.verify_union_variant::<flatbuffers::ForwardsUOffset<FixedSizeList>>("Type::FixedSizeList", pos),
- Type::Map => v.verify_union_variant::<flatbuffers::ForwardsUOffset<Map>>("Type::Map", pos),
- Type::Duration => v.verify_union_variant::<flatbuffers::ForwardsUOffset<Duration>>("Type::Duration", pos),
- Type::LargeBinary => v.verify_union_variant::<flatbuffers::ForwardsUOffset<LargeBinary>>("Type::LargeBinary", pos),
- Type::LargeUtf8 => v.verify_union_variant::<flatbuffers::ForwardsUOffset<LargeUtf8>>("Type::LargeUtf8", pos),
- Type::LargeList => v.verify_union_variant::<flatbuffers::ForwardsUOffset<LargeList>>("Type::LargeList", pos),
- _ => Ok(()),
- }
- })?
- .visit_field::<flatbuffers::ForwardsUOffset<flatbuffers::Vector<'_, flatbuffers::ForwardsUOffset<TensorDim>>>>(&"shape", Self::VT_SHAPE, true)?
- .visit_field::<flatbuffers::ForwardsUOffset<flatbuffers::Vector<'_, i64>>>(&"strides", Self::VT_STRIDES, false)?
- .visit_field::<Buffer>(&"data", Self::VT_DATA, true)?
- .finish();
- Ok(())
- }
-}
-pub struct TensorArgs<'a> {
- pub type_type: Type,
- pub type_: Option<flatbuffers::WIPOffset<flatbuffers::UnionWIPOffset>>,
- pub shape: Option<
- flatbuffers::WIPOffset<
- flatbuffers::Vector<'a, flatbuffers::ForwardsUOffset<TensorDim<'a>>>,
- >,
- >,
- pub strides: Option<flatbuffers::WIPOffset<flatbuffers::Vector<'a, i64>>>,
- pub data: Option<&'a Buffer>,
-}
-impl<'a> Default for TensorArgs<'a> {
- #[inline]
- fn default() -> Self {
- TensorArgs {
- type_type: Type::NONE,
- type_: None, // required field
- shape: None, // required field
- strides: None,
- data: None, // required field
- }
- }
-}
-pub struct TensorBuilder<'a: 'b, 'b> {
- fbb_: &'b mut flatbuffers::FlatBufferBuilder<'a>,
- start_: flatbuffers::WIPOffset<flatbuffers::TableUnfinishedWIPOffset>,
-}
-impl<'a: 'b, 'b> TensorBuilder<'a, 'b> {
- #[inline]
- pub fn add_type_type(&mut self, type_type: Type) {
- self.fbb_
- .push_slot::<Type>(Tensor::VT_TYPE_TYPE, type_type, Type::NONE);
- }
- #[inline]
- pub fn add_type_(
- &mut self,
- type_: flatbuffers::WIPOffset<flatbuffers::UnionWIPOffset>,
- ) {
- self.fbb_
- .push_slot_always::<flatbuffers::WIPOffset<_>>(Tensor::VT_TYPE_, type_);
- }
- #[inline]
- pub fn add_shape(
- &mut self,
- shape: flatbuffers::WIPOffset<
- flatbuffers::Vector<'b, flatbuffers::ForwardsUOffset<TensorDim<'b>>>,
- >,
- ) {
- self.fbb_
- .push_slot_always::<flatbuffers::WIPOffset<_>>(Tensor::VT_SHAPE, shape);
- }
- #[inline]
- pub fn add_strides(
- &mut self,
- strides: flatbuffers::WIPOffset<flatbuffers::Vector<'b, i64>>,
- ) {
- self.fbb_
- .push_slot_always::<flatbuffers::WIPOffset<_>>(Tensor::VT_STRIDES, strides);
- }
- #[inline]
- pub fn add_data(&mut self, data: &Buffer) {
- self.fbb_.push_slot_always::<&Buffer>(Tensor::VT_DATA, data);
- }
- #[inline]
- pub fn new(
- _fbb: &'b mut flatbuffers::FlatBufferBuilder<'a>,
- ) -> TensorBuilder<'a, 'b> {
- let start = _fbb.start_table();
- TensorBuilder {
- fbb_: _fbb,
- start_: start,
- }
- }
- #[inline]
- pub fn finish(self) -> flatbuffers::WIPOffset<Tensor<'a>> {
- let o = self.fbb_.end_table(self.start_);
- self.fbb_.required(o, Tensor::VT_TYPE_, "type_");
- self.fbb_.required(o, Tensor::VT_SHAPE, "shape");
- self.fbb_.required(o, Tensor::VT_DATA, "data");
- flatbuffers::WIPOffset::new(o.value())
- }
-}
-
-impl std::fmt::Debug for Tensor<'_> {
- fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
- let mut ds = f.debug_struct("Tensor");
- ds.field("type_type", &self.type_type());
- match self.type_type() {
- Type::Null => {
- if let Some(x) = self.type_as_null() {
- ds.field("type_", &x)
- } else {
- ds.field(
- "type_",
- &"InvalidFlatbuffer: Union discriminant does not match value.",
- )
- }
- }
- Type::Int => {
- if let Some(x) = self.type_as_int() {
- ds.field("type_", &x)
- } else {
- ds.field(
- "type_",
- &"InvalidFlatbuffer: Union discriminant does not match value.",
- )
- }
- }
- Type::FloatingPoint => {
- if let Some(x) = self.type_as_floating_point() {
- ds.field("type_", &x)
- } else {
- ds.field(
- "type_",
- &"InvalidFlatbuffer: Union discriminant does not match value.",
- )
- }
- }
- Type::Binary => {
- if let Some(x) = self.type_as_binary() {
- ds.field("type_", &x)
- } else {
- ds.field(
- "type_",
- &"InvalidFlatbuffer: Union discriminant does not match value.",
- )
- }
- }
- Type::Utf8 => {
- if let Some(x) = self.type_as_utf_8() {
- ds.field("type_", &x)
- } else {
- ds.field(
- "type_",
- &"InvalidFlatbuffer: Union discriminant does not match value.",
- )
- }
- }
- Type::Bool => {
- if let Some(x) = self.type_as_bool() {
- ds.field("type_", &x)
- } else {
- ds.field(
- "type_",
- &"InvalidFlatbuffer: Union discriminant does not match value.",
- )
- }
- }
- Type::Decimal => {
- if let Some(x) = self.type_as_decimal() {
- ds.field("type_", &x)
- } else {
- ds.field(
- "type_",
- &"InvalidFlatbuffer: Union discriminant does not match value.",
- )
- }
- }
- Type::Date => {
- if let Some(x) = self.type_as_date() {
- ds.field("type_", &x)
- } else {
- ds.field(
- "type_",
- &"InvalidFlatbuffer: Union discriminant does not match value.",
- )
- }
- }
- Type::Time => {
- if let Some(x) = self.type_as_time() {
- ds.field("type_", &x)
- } else {
- ds.field(
- "type_",
- &"InvalidFlatbuffer: Union discriminant does not match value.",
- )
- }
- }
- Type::Timestamp => {
- if let Some(x) = self.type_as_timestamp() {
- ds.field("type_", &x)
- } else {
- ds.field(
- "type_",
- &"InvalidFlatbuffer: Union discriminant does not match value.",
- )
- }
- }
- Type::Interval => {
- if let Some(x) = self.type_as_interval() {
- ds.field("type_", &x)
- } else {
- ds.field(
- "type_",
- &"InvalidFlatbuffer: Union discriminant does not match value.",
- )
- }
- }
- Type::List => {
- if let Some(x) = self.type_as_list() {
- ds.field("type_", &x)
- } else {
- ds.field(
- "type_",
- &"InvalidFlatbuffer: Union discriminant does not match value.",
- )
- }
- }
- Type::Struct_ => {
- if let Some(x) = self.type_as_struct_() {
- ds.field("type_", &x)
- } else {
- ds.field(
- "type_",
- &"InvalidFlatbuffer: Union discriminant does not match value.",
- )
- }
- }
- Type::Union => {
- if let Some(x) = self.type_as_union() {
- ds.field("type_", &x)
- } else {
- ds.field(
- "type_",
- &"InvalidFlatbuffer: Union discriminant does not match value.",
- )
- }
- }
- Type::FixedSizeBinary => {
- if let Some(x) = self.type_as_fixed_size_binary() {
- ds.field("type_", &x)
- } else {
- ds.field(
- "type_",
- &"InvalidFlatbuffer: Union discriminant does not match value.",
- )
- }
- }
- Type::FixedSizeList => {
- if let Some(x) = self.type_as_fixed_size_list() {
- ds.field("type_", &x)
- } else {
- ds.field(
- "type_",
- &"InvalidFlatbuffer: Union discriminant does not match value.",
- )
- }
- }
- Type::Map => {
- if let Some(x) = self.type_as_map() {
- ds.field("type_", &x)
- } else {
- ds.field(
- "type_",
- &"InvalidFlatbuffer: Union discriminant does not match value.",
- )
- }
- }
- Type::Duration => {
- if let Some(x) = self.type_as_duration() {
- ds.field("type_", &x)
- } else {
- ds.field(
- "type_",
- &"InvalidFlatbuffer: Union discriminant does not match value.",
- )
- }
- }
- Type::LargeBinary => {
- if let Some(x) = self.type_as_large_binary() {
- ds.field("type_", &x)
- } else {
- ds.field(
- "type_",
- &"InvalidFlatbuffer: Union discriminant does not match value.",
- )
- }
- }
- Type::LargeUtf8 => {
- if let Some(x) = self.type_as_large_utf_8() {
- ds.field("type_", &x)
- } else {
- ds.field(
- "type_",
- &"InvalidFlatbuffer: Union discriminant does not match value.",
- )
- }
- }
- Type::LargeList => {
- if let Some(x) = self.type_as_large_list() {
- ds.field("type_", &x)
- } else {
- ds.field(
- "type_",
- &"InvalidFlatbuffer: Union discriminant does not match value.",
- )
- }
- }
- _ => {
- let x: Option<()> = None;
- ds.field("type_", &x)
- }
- };
- ds.field("shape", &self.shape());
- ds.field("strides", &self.strides());
- ds.field("data", &self.data());
- ds.finish()
- }
-}
-#[inline]
-#[deprecated(since = "2.0.0", note = "Deprecated in favor of `root_as...` methods.")]
-pub fn get_root_as_tensor<'a>(buf: &'a [u8]) -> Tensor<'a> {
- unsafe { flatbuffers::root_unchecked::<Tensor<'a>>(buf) }
-}
-
-#[inline]
-#[deprecated(since = "2.0.0", note = "Deprecated in favor of `root_as...` methods.")]
-pub fn get_size_prefixed_root_as_tensor<'a>(buf: &'a [u8]) -> Tensor<'a> {
- unsafe { flatbuffers::size_prefixed_root_unchecked::<Tensor<'a>>(buf) }
-}
-
-#[inline]
-/// Verifies that a buffer of bytes contains a `Tensor`
-/// and returns it.
-/// Note that verification is still experimental and may not
-/// catch every error, or be maximally performant. For the
-/// previous, unchecked, behavior use
-/// `root_as_tensor_unchecked`.
-pub fn root_as_tensor(buf: &[u8]) -> Result<Tensor, flatbuffers::InvalidFlatbuffer> {
- flatbuffers::root::<Tensor>(buf)
-}
-#[inline]
-/// Verifies that a buffer of bytes contains a size prefixed
-/// `Tensor` and returns it.
-/// Note that verification is still experimental and may not
-/// catch every error, or be maximally performant. For the
-/// previous, unchecked, behavior use
-/// `size_prefixed_root_as_tensor_unchecked`.
-pub fn size_prefixed_root_as_tensor(
- buf: &[u8],
-) -> Result<Tensor, flatbuffers::InvalidFlatbuffer> {
- flatbuffers::size_prefixed_root::<Tensor>(buf)
-}
-#[inline]
-/// Verifies, with the given options, that a buffer of bytes
-/// contains a `Tensor` and returns it.
-/// Note that verification is still experimental and may not
-/// catch every error, or be maximally performant. For the
-/// previous, unchecked, behavior use
-/// `root_as_tensor_unchecked`.
-pub fn root_as_tensor_with_opts<'b, 'o>(
- opts: &'o flatbuffers::VerifierOptions,
- buf: &'b [u8],
-) -> Result<Tensor<'b>, flatbuffers::InvalidFlatbuffer> {
- flatbuffers::root_with_opts::<Tensor<'b>>(opts, buf)
-}
-#[inline]
-/// Verifies, with the given verifier options, that a buffer of
-/// bytes contains a size prefixed `Tensor` and returns
-/// it. Note that verification is still experimental and may not
-/// catch every error, or be maximally performant. For the
-/// previous, unchecked, behavior use
-/// `root_as_tensor_unchecked`.
-pub fn size_prefixed_root_as_tensor_with_opts<'b, 'o>(
- opts: &'o flatbuffers::VerifierOptions,
- buf: &'b [u8],
-) -> Result<Tensor<'b>, flatbuffers::InvalidFlatbuffer> {
- flatbuffers::size_prefixed_root_with_opts::<Tensor<'b>>(opts, buf)
-}
-#[inline]
-/// Assumes, without verification, that a buffer of bytes contains a Tensor and returns it.
-/// # Safety
-/// Callers must trust the given bytes do indeed contain a valid `Tensor`.
-pub unsafe fn root_as_tensor_unchecked(buf: &[u8]) -> Tensor {
- flatbuffers::root_unchecked::<Tensor>(buf)
-}
-#[inline]
-/// Assumes, without verification, that a buffer of bytes contains a size prefixed Tensor and returns it.
-/// # Safety
-/// Callers must trust the given bytes do indeed contain a valid size prefixed `Tensor`.
-pub unsafe fn size_prefixed_root_as_tensor_unchecked(buf: &[u8]) -> Tensor {
- flatbuffers::size_prefixed_root_unchecked::<Tensor>(buf)
-}
-#[inline]
-pub fn finish_tensor_buffer<'a, 'b>(
- fbb: &'b mut flatbuffers::FlatBufferBuilder<'a>,
- root: flatbuffers::WIPOffset<Tensor<'a>>,
-) {
- fbb.finish(root, None);
-}
-
-#[inline]
-pub fn finish_size_prefixed_tensor_buffer<'a, 'b>(
- fbb: &'b mut flatbuffers::FlatBufferBuilder<'a>,
- root: flatbuffers::WIPOffset<Tensor<'a>>,
-) {
- fbb.finish_size_prefixed(root, None);
-}
diff --git a/arrow/src/ipc/gen/mod.rs b/arrow/src/ipc/gen/mod.rs
deleted file mode 100644
index ceeb6b2..0000000
--- a/arrow/src/ipc/gen/mod.rs
+++ /dev/null
@@ -1,31 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-//! Generated code
-
-#![allow(non_snake_case)]
-
-#[allow(clippy::all)]
-pub mod File;
-#[allow(clippy::all)]
-pub mod Message;
-#[allow(clippy::all)]
-pub mod Schema;
-#[allow(clippy::all)]
-pub mod SparseTensor;
-#[allow(clippy::all)]
-pub mod Tensor;
diff --git a/arrow/src/ipc/mod.rs b/arrow/src/ipc/mod.rs
deleted file mode 100644
index a2d7103..0000000
--- a/arrow/src/ipc/mod.rs
+++ /dev/null
@@ -1,39 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-// TODO: (vcq): Protobuf codegen is not generating Debug impls.
-#![allow(missing_debug_implementations)]
-
-pub mod convert;
-pub mod reader;
-pub mod writer;
-
-#[allow(clippy::redundant_closure)]
-#[allow(clippy::needless_lifetimes)]
-#[allow(clippy::extra_unused_lifetimes)]
-#[allow(clippy::redundant_static_lifetimes)]
-#[allow(clippy::redundant_field_names)]
-pub mod gen;
-
-pub use self::gen::File::*;
-pub use self::gen::Message::*;
-pub use self::gen::Schema::*;
-pub use self::gen::SparseTensor::*;
-pub use self::gen::Tensor::*;
-
-const ARROW_MAGIC: [u8; 6] = [b'A', b'R', b'R', b'O', b'W', b'1'];
-const CONTINUATION_MARKER: [u8; 4] = [0xff; 4];
diff --git a/arrow/src/ipc/reader.rs b/arrow/src/ipc/reader.rs
deleted file mode 100644
index 3c893cd..0000000
--- a/arrow/src/ipc/reader.rs
+++ /dev/null
@@ -1,1160 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-//! Arrow IPC File and Stream Readers
-//!
-//! The `FileReader` and `StreamReader` have similar interfaces,
-//! however the `FileReader` expects a reader that supports `Seek`ing
-
-use std::collections::HashMap;
-use std::io::{BufReader, Read, Seek, SeekFrom};
-use std::sync::Arc;
-
-use crate::array::*;
-use crate::buffer::Buffer;
-use crate::compute::cast;
-use crate::datatypes::{DataType, Field, IntervalUnit, Schema, SchemaRef};
-use crate::error::{ArrowError, Result};
-use crate::ipc;
-use crate::record_batch::{RecordBatch, RecordBatchReader};
-
-use ipc::CONTINUATION_MARKER;
-use DataType::*;
-
-/// Read a buffer based on offset and length
-fn read_buffer(buf: &ipc::Buffer, a_data: &[u8]) -> Buffer {
- let start_offset = buf.offset() as usize;
- let end_offset = start_offset + buf.length() as usize;
- let buf_data = &a_data[start_offset..end_offset];
- Buffer::from(&buf_data)
-}
-
-/// Coordinates reading arrays based on data types.
-///
-/// Notes:
-/// * In the IPC format, null buffers are always set, but may be empty. We discard them if an array has 0 nulls
-/// * Numeric values inside list arrays are often stored as 64-bit values regardless of their data type size.
-/// We thus:
-/// - check if the bit width of non-64-bit numbers is 64, and
-/// - read the buffer as 64-bit (signed integer or float), and
-/// - cast the 64-bit array to the appropriate data type
-fn create_array(
- nodes: &[ipc::FieldNode],
- data_type: &DataType,
- data: &[u8],
- buffers: &[ipc::Buffer],
- dictionaries: &[Option<ArrayRef>],
- mut node_index: usize,
- mut buffer_index: usize,
-) -> (ArrayRef, usize, usize) {
- use DataType::*;
- let array = match data_type {
- Utf8 | Binary | LargeBinary | LargeUtf8 => {
- let array = create_primitive_array(
- &nodes[node_index],
- data_type,
- buffers[buffer_index..buffer_index + 3]
- .iter()
- .map(|buf| read_buffer(buf, data))
- .collect(),
- );
- node_index += 1;
- buffer_index += 3;
- array
- }
- FixedSizeBinary(_) => {
- let array = create_primitive_array(
- &nodes[node_index],
- data_type,
- buffers[buffer_index..buffer_index + 2]
- .iter()
- .map(|buf| read_buffer(buf, data))
- .collect(),
- );
- node_index += 1;
- buffer_index += 2;
- array
- }
- List(ref list_field) | LargeList(ref list_field) => {
- let list_node = &nodes[node_index];
- let list_buffers: Vec<Buffer> = buffers[buffer_index..buffer_index + 2]
- .iter()
- .map(|buf| read_buffer(buf, data))
- .collect();
- node_index += 1;
- buffer_index += 2;
- let triple = create_array(
- nodes,
- list_field.data_type(),
- data,
- buffers,
- dictionaries,
- node_index,
- buffer_index,
- );
- node_index = triple.1;
- buffer_index = triple.2;
-
- create_list_array(list_node, data_type, &list_buffers[..], triple.0)
- }
- FixedSizeList(ref list_field, _) => {
- let list_node = &nodes[node_index];
- let list_buffers: Vec<Buffer> = buffers[buffer_index..=buffer_index]
- .iter()
- .map(|buf| read_buffer(buf, data))
- .collect();
- node_index += 1;
- buffer_index += 1;
- let triple = create_array(
- nodes,
- list_field.data_type(),
- data,
- buffers,
- dictionaries,
- node_index,
- buffer_index,
- );
- node_index = triple.1;
- buffer_index = triple.2;
-
- create_list_array(list_node, data_type, &list_buffers[..], triple.0)
- }
- Struct(struct_fields) => {
- let struct_node = &nodes[node_index];
- let null_buffer: Buffer = read_buffer(&buffers[buffer_index], data);
- node_index += 1;
- buffer_index += 1;
-
- // read the arrays for each field
- let mut struct_arrays = vec![];
- // TODO investigate whether just knowing the number of buffers could
- // still work
- for struct_field in struct_fields {
- let triple = create_array(
- nodes,
- struct_field.data_type(),
- data,
- buffers,
- dictionaries,
- node_index,
- buffer_index,
- );
- node_index = triple.1;
- buffer_index = triple.2;
- struct_arrays.push((struct_field.clone(), triple.0));
- }
- let null_count = struct_node.null_count() as usize;
- let struct_array = if null_count > 0 {
- // create struct array from fields, arrays and null data
- StructArray::from((struct_arrays, null_buffer))
- } else {
- StructArray::from(struct_arrays)
- };
- Arc::new(struct_array)
- }
- // Create dictionary array from RecordBatch
- Dictionary(_, _) => {
- let index_node = &nodes[node_index];
- let index_buffers: Vec<Buffer> = buffers[buffer_index..buffer_index + 2]
- .iter()
- .map(|buf| read_buffer(buf, data))
- .collect();
- let value_array = dictionaries[node_index].clone().unwrap();
- node_index += 1;
- buffer_index += 2;
-
- create_dictionary_array(
- index_node,
- data_type,
- &index_buffers[..],
- value_array,
- )
- }
- Null => {
- let length = nodes[node_index].length() as usize;
- let data = ArrayData::builder(data_type.clone())
- .len(length)
- .offset(0)
- .build();
- node_index += 1;
- // no buffer increases
- make_array(data)
- }
- _ => {
- let array = create_primitive_array(
- &nodes[node_index],
- data_type,
- buffers[buffer_index..buffer_index + 2]
- .iter()
- .map(|buf| read_buffer(buf, data))
- .collect(),
- );
- node_index += 1;
- buffer_index += 2;
- array
- }
- };
- (array, node_index, buffer_index)
-}
-
-/// Reads the correct number of buffers based on data type and null_count, and creates a
-/// primitive array ref
-fn create_primitive_array(
- field_node: &ipc::FieldNode,
- data_type: &DataType,
- buffers: Vec<Buffer>,
-) -> ArrayRef {
- let length = field_node.length() as usize;
- let null_count = field_node.null_count() as usize;
- let array_data = match data_type {
- Utf8 | Binary | LargeBinary | LargeUtf8 => {
- // read 3 buffers
- let mut builder = ArrayData::builder(data_type.clone())
- .len(length)
- .buffers(buffers[1..3].to_vec())
- .offset(0);
- if null_count > 0 {
- builder = builder.null_bit_buffer(buffers[0].clone())
- }
- builder.build()
- }
- FixedSizeBinary(_) => {
- // read 3 buffers
- let mut builder = ArrayData::builder(data_type.clone())
- .len(length)
- .buffers(buffers[1..2].to_vec())
- .offset(0);
- if null_count > 0 {
- builder = builder.null_bit_buffer(buffers[0].clone())
- }
- builder.build()
- }
- Int8
- | Int16
- | Int32
- | UInt8
- | UInt16
- | UInt32
- | Time32(_)
- | Date32
- | Interval(IntervalUnit::YearMonth) => {
- if buffers[1].len() / 8 == length && length != 1 {
- // interpret as a signed i64, and cast appropriately
- let mut builder = ArrayData::builder(DataType::Int64)
- .len(length)
- .buffers(buffers[1..].to_vec())
- .offset(0);
- if null_count > 0 {
- builder = builder.null_bit_buffer(buffers[0].clone())
- }
- let values = Arc::new(Int64Array::from(builder.build())) as ArrayRef;
- // this cast is infallible, the unwrap is safe
- let casted = cast(&values, data_type).unwrap();
- casted.data().clone()
- } else {
- let mut builder = ArrayData::builder(data_type.clone())
- .len(length)
- .buffers(buffers[1..].to_vec())
- .offset(0);
- if null_count > 0 {
- builder = builder.null_bit_buffer(buffers[0].clone())
- }
- builder.build()
- }
- }
- Float32 => {
- if buffers[1].len() / 8 == length && length != 1 {
- // interpret as a f64, and cast appropriately
- let mut builder = ArrayData::builder(DataType::Float64)
- .len(length)
- .buffers(buffers[1..].to_vec())
- .offset(0);
- if null_count > 0 {
- builder = builder.null_bit_buffer(buffers[0].clone())
- }
- let values = Arc::new(Float64Array::from(builder.build())) as ArrayRef;
- // this cast is infallible, the unwrap is safe
- let casted = cast(&values, data_type).unwrap();
- casted.data().clone()
- } else {
- let mut builder = ArrayData::builder(data_type.clone())
- .len(length)
- .buffers(buffers[1..].to_vec())
- .offset(0);
- if null_count > 0 {
- builder = builder.null_bit_buffer(buffers[0].clone())
- }
- builder.build()
- }
- }
- Boolean
- | Int64
- | UInt64
- | Float64
- | Time64(_)
- | Timestamp(_, _)
- | Date64
- | Duration(_)
- | Interval(IntervalUnit::DayTime) => {
- let mut builder = ArrayData::builder(data_type.clone())
- .len(length)
- .buffers(buffers[1..].to_vec())
- .offset(0);
- if null_count > 0 {
- builder = builder.null_bit_buffer(buffers[0].clone())
- }
- builder.build()
- }
- Decimal(_, _) => {
- // read 3 buffers
- let mut builder = ArrayData::builder(data_type.clone())
- .len(length)
- .buffers(buffers[1..2].to_vec())
- .offset(0);
- if null_count > 0 {
- builder = builder.null_bit_buffer(buffers[0].clone())
- }
- builder.build()
- }
- t => panic!("Data type {:?} either unsupported or not primitive", t),
- };
-
- make_array(array_data)
-}
-
-/// Reads the correct number of buffers based on list type and null_count, and creates a
-/// list array ref
-fn create_list_array(
- field_node: &ipc::FieldNode,
- data_type: &DataType,
- buffers: &[Buffer],
- child_array: ArrayRef,
-) -> ArrayRef {
- if let DataType::List(_) = *data_type {
- let null_count = field_node.null_count() as usize;
- let mut builder = ArrayData::builder(data_type.clone())
- .len(field_node.length() as usize)
- .buffers(buffers[1..2].to_vec())
- .offset(0)
- .child_data(vec![child_array.data().clone()]);
- if null_count > 0 {
- builder = builder.null_bit_buffer(buffers[0].clone())
- }
- make_array(builder.build())
- } else if let DataType::LargeList(_) = *data_type {
- let null_count = field_node.null_count() as usize;
- let mut builder = ArrayData::builder(data_type.clone())
- .len(field_node.length() as usize)
- .buffers(buffers[1..2].to_vec())
- .offset(0)
- .child_data(vec![child_array.data().clone()]);
- if null_count > 0 {
- builder = builder.null_bit_buffer(buffers[0].clone())
- }
- make_array(builder.build())
- } else if let DataType::FixedSizeList(_, _) = *data_type {
- let null_count = field_node.null_count() as usize;
- let mut builder = ArrayData::builder(data_type.clone())
- .len(field_node.length() as usize)
- .buffers(buffers[1..1].to_vec())
- .offset(0)
- .child_data(vec![child_array.data().clone()]);
- if null_count > 0 {
- builder = builder.null_bit_buffer(buffers[0].clone())
- }
- make_array(builder.build())
- } else {
- panic!("Cannot create list array from {:?}", data_type)
- }
-}
-
-/// Reads the correct number of buffers based on list type and null_count, and creates a
-/// list array ref
-fn create_dictionary_array(
- field_node: &ipc::FieldNode,
- data_type: &DataType,
- buffers: &[Buffer],
- value_array: ArrayRef,
-) -> ArrayRef {
- if let DataType::Dictionary(_, _) = *data_type {
- let null_count = field_node.null_count() as usize;
- let mut builder = ArrayData::builder(data_type.clone())
- .len(field_node.length() as usize)
- .buffers(buffers[1..2].to_vec())
- .offset(0)
- .child_data(vec![value_array.data().clone()]);
- if null_count > 0 {
- builder = builder.null_bit_buffer(buffers[0].clone())
- }
- make_array(builder.build())
- } else {
- unreachable!("Cannot create dictionary array from {:?}", data_type)
- }
-}
-
-/// Creates a record batch from binary data using the `ipc::RecordBatch` indexes and the `Schema`
-pub fn read_record_batch(
- buf: &[u8],
- batch: ipc::RecordBatch,
- schema: SchemaRef,
- dictionaries: &[Option<ArrayRef>],
-) -> Result<RecordBatch> {
- let buffers = batch.buffers().ok_or_else(|| {
- ArrowError::IoError("Unable to get buffers from IPC RecordBatch".to_string())
- })?;
- let field_nodes = batch.nodes().ok_or_else(|| {
- ArrowError::IoError("Unable to get field nodes from IPC RecordBatch".to_string())
- })?;
- // keep track of buffer and node index, the functions that create arrays mutate these
- let mut buffer_index = 0;
- let mut node_index = 0;
- let mut arrays = vec![];
-
- // keep track of index as lists require more than one node
- for field in schema.fields() {
- let triple = create_array(
- field_nodes,
- field.data_type(),
- &buf,
- buffers,
- dictionaries,
- node_index,
- buffer_index,
- );
- node_index = triple.1;
- buffer_index = triple.2;
- arrays.push(triple.0);
- }
-
- RecordBatch::try_new(schema, arrays)
-}
-
-/// Read the dictionary from the buffer and provided metadata,
-/// updating the `dictionaries_by_field` with the resulting dictionary
-pub fn read_dictionary(
- buf: &[u8],
- batch: ipc::DictionaryBatch,
- schema: &Schema,
- dictionaries_by_field: &mut [Option<ArrayRef>],
-) -> Result<()> {
- if batch.isDelta() {
- return Err(ArrowError::IoError(
- "delta dictionary batches not supported".to_string(),
- ));
- }
-
- let id = batch.id();
- let fields_using_this_dictionary = schema.fields_with_dict_id(id);
- let first_field = fields_using_this_dictionary.first().ok_or_else(|| {
- ArrowError::InvalidArgumentError("dictionary id not found in schema".to_string())
- })?;
-
- // As the dictionary batch does not contain the type of the
- // values array, we need to retrieve this from the schema.
- // Get an array representing this dictionary's values.
- let dictionary_values: ArrayRef = match first_field.data_type() {
- DataType::Dictionary(_, ref value_type) => {
- // Make a fake schema for the dictionary batch.
- let schema = Schema {
- fields: vec![Field::new("", value_type.as_ref().clone(), false)],
- metadata: HashMap::new(),
- };
- // Read a single column
- let record_batch = read_record_batch(
- &buf,
- batch.data().unwrap(),
- Arc::new(schema),
- &dictionaries_by_field,
- )?;
- Some(record_batch.column(0).clone())
- }
- _ => None,
- }
- .ok_or_else(|| {
- ArrowError::InvalidArgumentError("dictionary id not found in schema".to_string())
- })?;
-
- // for all fields with this dictionary id, update the dictionaries vector
- // in the reader. Note that a dictionary batch may be shared between many fields.
- // We don't currently record the isOrdered field. This could be general
- // attributes of arrays.
- for (i, field) in schema.fields().iter().enumerate() {
- if field.dict_id() == Some(id) {
- // Add (possibly multiple) array refs to the dictionaries array.
- dictionaries_by_field[i] = Some(dictionary_values.clone());
- }
- }
-
- Ok(())
-}
-
-/// Arrow File reader
-pub struct FileReader<R: Read + Seek> {
- /// Buffered file reader that supports reading and seeking
- reader: BufReader<R>,
-
- /// The schema that is read from the file header
- schema: SchemaRef,
-
- /// The blocks in the file
- ///
- /// A block indicates the regions in the file to read to get data
- blocks: Vec<ipc::Block>,
-
- /// A counter to keep track of the current block that should be read
- current_block: usize,
-
- /// The total number of blocks, which may contain record batches and other types
- total_blocks: usize,
-
- /// Optional dictionaries for each schema field.
- ///
- /// Dictionaries may be appended to in the streaming format.
- dictionaries_by_field: Vec<Option<ArrayRef>>,
-
- /// Metadata version
- metadata_version: ipc::MetadataVersion,
-}
-
-impl<R: Read + Seek> FileReader<R> {
- /// Try to create a new file reader
- ///
- /// Returns errors if the file does not meet the Arrow Format header and footer
- /// requirements
- pub fn try_new(reader: R) -> Result<Self> {
- let mut reader = BufReader::new(reader);
- // check if header and footer contain correct magic bytes
- let mut magic_buffer: [u8; 6] = [0; 6];
- reader.read_exact(&mut magic_buffer)?;
- if magic_buffer != super::ARROW_MAGIC {
- return Err(ArrowError::IoError(
- "Arrow file does not contain correct header".to_string(),
- ));
- }
- reader.seek(SeekFrom::End(-6))?;
- reader.read_exact(&mut magic_buffer)?;
- if magic_buffer != super::ARROW_MAGIC {
- return Err(ArrowError::IoError(
- "Arrow file does not contain correct footer".to_string(),
- ));
- }
- // read footer length
- let mut footer_size: [u8; 4] = [0; 4];
- reader.seek(SeekFrom::End(-10))?;
- reader.read_exact(&mut footer_size)?;
- let footer_len = i32::from_le_bytes(footer_size);
-
- // read footer
- let mut footer_data = vec![0; footer_len as usize];
- reader.seek(SeekFrom::End(-10 - footer_len as i64))?;
- reader.read_exact(&mut footer_data)?;
-
- let footer = ipc::root_as_footer(&footer_data[..]).map_err(|err| {
- ArrowError::IoError(format!("Unable to get root as footer: {:?}", err))
- })?;
-
- let blocks = footer.recordBatches().ok_or_else(|| {
- ArrowError::IoError(
- "Unable to get record batches from IPC Footer".to_string(),
- )
- })?;
-
- let total_blocks = blocks.len();
-
- let ipc_schema = footer.schema().unwrap();
- let schema = ipc::convert::fb_to_schema(ipc_schema);
-
- // Create an array of optional dictionary value arrays, one per field.
- let mut dictionaries_by_field = vec![None; schema.fields().len()];
- for block in footer.dictionaries().unwrap() {
- // read length from end of offset
- let mut message_size: [u8; 4] = [0; 4];
- reader.seek(SeekFrom::Start(block.offset() as u64))?;
- reader.read_exact(&mut message_size)?;
- let footer_len = if message_size == CONTINUATION_MARKER {
- reader.read_exact(&mut message_size)?;
- i32::from_le_bytes(message_size)
- } else {
- i32::from_le_bytes(message_size)
- };
-
- let mut block_data = vec![0; footer_len as usize];
-
- reader.read_exact(&mut block_data)?;
-
- let message = ipc::root_as_message(&block_data[..]).map_err(|err| {
- ArrowError::IoError(format!("Unable to get root as message: {:?}", err))
- })?;
-
- match message.header_type() {
- ipc::MessageHeader::DictionaryBatch => {
- let batch = message.header_as_dictionary_batch().unwrap();
-
- // read the block that makes up the dictionary batch into a buffer
- let mut buf = vec![0; block.bodyLength() as usize];
- reader.seek(SeekFrom::Start(
- block.offset() as u64 + block.metaDataLength() as u64,
- ))?;
- reader.read_exact(&mut buf)?;
-
- read_dictionary(&buf, batch, &schema, &mut dictionaries_by_field)?;
- }
- t => {
- return Err(ArrowError::IoError(format!(
- "Expecting DictionaryBatch in dictionary blocks, found {:?}.",
- t
- )));
- }
- };
- }
-
- Ok(Self {
- reader,
- schema: Arc::new(schema),
- blocks: blocks.to_vec(),
- current_block: 0,
- total_blocks,
- dictionaries_by_field,
- metadata_version: footer.version(),
- })
- }
-
- /// Return the number of batches in the file
- pub fn num_batches(&self) -> usize {
- self.total_blocks
- }
-
- /// Return the schema of the file
- pub fn schema(&self) -> SchemaRef {
- self.schema.clone()
- }
-
- /// Read a specific record batch
- ///
- /// Sets the current block to the index, allowing random reads
- pub fn set_index(&mut self, index: usize) -> Result<()> {
- if index >= self.total_blocks {
- Err(ArrowError::IoError(format!(
- "Cannot set batch to index {} from {} total batches",
- index, self.total_blocks
- )))
- } else {
- self.current_block = index;
- Ok(())
- }
- }
-
- fn maybe_next(&mut self) -> Result<Option<RecordBatch>> {
- let block = self.blocks[self.current_block];
- self.current_block += 1;
-
- // read length
- self.reader.seek(SeekFrom::Start(block.offset() as u64))?;
- let mut meta_buf = [0; 4];
- self.reader.read_exact(&mut meta_buf)?;
- if meta_buf == CONTINUATION_MARKER {
- // continuation marker encountered, read message next
- self.reader.read_exact(&mut meta_buf)?;
- }
- let meta_len = i32::from_le_bytes(meta_buf);
-
- let mut block_data = vec![0; meta_len as usize];
- self.reader.read_exact(&mut block_data)?;
-
- let message = ipc::root_as_message(&block_data[..]).map_err(|err| {
- ArrowError::IoError(format!("Unable to get root as footer: {:?}", err))
- })?;
-
- // some old test data's footer metadata is not set, so we account for that
- if self.metadata_version != ipc::MetadataVersion::V1
- && message.version() != self.metadata_version
- {
- return Err(ArrowError::IoError(
- "Could not read IPC message as metadata versions mismatch".to_string(),
- ));
- }
-
- match message.header_type() {
- ipc::MessageHeader::Schema => Err(ArrowError::IoError(
- "Not expecting a schema when messages are read".to_string(),
- )),
- ipc::MessageHeader::RecordBatch => {
- let batch = message.header_as_record_batch().ok_or_else(|| {
- ArrowError::IoError(
- "Unable to read IPC message as record batch".to_string(),
- )
- })?;
- // read the block that makes up the record batch into a buffer
- let mut buf = vec![0; block.bodyLength() as usize];
- self.reader.seek(SeekFrom::Start(
- block.offset() as u64 + block.metaDataLength() as u64,
- ))?;
- self.reader.read_exact(&mut buf)?;
-
- read_record_batch(
- &buf,
- batch,
- self.schema(),
- &self.dictionaries_by_field,
- ).map(Some)
- }
- ipc::MessageHeader::NONE => {
- Ok(None)
- }
- t => Err(ArrowError::IoError(format!(
- "Reading types other than record batches not yet supported, unable to read {:?}", t
- ))),
- }
- }
-}
-
-impl<R: Read + Seek> Iterator for FileReader<R> {
- type Item = Result<RecordBatch>;
-
- fn next(&mut self) -> Option<Self::Item> {
- // get current block
- if self.current_block < self.total_blocks {
- self.maybe_next().transpose()
- } else {
- None
- }
- }
-}
-
-impl<R: Read + Seek> RecordBatchReader for FileReader<R> {
- fn schema(&self) -> SchemaRef {
- self.schema.clone()
- }
-}
-
-/// Arrow Stream reader
-pub struct StreamReader<R: Read> {
- /// Buffered stream reader
- reader: BufReader<R>,
-
- /// The schema that is read from the stream's first message
- schema: SchemaRef,
-
- /// Optional dictionaries for each schema field.
- ///
- /// Dictionaries may be appended to in the streaming format.
- dictionaries_by_field: Vec<Option<ArrayRef>>,
-
- /// An indicator of whether the stream is complete.
- ///
- /// This value is set to `true` the first time the reader's `next()` returns `None`.
- finished: bool,
-}
-
-impl<R: Read> StreamReader<R> {
- /// Try to create a new stream reader
- ///
- /// The first message in the stream is the schema, the reader will fail if it does not
- /// encounter a schema.
- /// To check if the reader is done, use `is_finished(self)`
- pub fn try_new(reader: R) -> Result<Self> {
- let mut reader = BufReader::new(reader);
- // determine metadata length
- let mut meta_size: [u8; 4] = [0; 4];
- reader.read_exact(&mut meta_size)?;
- let meta_len = {
- // If a continuation marker is encountered, skip over it and read
- // the size from the next four bytes.
- if meta_size == CONTINUATION_MARKER {
- reader.read_exact(&mut meta_size)?;
- }
- i32::from_le_bytes(meta_size)
- };
-
- let mut meta_buffer = vec![0; meta_len as usize];
- reader.read_exact(&mut meta_buffer)?;
-
- let message = ipc::root_as_message(meta_buffer.as_slice()).map_err(|err| {
- ArrowError::IoError(format!("Unable to get root as message: {:?}", err))
- })?;
- // message header is a Schema, so read it
- let ipc_schema: ipc::Schema = message.header_as_schema().ok_or_else(|| {
- ArrowError::IoError("Unable to read IPC message as schema".to_string())
- })?;
- let schema = ipc::convert::fb_to_schema(ipc_schema);
-
- // Create an array of optional dictionary value arrays, one per field.
- let dictionaries_by_field = vec![None; schema.fields().len()];
-
- Ok(Self {
- reader,
- schema: Arc::new(schema),
- finished: false,
- dictionaries_by_field,
- })
- }
-
- /// Return the schema of the stream
- pub fn schema(&self) -> SchemaRef {
- self.schema.clone()
- }
-
- /// Check if the stream is finished
- pub fn is_finished(&self) -> bool {
- self.finished
- }
-
- fn maybe_next(&mut self) -> Result<Option<RecordBatch>> {
- if self.finished {
- return Ok(None);
- }
- // determine metadata length
- let mut meta_size: [u8; 4] = [0; 4];
-
- match self.reader.read_exact(&mut meta_size) {
- Ok(()) => (),
- Err(e) => {
- return if e.kind() == std::io::ErrorKind::UnexpectedEof {
- // Handle EOF without the "0xFFFFFFFF 0x00000000"
- // valid according to:
- // https://arrow.apache.org/docs/format/Columnar.html#ipc-streaming-format
- self.finished = true;
- Ok(None)
- } else {
- Err(ArrowError::from(e))
- };
- }
- }
-
- let meta_len = {
- // If a continuation marker is encountered, skip over it and read
- // the size from the next four bytes.
- if meta_size == CONTINUATION_MARKER {
- self.reader.read_exact(&mut meta_size)?;
- }
- i32::from_le_bytes(meta_size)
- };
-
- if meta_len == 0 {
- // the stream has ended, mark the reader as finished
- self.finished = true;
- return Ok(None);
- }
-
- let mut meta_buffer = vec![0; meta_len as usize];
- self.reader.read_exact(&mut meta_buffer)?;
-
- let vecs = &meta_buffer.to_vec();
- let message = ipc::root_as_message(vecs).map_err(|err| {
- ArrowError::IoError(format!("Unable to get root as message: {:?}", err))
- })?;
-
- match message.header_type() {
- ipc::MessageHeader::Schema => Err(ArrowError::IoError(
- "Not expecting a schema when messages are read".to_string(),
- )),
- ipc::MessageHeader::RecordBatch => {
- let batch = message.header_as_record_batch().ok_or_else(|| {
- ArrowError::IoError(
- "Unable to read IPC message as record batch".to_string(),
- )
- })?;
- // read the block that makes up the record batch into a buffer
- let mut buf = vec![0; message.bodyLength() as usize];
- self.reader.read_exact(&mut buf)?;
-
- read_record_batch(&buf, batch, self.schema(), &self.dictionaries_by_field).map(Some)
- }
- ipc::MessageHeader::DictionaryBatch => {
- let batch = message.header_as_dictionary_batch().ok_or_else(|| {
- ArrowError::IoError(
- "Unable to read IPC message as dictionary batch".to_string(),
- )
- })?;
- // read the block that makes up the dictionary batch into a buffer
- let mut buf = vec![0; message.bodyLength() as usize];
- self.reader.read_exact(&mut buf)?;
-
- read_dictionary(
- &buf, batch, &self.schema, &mut self.dictionaries_by_field
- )?;
-
- // read the next message until we encounter a RecordBatch
- self.maybe_next()
- }
- ipc::MessageHeader::NONE => {
- Ok(None)
- }
- t => Err(ArrowError::IoError(
- format!("Reading types other than record batches not yet supported, unable to read {:?} ", t)
- )),
- }
- }
-}
-
-impl<R: Read> Iterator for StreamReader<R> {
- type Item = Result<RecordBatch>;
-
- fn next(&mut self) -> Option<Self::Item> {
- self.maybe_next().transpose()
- }
-}
-
-impl<R: Read> RecordBatchReader for StreamReader<R> {
- fn schema(&self) -> SchemaRef {
- self.schema.clone()
- }
-}
-
-#[cfg(test)]
-mod tests {
- use super::*;
-
- use std::fs::File;
-
- use flate2::read::GzDecoder;
-
- use crate::util::integration_util::*;
-
- #[test]
- fn read_generated_files_014() {
- let testdata = crate::util::test_util::arrow_test_data();
- let version = "0.14.1";
- // the test is repetitive, thus we can read all supported files at once
- let paths = vec![
- "generated_interval",
- "generated_datetime",
- "generated_dictionary",
- "generated_nested",
- "generated_primitive_no_batches",
- "generated_primitive_zerolength",
- "generated_primitive",
- "generated_decimal",
- ];
- paths.iter().for_each(|path| {
- let file = File::open(format!(
- "{}/arrow-ipc-stream/integration/{}/{}.arrow_file",
- testdata, version, path
- ))
- .unwrap();
-
- let mut reader = FileReader::try_new(file).unwrap();
-
- // read expected JSON output
- let arrow_json = read_gzip_json(version, path);
- assert!(arrow_json.equals_reader(&mut reader));
- });
- }
-
- #[test]
- #[should_panic(expected = "Big Endian is not supported for Decimal!")]
- fn read_decimal_be_file_should_panic() {
- let testdata = crate::util::test_util::arrow_test_data();
- let file = File::open(format!(
- "{}/arrow-ipc-stream/integration/1.0.0-bigendian/generated_decimal.arrow_file",
- testdata
- ))
- .unwrap();
- FileReader::try_new(file).unwrap();
- }
-
- #[test]
- fn read_generated_be_files_should_work() {
- // complementary to the previous test
- let testdata = crate::util::test_util::arrow_test_data();
- let paths = vec![
- "generated_interval",
- "generated_datetime",
- "generated_dictionary",
- "generated_nested",
- "generated_null_trivial",
- "generated_null",
- "generated_primitive_no_batches",
- "generated_primitive_zerolength",
- "generated_primitive",
- ];
- paths.iter().for_each(|path| {
- let file = File::open(format!(
- "{}/arrow-ipc-stream/integration/1.0.0-bigendian/{}.arrow_file",
- testdata, path
- ))
- .unwrap();
-
- FileReader::try_new(file).unwrap();
- });
- }
-
- #[test]
- fn read_generated_streams_014() {
- let testdata = crate::util::test_util::arrow_test_data();
- let version = "0.14.1";
- // the test is repetitive, thus we can read all supported files at once
- let paths = vec![
- "generated_interval",
- "generated_datetime",
- "generated_dictionary",
- "generated_nested",
- "generated_primitive_no_batches",
- "generated_primitive_zerolength",
- "generated_primitive",
- "generated_decimal",
- ];
- paths.iter().for_each(|path| {
- let file = File::open(format!(
- "{}/arrow-ipc-stream/integration/{}/{}.stream",
- testdata, version, path
- ))
- .unwrap();
-
- let mut reader = StreamReader::try_new(file).unwrap();
-
- // read expected JSON output
- let arrow_json = read_gzip_json(version, path);
- assert!(arrow_json.equals_reader(&mut reader));
- // the next batch must be empty
- assert!(reader.next().is_none());
- // the stream must indicate that it's finished
- assert!(reader.is_finished());
- });
- }
-
- #[test]
- fn read_generated_files_100() {
- let testdata = crate::util::test_util::arrow_test_data();
- let version = "1.0.0-littleendian";
- // the test is repetitive, thus we can read all supported files at once
- let paths = vec![
- "generated_interval",
- "generated_datetime",
- "generated_dictionary",
- "generated_nested",
- "generated_null_trivial",
- "generated_null",
- "generated_primitive_no_batches",
- "generated_primitive_zerolength",
- "generated_primitive",
- ];
- paths.iter().for_each(|path| {
- let file = File::open(format!(
- "{}/arrow-ipc-stream/integration/{}/{}.arrow_file",
- testdata, version, path
- ))
- .unwrap();
-
- let mut reader = FileReader::try_new(file).unwrap();
-
- // read expected JSON output
- let arrow_json = read_gzip_json(version, path);
- assert!(arrow_json.equals_reader(&mut reader));
- });
- }
-
- #[test]
- fn read_generated_streams_100() {
- let testdata = crate::util::test_util::arrow_test_data();
- let version = "1.0.0-littleendian";
- // the test is repetitive, thus we can read all supported files at once
- let paths = vec![
- "generated_interval",
- "generated_datetime",
- "generated_dictionary",
- "generated_nested",
- "generated_null_trivial",
- "generated_null",
- "generated_primitive_no_batches",
- "generated_primitive_zerolength",
- "generated_primitive",
- ];
- paths.iter().for_each(|path| {
- let file = File::open(format!(
- "{}/arrow-ipc-stream/integration/{}/{}.stream",
- testdata, version, path
- ))
- .unwrap();
-
- let mut reader = StreamReader::try_new(file).unwrap();
-
- // read expected JSON output
- let arrow_json = read_gzip_json(version, path);
- assert!(arrow_json.equals_reader(&mut reader));
- // the next batch must be empty
- assert!(reader.next().is_none());
- // the stream must indicate that it's finished
- assert!(reader.is_finished());
- });
- }
-
- #[test]
- fn test_arrow_single_float_row() {
- let schema = Schema::new(vec![
- Field::new("a", DataType::Float32, false),
- Field::new("b", DataType::Float32, false),
- Field::new("c", DataType::Int32, false),
- Field::new("d", DataType::Int32, false),
- ]);
- let arrays = vec![
- Arc::new(Float32Array::from(vec![1.23])) as ArrayRef,
- Arc::new(Float32Array::from(vec![-6.50])) as ArrayRef,
- Arc::new(Int32Array::from(vec![2])) as ArrayRef,
- Arc::new(Int32Array::from(vec![1])) as ArrayRef,
- ];
- let batch = RecordBatch::try_new(Arc::new(schema.clone()), arrays).unwrap();
- // create stream writer
- let file = File::create("target/debug/testdata/float.stream").unwrap();
- let mut stream_writer =
- crate::ipc::writer::StreamWriter::try_new(file, &schema).unwrap();
- stream_writer.write(&batch).unwrap();
- stream_writer.finish().unwrap();
-
- // read stream back
- let file = File::open("target/debug/testdata/float.stream").unwrap();
- let reader = StreamReader::try_new(file).unwrap();
-
- reader.for_each(|batch| {
- let batch = batch.unwrap();
- assert!(
- batch
- .column(0)
- .as_any()
- .downcast_ref::<Float32Array>()
- .unwrap()
- .value(0)
- != 0.0
- );
- assert!(
- batch
- .column(1)
- .as_any()
- .downcast_ref::<Float32Array>()
- .unwrap()
- .value(0)
- != 0.0
- );
- })
- }
-
- /// Read gzipped JSON file
- fn read_gzip_json(version: &str, path: &str) -> ArrowJson {
- let testdata = crate::util::test_util::arrow_test_data();
- let file = File::open(format!(
- "{}/arrow-ipc-stream/integration/{}/{}.json.gz",
- testdata, version, path
- ))
- .unwrap();
- let mut gz = GzDecoder::new(&file);
- let mut s = String::new();
- gz.read_to_string(&mut s).unwrap();
- // convert to Arrow JSON
- let arrow_json: ArrowJson = serde_json::from_str(&s).unwrap();
- arrow_json
- }
-}
diff --git a/arrow/src/ipc/writer.rs b/arrow/src/ipc/writer.rs
deleted file mode 100644
index a6df7b8..0000000
--- a/arrow/src/ipc/writer.rs
+++ /dev/null
@@ -1,1160 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-//! Arrow IPC File and Stream Writers
-//!
-//! The `FileWriter` and `StreamWriter` have similar interfaces,
-//! however the `FileWriter` expects a reader that supports `Seek`ing
-
-use std::collections::HashMap;
-use std::io::{BufWriter, Write};
-
-use flatbuffers::FlatBufferBuilder;
-
-use crate::array::{ArrayData, ArrayRef};
-use crate::buffer::{Buffer, MutableBuffer};
-use crate::datatypes::*;
-use crate::error::{ArrowError, Result};
-use crate::ipc;
-use crate::record_batch::RecordBatch;
-use crate::util::bit_util;
-
-use ipc::CONTINUATION_MARKER;
-
-/// IPC write options used to control the behaviour of the writer
-#[derive(Debug)]
-pub struct IpcWriteOptions {
- /// Write padding after memory buffers to this multiple of bytes.
- /// Generally 8 or 64, defaults to 8
- alignment: usize,
- /// The legacy format is for releases before 0.15.0, and uses metadata V4
- write_legacy_ipc_format: bool,
- /// The metadata version to write. The Rust IPC writer supports V4+
- ///
- /// *Default versions per crate*
- ///
- /// When creating the default IpcWriteOptions, the following metadata versions are used:
- ///
- /// version 2.0.0: V4, with legacy format enabled
- /// version 4.0.0: V5
- metadata_version: ipc::MetadataVersion,
-}
-
-impl IpcWriteOptions {
- /// Try create IpcWriteOptions, checking for incompatible settings
- pub fn try_new(
- alignment: usize,
- write_legacy_ipc_format: bool,
- metadata_version: ipc::MetadataVersion,
- ) -> Result<Self> {
- if alignment == 0 || alignment % 8 != 0 {
- return Err(ArrowError::InvalidArgumentError(
- "Alignment should be greater than 0 and be a multiple of 8".to_string(),
- ));
- }
- match metadata_version {
- ipc::MetadataVersion::V1
- | ipc::MetadataVersion::V2
- | ipc::MetadataVersion::V3 => Err(ArrowError::InvalidArgumentError(
- "Writing IPC metadata version 3 and lower not supported".to_string(),
- )),
- ipc::MetadataVersion::V4 => Ok(Self {
- alignment,
- write_legacy_ipc_format,
- metadata_version,
- }),
- ipc::MetadataVersion::V5 => {
- if write_legacy_ipc_format {
- Err(ArrowError::InvalidArgumentError(
- "Legacy IPC format only supported on metadata version 4"
- .to_string(),
- ))
- } else {
- Ok(Self {
- alignment,
- write_legacy_ipc_format,
- metadata_version,
- })
- }
- }
- z => panic!("Unsupported ipc::MetadataVersion {:?}", z),
- }
- }
-}
-
-impl Default for IpcWriteOptions {
- fn default() -> Self {
- Self {
- alignment: 8,
- write_legacy_ipc_format: false,
- metadata_version: ipc::MetadataVersion::V5,
- }
- }
-}
-
-#[derive(Debug, Default)]
-pub struct IpcDataGenerator {}
-
-impl IpcDataGenerator {
- pub fn schema_to_bytes(
- &self,
- schema: &Schema,
- write_options: &IpcWriteOptions,
- ) -> EncodedData {
- let mut fbb = FlatBufferBuilder::new();
- let schema = {
- let fb = ipc::convert::schema_to_fb_offset(&mut fbb, schema);
- fb.as_union_value()
- };
-
- let mut message = ipc::MessageBuilder::new(&mut fbb);
- message.add_version(write_options.metadata_version);
- message.add_header_type(ipc::MessageHeader::Schema);
- message.add_bodyLength(0);
- message.add_header(schema);
- // TODO: custom metadata
- let data = message.finish();
- fbb.finish(data, None);
-
- let data = fbb.finished_data();
- EncodedData {
- ipc_message: data.to_vec(),
- arrow_data: vec![],
- }
- }
-
- pub fn encoded_batch(
- &self,
- batch: &RecordBatch,
- dictionary_tracker: &mut DictionaryTracker,
- write_options: &IpcWriteOptions,
- ) -> Result<(Vec<EncodedData>, EncodedData)> {
- // TODO: handle nested dictionaries
- let schema = batch.schema();
- let mut encoded_dictionaries = Vec::with_capacity(schema.fields().len());
-
- for (i, field) in schema.fields().iter().enumerate() {
- let column = batch.column(i);
-
- if let DataType::Dictionary(_key_type, _value_type) = column.data_type() {
- let dict_id = field
- .dict_id()
- .expect("All Dictionary types have `dict_id`");
- let dict_data = column.data();
- let dict_values = &dict_data.child_data()[0];
-
- let emit = dictionary_tracker.insert(dict_id, column)?;
-
- if emit {
- encoded_dictionaries.push(self.dictionary_batch_to_bytes(
- dict_id,
- dict_values,
- write_options,
- ));
- }
- }
- }
-
- let encoded_message = self.record_batch_to_bytes(batch, write_options);
-
- Ok((encoded_dictionaries, encoded_message))
- }
-
- /// Write a `RecordBatch` into two sets of bytes, one for the header (ipc::Message) and the
- /// other for the batch's data
- fn record_batch_to_bytes(
- &self,
- batch: &RecordBatch,
- write_options: &IpcWriteOptions,
- ) -> EncodedData {
- let mut fbb = FlatBufferBuilder::new();
-
- let mut nodes: Vec<ipc::FieldNode> = vec![];
- let mut buffers: Vec<ipc::Buffer> = vec![];
- let mut arrow_data: Vec<u8> = vec![];
- let mut offset = 0;
- for array in batch.columns() {
- let array_data = array.data();
- offset = write_array_data(
- &array_data,
- &mut buffers,
- &mut arrow_data,
- &mut nodes,
- offset,
- array.len(),
- array.null_count(),
- );
- }
-
- // write data
- let buffers = fbb.create_vector(&buffers);
- let nodes = fbb.create_vector(&nodes);
-
- let root = {
- let mut batch_builder = ipc::RecordBatchBuilder::new(&mut fbb);
- batch_builder.add_length(batch.num_rows() as i64);
- batch_builder.add_nodes(nodes);
- batch_builder.add_buffers(buffers);
- let b = batch_builder.finish();
- b.as_union_value()
- };
- // create an ipc::Message
- let mut message = ipc::MessageBuilder::new(&mut fbb);
- message.add_version(write_options.metadata_version);
- message.add_header_type(ipc::MessageHeader::RecordBatch);
- message.add_bodyLength(arrow_data.len() as i64);
- message.add_header(root);
- let root = message.finish();
- fbb.finish(root, None);
- let finished_data = fbb.finished_data();
-
- EncodedData {
- ipc_message: finished_data.to_vec(),
- arrow_data,
- }
- }
-
- /// Write dictionary values into two sets of bytes, one for the header (ipc::Message) and the
- /// other for the data
- fn dictionary_batch_to_bytes(
- &self,
- dict_id: i64,
- array_data: &ArrayData,
- write_options: &IpcWriteOptions,
- ) -> EncodedData {
- let mut fbb = FlatBufferBuilder::new();
-
- let mut nodes: Vec<ipc::FieldNode> = vec![];
- let mut buffers: Vec<ipc::Buffer> = vec![];
- let mut arrow_data: Vec<u8> = vec![];
-
- write_array_data(
- &array_data,
- &mut buffers,
- &mut arrow_data,
- &mut nodes,
- 0,
- array_data.len(),
- array_data.null_count(),
- );
-
- // write data
- let buffers = fbb.create_vector(&buffers);
- let nodes = fbb.create_vector(&nodes);
-
- let root = {
- let mut batch_builder = ipc::RecordBatchBuilder::new(&mut fbb);
- batch_builder.add_length(array_data.len() as i64);
- batch_builder.add_nodes(nodes);
- batch_builder.add_buffers(buffers);
- batch_builder.finish()
- };
-
- let root = {
- let mut batch_builder = ipc::DictionaryBatchBuilder::new(&mut fbb);
- batch_builder.add_id(dict_id);
- batch_builder.add_data(root);
- batch_builder.finish().as_union_value()
- };
-
- let root = {
- let mut message_builder = ipc::MessageBuilder::new(&mut fbb);
- message_builder.add_version(write_options.metadata_version);
- message_builder.add_header_type(ipc::MessageHeader::DictionaryBatch);
- message_builder.add_bodyLength(arrow_data.len() as i64);
- message_builder.add_header(root);
- message_builder.finish()
- };
-
- fbb.finish(root, None);
- let finished_data = fbb.finished_data();
-
- EncodedData {
- ipc_message: finished_data.to_vec(),
- arrow_data,
- }
- }
-}
-
-/// Keeps track of dictionaries that have been written, to avoid emitting the same dictionary
-/// multiple times. Can optionally error if an update to an existing dictionary is attempted, which
-/// isn't allowed in the `FileWriter`.
-pub struct DictionaryTracker {
- written: HashMap<i64, ArrayRef>,
- error_on_replacement: bool,
-}
-
-impl DictionaryTracker {
- pub fn new(error_on_replacement: bool) -> Self {
- Self {
- written: HashMap::new(),
- error_on_replacement,
- }
- }
-
- /// Keep track of the dictionary with the given ID and values. Behavior:
- ///
- /// * If this ID has been written already and has the same data, return `Ok(false)` to indicate
- /// that the dictionary was not actually inserted (because it's already been seen).
- /// * If this ID has been written already but with different data, and this tracker is
- /// configured to return an error, return an error.
- /// * If the tracker has not been configured to error on replacement or this dictionary
- /// has never been seen before, return `Ok(true)` to indicate that the dictionary was just
- /// inserted.
- pub fn insert(&mut self, dict_id: i64, column: &ArrayRef) -> Result<bool> {
- let dict_data = column.data();
- let dict_values = &dict_data.child_data()[0];
-
- // If a dictionary with this id was already emitted, check if it was the same.
- if let Some(last) = self.written.get(&dict_id) {
- if last.data().child_data()[0] == *dict_values {
- // Same dictionary values => no need to emit it again
- return Ok(false);
- } else if self.error_on_replacement {
- return Err(ArrowError::InvalidArgumentError(
- "Dictionary replacement detected when writing IPC file format. \
- Arrow IPC files only support a single dictionary for a given field \
- across all batches."
- .to_string(),
- ));
- }
- }
-
- self.written.insert(dict_id, column.clone());
- Ok(true)
- }
-}
-
-pub struct FileWriter<W: Write> {
- /// The object to write to
- writer: BufWriter<W>,
- /// IPC write options
- write_options: IpcWriteOptions,
- /// A reference to the schema, used in validating record batches
- schema: Schema,
- /// The number of bytes between each block of bytes, as an offset for random access
- block_offsets: usize,
- /// Dictionary blocks that will be written as part of the IPC footer
- dictionary_blocks: Vec<ipc::Block>,
- /// Record blocks that will be written as part of the IPC footer
- record_blocks: Vec<ipc::Block>,
- /// Whether the writer footer has been written, and the writer is finished
- finished: bool,
- /// Keeps track of dictionaries that have been written
- dictionary_tracker: DictionaryTracker,
-
- data_gen: IpcDataGenerator,
-}
-
-impl<W: Write> FileWriter<W> {
- /// Try create a new writer, with the schema written as part of the header
- pub fn try_new(writer: W, schema: &Schema) -> Result<Self> {
- let write_options = IpcWriteOptions::default();
- Self::try_new_with_options(writer, schema, write_options)
- }
-
- /// Try create a new writer with IpcWriteOptions
- pub fn try_new_with_options(
- writer: W,
- schema: &Schema,
- write_options: IpcWriteOptions,
- ) -> Result<Self> {
- let data_gen = IpcDataGenerator::default();
- let mut writer = BufWriter::new(writer);
- // write magic to header
- writer.write_all(&super::ARROW_MAGIC[..])?;
- // create an 8-byte boundary after the header
- writer.write_all(&[0, 0])?;
- // write the schema, set the written bytes to the schema + header
- let encoded_message = data_gen.schema_to_bytes(schema, &write_options);
- let (meta, data) = write_message(&mut writer, encoded_message, &write_options)?;
- Ok(Self {
- writer,
- write_options,
- schema: schema.clone(),
- block_offsets: meta + data + 8,
- dictionary_blocks: vec![],
- record_blocks: vec![],
- finished: false,
- dictionary_tracker: DictionaryTracker::new(true),
- data_gen,
- })
- }
-
- /// Write a record batch to the file
- pub fn write(&mut self, batch: &RecordBatch) -> Result<()> {
- if self.finished {
- return Err(ArrowError::IoError(
- "Cannot write record batch to file writer as it is closed".to_string(),
- ));
- }
-
- let (encoded_dictionaries, encoded_message) = self.data_gen.encoded_batch(
- batch,
- &mut self.dictionary_tracker,
- &self.write_options,
- )?;
-
- for encoded_dictionary in encoded_dictionaries {
- let (meta, data) =
- write_message(&mut self.writer, encoded_dictionary, &self.write_options)?;
-
- let block =
- ipc::Block::new(self.block_offsets as i64, meta as i32, data as i64);
- self.dictionary_blocks.push(block);
- self.block_offsets += meta + data;
- }
-
- let (meta, data) =
- write_message(&mut self.writer, encoded_message, &self.write_options)?;
- // add a record block for the footer
- let block = ipc::Block::new(
- self.block_offsets as i64,
- meta as i32, // TODO: is this still applicable?
- data as i64,
- );
- self.record_blocks.push(block);
- self.block_offsets += meta + data;
- Ok(())
- }
-
- /// Write footer and closing tag, then mark the writer as done
- pub fn finish(&mut self) -> Result<()> {
- if self.finished {
- return Err(ArrowError::IoError(
- "Cannot write footer to file writer as it is closed".to_string(),
- ));
- }
-
- // write EOS
- write_continuation(&mut self.writer, &self.write_options, 0)?;
-
- let mut fbb = FlatBufferBuilder::new();
- let dictionaries = fbb.create_vector(&self.dictionary_blocks);
- let record_batches = fbb.create_vector(&self.record_blocks);
- let schema = ipc::convert::schema_to_fb_offset(&mut fbb, &self.schema);
-
- let root = {
- let mut footer_builder = ipc::FooterBuilder::new(&mut fbb);
- footer_builder.add_version(self.write_options.metadata_version);
- footer_builder.add_schema(schema);
- footer_builder.add_dictionaries(dictionaries);
- footer_builder.add_recordBatches(record_batches);
- footer_builder.finish()
- };
- fbb.finish(root, None);
- let footer_data = fbb.finished_data();
- self.writer.write_all(footer_data)?;
- self.writer
- .write_all(&(footer_data.len() as i32).to_le_bytes())?;
- self.writer.write_all(&super::ARROW_MAGIC)?;
- self.writer.flush()?;
- self.finished = true;
-
- Ok(())
- }
-}
-
-pub struct StreamWriter<W: Write> {
- /// The object to write to
- writer: BufWriter<W>,
- /// IPC write options
- write_options: IpcWriteOptions,
- /// A reference to the schema, used in validating record batches
- schema: Schema,
- /// Whether the writer footer has been written, and the writer is finished
- finished: bool,
- /// Keeps track of dictionaries that have been written
- dictionary_tracker: DictionaryTracker,
-
- data_gen: IpcDataGenerator,
-}
-
-impl<W: Write> StreamWriter<W> {
- /// Try create a new writer, with the schema written as part of the header
- pub fn try_new(writer: W, schema: &Schema) -> Result<Self> {
- let write_options = IpcWriteOptions::default();
- Self::try_new_with_options(writer, schema, write_options)
- }
-
- pub fn try_new_with_options(
- writer: W,
- schema: &Schema,
- write_options: IpcWriteOptions,
- ) -> Result<Self> {
- let data_gen = IpcDataGenerator::default();
- let mut writer = BufWriter::new(writer);
- // write the schema, set the written bytes to the schema
- let encoded_message = data_gen.schema_to_bytes(schema, &write_options);
- write_message(&mut writer, encoded_message, &write_options)?;
- Ok(Self {
- writer,
- write_options,
- schema: schema.clone(),
- finished: false,
- dictionary_tracker: DictionaryTracker::new(false),
- data_gen,
- })
- }
-
- /// Write a record batch to the stream
- pub fn write(&mut self, batch: &RecordBatch) -> Result<()> {
- if self.finished {
- return Err(ArrowError::IoError(
- "Cannot write record batch to stream writer as it is closed".to_string(),
- ));
- }
-
- let (encoded_dictionaries, encoded_message) = self
- .data_gen
- .encoded_batch(batch, &mut self.dictionary_tracker, &self.write_options)
- .expect("StreamWriter is configured to not error on dictionary replacement");
-
- for encoded_dictionary in encoded_dictionaries {
- write_message(&mut self.writer, encoded_dictionary, &self.write_options)?;
- }
-
- write_message(&mut self.writer, encoded_message, &self.write_options)?;
- Ok(())
- }
-
- /// Write continuation bytes, and mark the stream as done
- pub fn finish(&mut self) -> Result<()> {
- if self.finished {
- return Err(ArrowError::IoError(
- "Cannot write footer to stream writer as it is closed".to_string(),
- ));
- }
-
- write_continuation(&mut self.writer, &self.write_options, 0)?;
-
- self.finished = true;
-
- Ok(())
- }
-
- /// Unwraps the BufWriter housed in StreamWriter.writer, returning the underlying
- /// writer
- ///
- /// The buffer is flushed and the StreamWriter is finished before returning the
- /// writer.
- ///
- /// # Errors
- ///
- /// An ['Err'] may be returned if an error occurs while finishing the StreamWriter
- /// or while flushing the buffer.
- ///
- /// # Example
- ///
- /// ```
- /// # use arrow::datatypes::Schema;
- /// # use arrow::ipc::writer::StreamWriter;
- /// # use arrow::error::ArrowError;
- /// # fn main() -> Result<(), ArrowError> {
- /// // The result we expect from an empty schema
- /// let expected = vec![
- /// 255, 255, 255, 255, 64, 0, 0, 0,
- /// 16, 0, 0, 0, 0, 0, 10, 0,
- /// 14, 0, 12, 0, 11, 0, 4, 0,
- /// 10, 0, 0, 0, 20, 0, 0, 0,
- /// 0, 0, 0, 1, 4, 0, 10, 0,
- /// 12, 0, 0, 0, 8, 0, 4, 0,
- /// 10, 0, 0, 0, 8, 0, 0, 0,
- /// 8, 0, 0, 0, 0, 0, 0, 0,
- /// 0, 0, 0, 0, 0, 0, 0, 0,
- /// 255, 255, 255, 255, 0, 0, 0, 0
- /// ];
- ///
- /// let schema = Schema::new(vec![]);
- /// let buffer: Vec<u8> = Vec::new();
- /// let stream_writer = StreamWriter::try_new(buffer, &schema)?;
- ///
- /// assert_eq!(stream_writer.into_inner()?, expected);
- /// # Ok(())
- /// # }
- /// ```
- pub fn into_inner(mut self) -> Result<W> {
- if !self.finished {
- self.finish()?;
- }
- self.writer.into_inner().map_err(ArrowError::from)
- }
-}
-
-/// Stores the encoded data, which is an ipc::Message, and optional Arrow data
-pub struct EncodedData {
- /// An encoded ipc::Message
- pub ipc_message: Vec<u8>,
- /// Arrow buffers to be written, should be an empty vec for schema messages
- pub arrow_data: Vec<u8>,
-}
-/// Write a message's IPC data and buffers, returning metadata and buffer data lengths written
-pub fn write_message<W: Write>(
- mut writer: W,
- encoded: EncodedData,
- write_options: &IpcWriteOptions,
-) -> Result<(usize, usize)> {
- let arrow_data_len = encoded.arrow_data.len();
- if arrow_data_len % 8 != 0 {
- return Err(ArrowError::MemoryError(
- "Arrow data not aligned".to_string(),
- ));
- }
-
- let a = write_options.alignment - 1;
- let buffer = encoded.ipc_message;
- let flatbuf_size = buffer.len();
- let prefix_size = if write_options.write_legacy_ipc_format {
- 4
- } else {
- 8
- };
- let aligned_size = (flatbuf_size + prefix_size + a) & !a;
- let padding_bytes = aligned_size - flatbuf_size - prefix_size;
-
- write_continuation(
- &mut writer,
- &write_options,
- (aligned_size - prefix_size) as i32,
- )?;
-
- // write the flatbuf
- if flatbuf_size > 0 {
- writer.write_all(&buffer)?;
- }
- // write padding
- writer.write_all(&vec![0; padding_bytes])?;
-
- // write arrow data
- let body_len = if arrow_data_len > 0 {
- write_body_buffers(&mut writer, &encoded.arrow_data)?
- } else {
- 0
- };
-
- Ok((aligned_size, body_len))
-}
-
-fn write_body_buffers<W: Write>(mut writer: W, data: &[u8]) -> Result<usize> {
- let len = data.len() as u32;
- let pad_len = pad_to_8(len) as u32;
- let total_len = len + pad_len;
-
- // write body buffer
- writer.write_all(data)?;
- if pad_len > 0 {
- writer.write_all(&vec![0u8; pad_len as usize][..])?;
- }
-
- writer.flush()?;
- Ok(total_len as usize)
-}
-
-/// Write a record batch to the writer, writing the message size before the message
-/// if the record batch is being written to a stream
-fn write_continuation<W: Write>(
- mut writer: W,
- write_options: &IpcWriteOptions,
- total_len: i32,
-) -> Result<usize> {
- let mut written = 8;
-
- // the version of the writer determines whether continuation markers should be added
- match write_options.metadata_version {
- ipc::MetadataVersion::V1
- | ipc::MetadataVersion::V2
- | ipc::MetadataVersion::V3 => {
- unreachable!("Options with the metadata version cannot be created")
- }
- ipc::MetadataVersion::V4 => {
- if !write_options.write_legacy_ipc_format {
- // v0.15.0 format
- writer.write_all(&CONTINUATION_MARKER)?;
- written = 4;
- }
- writer.write_all(&total_len.to_le_bytes()[..])?;
- }
- ipc::MetadataVersion::V5 => {
- // write continuation marker and message length
- writer.write_all(&CONTINUATION_MARKER)?;
- writer.write_all(&total_len.to_le_bytes()[..])?;
- }
- z => panic!("Unsupported ipc::MetadataVersion {:?}", z),
- };
-
- writer.flush()?;
-
- Ok(written)
-}
-
-/// Write array data to a vector of bytes
-fn write_array_data(
- array_data: &ArrayData,
- mut buffers: &mut Vec<ipc::Buffer>,
- mut arrow_data: &mut Vec<u8>,
- mut nodes: &mut Vec<ipc::FieldNode>,
- offset: i64,
- num_rows: usize,
- null_count: usize,
-) -> i64 {
- let mut offset = offset;
- nodes.push(ipc::FieldNode::new(num_rows as i64, null_count as i64));
- // NullArray does not have any buffers, thus the null buffer is not generated
- if array_data.data_type() != &DataType::Null {
- // write null buffer if exists
- let null_buffer = match array_data.null_buffer() {
- None => {
- // create a buffer and fill it with valid bits
- let num_bytes = bit_util::ceil(num_rows, 8);
- let buffer = MutableBuffer::new(num_bytes);
- let buffer = buffer.with_bitset(num_bytes, true);
- buffer.into()
- }
- Some(buffer) => buffer.clone(),
- };
-
- offset = write_buffer(&null_buffer, &mut buffers, &mut arrow_data, offset);
- }
-
- array_data.buffers().iter().for_each(|buffer| {
- offset = write_buffer(buffer, &mut buffers, &mut arrow_data, offset);
- });
-
- if !matches!(array_data.data_type(), DataType::Dictionary(_, _)) {
- // recursively write out nested structures
- array_data.child_data().iter().for_each(|data_ref| {
- // write the nested data (e.g list data)
- offset = write_array_data(
- data_ref,
- &mut buffers,
- &mut arrow_data,
- &mut nodes,
- offset,
- data_ref.len(),
- data_ref.null_count(),
- );
- });
- }
-
- offset
-}
-
-/// Write a buffer to a vector of bytes, and add its ipc::Buffer to a vector
-fn write_buffer(
- buffer: &Buffer,
- buffers: &mut Vec<ipc::Buffer>,
- arrow_data: &mut Vec<u8>,
- offset: i64,
-) -> i64 {
- let len = buffer.len();
- let pad_len = pad_to_8(len as u32);
- let total_len: i64 = (len + pad_len) as i64;
- // assert_eq!(len % 8, 0, "Buffer width not a multiple of 8 bytes");
- buffers.push(ipc::Buffer::new(offset, total_len));
- arrow_data.extend_from_slice(buffer.as_slice());
- arrow_data.extend_from_slice(&vec![0u8; pad_len][..]);
- offset + total_len
-}
-
-/// Calculate an 8-byte boundary and return the number of bytes needed to pad to 8 bytes
-#[inline]
-fn pad_to_8(len: u32) -> usize {
- (((len + 7) & !7) - len) as usize
-}
-
-#[cfg(test)]
-mod tests {
- use super::*;
-
- use std::fs::File;
- use std::io::Read;
- use std::sync::Arc;
-
- use flate2::read::GzDecoder;
- use ipc::MetadataVersion;
-
- use crate::array::*;
- use crate::datatypes::Field;
- use crate::ipc::reader::*;
- use crate::util::integration_util::*;
-
- #[test]
- fn test_write_file() {
- let schema = Schema::new(vec![Field::new("field1", DataType::UInt32, false)]);
- let values: Vec<Option<u32>> = vec![
- Some(999),
- None,
- Some(235),
- Some(123),
- None,
- None,
- None,
- None,
- None,
- ];
- let array1 = UInt32Array::from(values);
- let batch = RecordBatch::try_new(
- Arc::new(schema.clone()),
- vec![Arc::new(array1) as ArrayRef],
- )
- .unwrap();
- {
- let file = File::create("target/debug/testdata/arrow.arrow_file").unwrap();
- let mut writer = FileWriter::try_new(file, &schema).unwrap();
-
- writer.write(&batch).unwrap();
- writer.finish().unwrap();
- }
-
- {
- let file =
- File::open(format!("target/debug/testdata/{}.arrow_file", "arrow"))
- .unwrap();
- let mut reader = FileReader::try_new(file).unwrap();
- while let Some(Ok(read_batch)) = reader.next() {
- read_batch
- .columns()
- .iter()
- .zip(batch.columns())
- .for_each(|(a, b)| {
- assert_eq!(a.data_type(), b.data_type());
- assert_eq!(a.len(), b.len());
- assert_eq!(a.null_count(), b.null_count());
- });
- }
- }
- }
-
- fn write_null_file(options: IpcWriteOptions, suffix: &str) {
- let schema = Schema::new(vec![
- Field::new("nulls", DataType::Null, true),
- Field::new("int32s", DataType::Int32, false),
- Field::new("nulls2", DataType::Null, false),
- Field::new("f64s", DataType::Float64, false),
- ]);
- let array1 = NullArray::new(32);
- let array2 = Int32Array::from(vec![1; 32]);
- let array3 = NullArray::new(32);
- let array4 = Float64Array::from(vec![std::f64::NAN; 32]);
- let batch = RecordBatch::try_new(
- Arc::new(schema.clone()),
- vec![
- Arc::new(array1) as ArrayRef,
- Arc::new(array2) as ArrayRef,
- Arc::new(array3) as ArrayRef,
- Arc::new(array4) as ArrayRef,
- ],
- )
- .unwrap();
- let file_name = format!("target/debug/testdata/nulls_{}.arrow_file", suffix);
- {
- let file = File::create(&file_name).unwrap();
- let mut writer =
- FileWriter::try_new_with_options(file, &schema, options).unwrap();
-
- writer.write(&batch).unwrap();
- writer.finish().unwrap();
- }
-
- {
- let file = File::open(&file_name).unwrap();
- let reader = FileReader::try_new(file).unwrap();
- reader.for_each(|maybe_batch| {
- maybe_batch
- .unwrap()
- .columns()
- .iter()
- .zip(batch.columns())
- .for_each(|(a, b)| {
- assert_eq!(a.data_type(), b.data_type());
- assert_eq!(a.len(), b.len());
- assert_eq!(a.null_count(), b.null_count());
- });
- });
- }
- }
- #[test]
- fn test_write_null_file_v4() {
- write_null_file(
- IpcWriteOptions::try_new(8, false, MetadataVersion::V4).unwrap(),
- "v4_a8",
- );
- write_null_file(
- IpcWriteOptions::try_new(8, true, MetadataVersion::V4).unwrap(),
- "v4_a8l",
- );
- write_null_file(
- IpcWriteOptions::try_new(64, false, MetadataVersion::V4).unwrap(),
- "v4_a64",
- );
- write_null_file(
- IpcWriteOptions::try_new(64, true, MetadataVersion::V4).unwrap(),
- "v4_a64l",
- );
- }
-
- #[test]
- fn test_write_null_file_v5() {
- write_null_file(
- IpcWriteOptions::try_new(8, false, MetadataVersion::V5).unwrap(),
- "v5_a8",
- );
- write_null_file(
- IpcWriteOptions::try_new(64, false, MetadataVersion::V5).unwrap(),
- "v5_a64",
- );
- }
-
- #[test]
- fn read_and_rewrite_generated_files_014() {
- let testdata = crate::util::test_util::arrow_test_data();
- let version = "0.14.1";
- // the test is repetitive, thus we can read all supported files at once
- let paths = vec![
- "generated_interval",
- "generated_datetime",
- "generated_dictionary",
- "generated_nested",
- "generated_primitive_no_batches",
- "generated_primitive_zerolength",
- "generated_primitive",
- "generated_decimal",
- ];
- paths.iter().for_each(|path| {
- let file = File::open(format!(
- "{}/arrow-ipc-stream/integration/{}/{}.arrow_file",
- testdata, version, path
- ))
- .unwrap();
-
- let mut reader = FileReader::try_new(file).unwrap();
-
- // read and rewrite the file to a temp location
- {
- let file = File::create(format!(
- "target/debug/testdata/{}-{}.arrow_file",
- version, path
- ))
- .unwrap();
- let mut writer = FileWriter::try_new(file, &reader.schema()).unwrap();
- while let Some(Ok(batch)) = reader.next() {
- writer.write(&batch).unwrap();
- }
- writer.finish().unwrap();
- }
-
- let file = File::open(format!(
- "target/debug/testdata/{}-{}.arrow_file",
- version, path
- ))
- .unwrap();
- let mut reader = FileReader::try_new(file).unwrap();
-
- // read expected JSON output
- let arrow_json = read_gzip_json(version, path);
- assert!(arrow_json.equals_reader(&mut reader));
- });
- }
-
- #[test]
- fn read_and_rewrite_generated_streams_014() {
- let testdata = crate::util::test_util::arrow_test_data();
- let version = "0.14.1";
- // the test is repetitive, thus we can read all supported files at once
- let paths = vec![
- "generated_interval",
- "generated_datetime",
- "generated_dictionary",
- "generated_nested",
- "generated_primitive_no_batches",
- "generated_primitive_zerolength",
- "generated_primitive",
- "generated_decimal",
- ];
- paths.iter().for_each(|path| {
- let file = File::open(format!(
- "{}/arrow-ipc-stream/integration/{}/{}.stream",
- testdata, version, path
- ))
- .unwrap();
-
- let reader = StreamReader::try_new(file).unwrap();
-
- // read and rewrite the stream to a temp location
- {
- let file = File::create(format!(
- "target/debug/testdata/{}-{}.stream",
- version, path
- ))
- .unwrap();
- let mut writer = StreamWriter::try_new(file, &reader.schema()).unwrap();
- reader.for_each(|batch| {
- writer.write(&batch.unwrap()).unwrap();
- });
- writer.finish().unwrap();
- }
-
- let file =
- File::open(format!("target/debug/testdata/{}-{}.stream", version, path))
- .unwrap();
- let mut reader = StreamReader::try_new(file).unwrap();
-
- // read expected JSON output
- let arrow_json = read_gzip_json(version, path);
- assert!(arrow_json.equals_reader(&mut reader));
- });
- }
-
- #[test]
- fn read_and_rewrite_generated_files_100() {
- let testdata = crate::util::test_util::arrow_test_data();
- let version = "1.0.0-littleendian";
- // the test is repetitive, thus we can read all supported files at once
- let paths = vec![
- "generated_custom_metadata",
- "generated_datetime",
- "generated_dictionary_unsigned",
- "generated_dictionary",
- // "generated_duplicate_fieldnames",
- "generated_interval",
- "generated_nested",
- // "generated_nested_large_offsets",
- "generated_null_trivial",
- "generated_null",
- "generated_primitive_large_offsets",
- "generated_primitive_no_batches",
- "generated_primitive_zerolength",
- "generated_primitive",
- // "generated_recursive_nested",
- ];
- paths.iter().for_each(|path| {
- let file = File::open(format!(
- "{}/arrow-ipc-stream/integration/{}/{}.arrow_file",
- testdata, version, path
- ))
- .unwrap();
-
- let mut reader = FileReader::try_new(file).unwrap();
-
- // read and rewrite the file to a temp location
- {
- let file = File::create(format!(
- "target/debug/testdata/{}-{}.arrow_file",
- version, path
- ))
- .unwrap();
- // write IPC version 5
- let options =
- IpcWriteOptions::try_new(8, false, ipc::MetadataVersion::V5).unwrap();
- let mut writer =
- FileWriter::try_new_with_options(file, &reader.schema(), options)
- .unwrap();
- while let Some(Ok(batch)) = reader.next() {
- writer.write(&batch).unwrap();
- }
- writer.finish().unwrap();
- }
-
- let file = File::open(format!(
- "target/debug/testdata/{}-{}.arrow_file",
- version, path
- ))
- .unwrap();
- let mut reader = FileReader::try_new(file).unwrap();
-
- // read expected JSON output
- let arrow_json = read_gzip_json(version, path);
- assert!(arrow_json.equals_reader(&mut reader));
- });
- }
-
- #[test]
- fn read_and_rewrite_generated_streams_100() {
- let testdata = crate::util::test_util::arrow_test_data();
- let version = "1.0.0-littleendian";
- // the test is repetitive, thus we can read all supported files at once
- let paths = vec![
- "generated_custom_metadata",
- "generated_datetime",
- "generated_dictionary_unsigned",
- "generated_dictionary",
- // "generated_duplicate_fieldnames",
- "generated_interval",
- "generated_nested",
- // "generated_nested_large_offsets",
- "generated_null_trivial",
- "generated_null",
- "generated_primitive_large_offsets",
- "generated_primitive_no_batches",
- "generated_primitive_zerolength",
- "generated_primitive",
- // "generated_recursive_nested",
- ];
- paths.iter().for_each(|path| {
- let file = File::open(format!(
- "{}/arrow-ipc-stream/integration/{}/{}.stream",
- testdata, version, path
- ))
- .unwrap();
-
- let reader = StreamReader::try_new(file).unwrap();
-
- // read and rewrite the stream to a temp location
- {
- let file = File::create(format!(
- "target/debug/testdata/{}-{}.stream",
- version, path
- ))
- .unwrap();
- let options =
- IpcWriteOptions::try_new(8, false, ipc::MetadataVersion::V5).unwrap();
- let mut writer =
- StreamWriter::try_new_with_options(file, &reader.schema(), options)
- .unwrap();
- reader.for_each(|batch| {
- writer.write(&batch.unwrap()).unwrap();
- });
- writer.finish().unwrap();
- }
-
- let file =
- File::open(format!("target/debug/testdata/{}-{}.stream", version, path))
- .unwrap();
- let mut reader = StreamReader::try_new(file).unwrap();
-
- // read expected JSON output
- let arrow_json = read_gzip_json(version, path);
- assert!(arrow_json.equals_reader(&mut reader));
- });
- }
-
- /// Read gzipped JSON file
- fn read_gzip_json(version: &str, path: &str) -> ArrowJson {
- let testdata = crate::util::test_util::arrow_test_data();
- let file = File::open(format!(
- "{}/arrow-ipc-stream/integration/{}/{}.json.gz",
- testdata, version, path
- ))
- .unwrap();
- let mut gz = GzDecoder::new(&file);
- let mut s = String::new();
- gz.read_to_string(&mut s).unwrap();
- // convert to Arrow JSON
- let arrow_json: ArrowJson = serde_json::from_str(&s).unwrap();
- arrow_json
- }
-}
diff --git a/arrow/src/json/mod.rs b/arrow/src/json/mod.rs
deleted file mode 100644
index 6b3df18..0000000
--- a/arrow/src/json/mod.rs
+++ /dev/null
@@ -1,27 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-//! Transfer data between the Arrow memory format and JSON
-//! line-delimited records. See the module level documentation for the
-//! [`reader`] and [`writer`] for usage examples.
-
-pub mod reader;
-pub mod writer;
-
-pub use self::reader::Reader;
-pub use self::reader::ReaderBuilder;
-pub use self::writer::{ArrayWriter, LineDelimitedWriter, Writer};
diff --git a/arrow/src/json/reader.rs b/arrow/src/json/reader.rs
deleted file mode 100644
index 9235142..0000000
--- a/arrow/src/json/reader.rs
+++ /dev/null
@@ -1,2988 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-//! # JSON Reader
-//!
-//! This JSON reader allows JSON line-delimited files to be read into the Arrow memory
-//! model. Records are loaded in batches and are then converted from row-based data to
-//! columnar data.
-//!
-//! Example:
-//!
-//! ```
-//! use arrow::datatypes::{DataType, Field, Schema};
-//! use arrow::json;
-//! use std::fs::File;
-//! use std::io::BufReader;
-//! use std::sync::Arc;
-//!
-//! let schema = Schema::new(vec![
-//! Field::new("a", DataType::Float64, false),
-//! Field::new("b", DataType::Float64, false),
-//! Field::new("c", DataType::Float64, false),
-//! ]);
-//!
-//! let file = File::open("test/data/basic.json").unwrap();
-//!
-//! let mut json = json::Reader::new(BufReader::new(file), Arc::new(schema), 1024, None);
-//! let batch = json.next().unwrap().unwrap();
-//! ```
-
-use std::io::{BufRead, BufReader, Read, Seek, SeekFrom};
-use std::iter::FromIterator;
-use std::sync::Arc;
-
-use indexmap::map::IndexMap as HashMap;
-use indexmap::set::IndexSet as HashSet;
-use serde_json::{map::Map as JsonMap, Value};
-
-use crate::buffer::MutableBuffer;
-use crate::datatypes::*;
-use crate::error::{ArrowError, Result};
-use crate::record_batch::RecordBatch;
-use crate::util::bit_util;
-use crate::{array::*, buffer::Buffer};
-
-#[derive(Debug, Clone)]
-enum InferredType {
- Scalar(HashSet<DataType>),
- Array(Box<InferredType>),
- Object(HashMap<String, InferredType>),
- Any,
-}
-
-impl InferredType {
- fn merge(&mut self, other: InferredType) -> Result<()> {
- match (self, other) {
- (InferredType::Array(s), InferredType::Array(o)) => {
- s.merge(*o)?;
- }
- (InferredType::Scalar(self_hs), InferredType::Scalar(other_hs)) => {
- other_hs.into_iter().for_each(|v| {
- self_hs.insert(v);
- });
- }
- (InferredType::Object(self_map), InferredType::Object(other_map)) => {
- for (k, v) in other_map {
- self_map.entry(k).or_insert(InferredType::Any).merge(v)?;
- }
- }
- (s @ InferredType::Any, v) => {
- *s = v;
- }
- (_, InferredType::Any) => {}
- // convert a scalar type to a single-item scalar array type.
- (
- InferredType::Array(self_inner_type),
- other_scalar @ InferredType::Scalar(_),
- ) => {
- self_inner_type.merge(other_scalar)?;
- }
- (s @ InferredType::Scalar(_), InferredType::Array(mut other_inner_type)) => {
- other_inner_type.merge(s.clone())?;
- *s = InferredType::Array(other_inner_type);
- }
- // incompatible types
- (s, o) => {
- return Err(ArrowError::JsonError(format!(
- "Incompatible type found during schema inference: {:?} v.s. {:?}",
- s, o,
- )));
- }
- }
-
- Ok(())
- }
-}
-
-/// Coerce data type during inference
-///
-/// * `Int64` and `Float64` should be `Float64`
-/// * Lists and scalars are coerced to a list of a compatible scalar
-/// * All other types are coerced to `Utf8`
-fn coerce_data_type(dt: Vec<&DataType>) -> DataType {
- let mut dt_iter = dt.into_iter().cloned();
- let dt_init = dt_iter.next().unwrap_or(DataType::Utf8);
-
- dt_iter.fold(dt_init, |l, r| match (l, r) {
- (DataType::Boolean, DataType::Boolean) => DataType::Boolean,
- (DataType::Int64, DataType::Int64) => DataType::Int64,
- (DataType::Float64, DataType::Float64)
- | (DataType::Float64, DataType::Int64)
- | (DataType::Int64, DataType::Float64) => DataType::Float64,
- (DataType::List(l), DataType::List(r)) => DataType::List(Box::new(Field::new(
- "item",
- coerce_data_type(vec![l.data_type(), r.data_type()]),
- true,
- ))),
- // coerce scalar and scalar array into scalar array
- (DataType::List(e), not_list) | (not_list, DataType::List(e)) => {
- DataType::List(Box::new(Field::new(
- "item",
- coerce_data_type(vec![e.data_type(), ¬_list]),
- true,
- )))
- }
- _ => DataType::Utf8,
- })
-}
-
-fn generate_datatype(t: &InferredType) -> Result<DataType> {
- Ok(match t {
- InferredType::Scalar(hs) => coerce_data_type(hs.iter().collect()),
- InferredType::Object(spec) => DataType::Struct(generate_fields(spec)?),
- InferredType::Array(ele_type) => DataType::List(Box::new(Field::new(
- "item",
- generate_datatype(ele_type)?,
- true,
- ))),
- InferredType::Any => DataType::Null,
- })
-}
-
-fn generate_fields(spec: &HashMap<String, InferredType>) -> Result<Vec<Field>> {
- spec.iter()
- .map(|(k, types)| Ok(Field::new(k, generate_datatype(types)?, true)))
- .collect()
-}
-
-/// Generate schema from JSON field names and inferred data types
-fn generate_schema(spec: HashMap<String, InferredType>) -> Result<Schema> {
- Ok(Schema::new(generate_fields(&spec)?))
-}
-
-/// JSON file reader that produces a serde_json::Value iterator from a Read trait
-///
-/// # Example
-///
-/// ```
-/// use std::fs::File;
-/// use std::io::BufReader;
-/// use arrow::json::reader::ValueIter;
-///
-/// let mut reader =
-/// BufReader::new(File::open("test/data/mixed_arrays.json").unwrap());
-/// let mut value_reader = ValueIter::new(&mut reader, None);
-/// for value in value_reader {
-/// println!("JSON value: {}", value.unwrap());
-/// }
-/// ```
-#[derive(Debug)]
-pub struct ValueIter<'a, R: Read> {
- reader: &'a mut BufReader<R>,
- max_read_records: Option<usize>,
- record_count: usize,
- // reuse line buffer to avoid allocation on each record
- line_buf: String,
-}
-
-impl<'a, R: Read> ValueIter<'a, R> {
- pub fn new(reader: &'a mut BufReader<R>, max_read_records: Option<usize>) -> Self {
- Self {
- reader,
- max_read_records,
- record_count: 0,
- line_buf: String::new(),
- }
- }
-}
-
-impl<'a, R: Read> Iterator for ValueIter<'a, R> {
- type Item = Result<Value>;
-
- fn next(&mut self) -> Option<Self::Item> {
- if let Some(max) = self.max_read_records {
- if self.record_count >= max {
- return None;
- }
- }
-
- loop {
- self.line_buf.truncate(0);
- match self.reader.read_line(&mut self.line_buf) {
- Ok(0) => {
- // read_line returns 0 when stream reached EOF
- return None;
- }
- Err(e) => {
- return Some(Err(ArrowError::JsonError(format!(
- "Failed to read JSON record: {}",
- e
- ))));
- }
- _ => {
- let trimmed_s = self.line_buf.trim();
- if trimmed_s.is_empty() {
- // ignore empty lines
- continue;
- }
-
- self.record_count += 1;
- return Some(serde_json::from_str(trimmed_s).map_err(|e| {
- ArrowError::JsonError(format!("Not valid JSON: {}", e))
- }));
- }
- }
- }
- }
-}
-
-/// Infer the fields of a JSON file by reading the first n records of the file, with
-/// `max_read_records` controlling the maximum number of records to read.
-///
-/// If `max_read_records` is not set, the whole file is read to infer its field types.
-///
-/// Contrary to [`infer_json_schema`], this function will seek back to the start of the `reader`.
-/// That way, the `reader` can be used immediately afterwards to create a [`Reader`].
-///
-/// # Examples
-/// ```
-/// use std::fs::File;
-/// use std::io::BufReader;
-/// use arrow::json::reader::infer_json_schema_from_seekable;
-///
-/// let file = File::open("test/data/mixed_arrays.json").unwrap();
-/// // file's cursor's offset at 0
-/// let mut reader = BufReader::new(file);
-/// let inferred_schema = infer_json_schema_from_seekable(&mut reader, None).unwrap();
-/// // file's cursor's offset automatically set at 0
-/// ```
-pub fn infer_json_schema_from_seekable<R: Read + Seek>(
- reader: &mut BufReader<R>,
- max_read_records: Option<usize>,
-) -> Result<Schema> {
- let schema = infer_json_schema(reader, max_read_records);
- // return the reader seek back to the start
- reader.seek(SeekFrom::Start(0))?;
-
- schema
-}
-
-/// Infer the fields of a JSON file by reading the first n records of the buffer, with
-/// `max_read_records` controlling the maximum number of records to read.
-///
-/// If `max_read_records` is not set, the whole file is read to infer its field types.
-///
-/// This function will not seek back to the start of the `reader`. The user has to manage the
-/// original file's cursor. This function is useful when the `reader`'s cursor is not available
-/// (does not implement [`Seek`]), such is the case for compressed streams decoders.
-///
-/// # Examples
-/// ```
-/// use std::fs::File;
-/// use std::io::{BufReader, SeekFrom, Seek};
-/// use flate2::read::GzDecoder;
-/// use arrow::json::reader::infer_json_schema;
-///
-/// let mut file = File::open("test/data/mixed_arrays.json.gz").unwrap();
-///
-/// // file's cursor's offset at 0
-/// let mut reader = BufReader::new(GzDecoder::new(&file));
-/// let inferred_schema = infer_json_schema(&mut reader, None).unwrap();
-/// // cursor's offset at end of file
-///
-/// // seek back to start so that the original file is usable again
-/// file.seek(SeekFrom::Start(0)).unwrap();
-/// ```
-pub fn infer_json_schema<R: Read>(
- reader: &mut BufReader<R>,
- max_read_records: Option<usize>,
-) -> Result<Schema> {
- infer_json_schema_from_iterator(ValueIter::new(reader, max_read_records))
-}
-
-fn set_object_scalar_field_type(
- field_types: &mut HashMap<String, InferredType>,
- key: &str,
- ftype: DataType,
-) -> Result<()> {
- if !field_types.contains_key(key) {
- field_types.insert(key.to_string(), InferredType::Scalar(HashSet::new()));
- }
-
- match field_types.get_mut(key).unwrap() {
- InferredType::Scalar(hs) => {
- hs.insert(ftype);
- Ok(())
- }
- // in case of column contains both scalar type and scalar array type, we convert type of
- // this column to scalar array.
- scalar_array @ InferredType::Array(_) => {
- let mut hs = HashSet::new();
- hs.insert(ftype);
- scalar_array.merge(InferredType::Scalar(hs))?;
- Ok(())
- }
- t => Err(ArrowError::JsonError(format!(
- "Expected scalar or scalar array JSON type, found: {:?}",
- t,
- ))),
- }
-}
-
-fn infer_scalar_array_type(array: &[Value]) -> Result<InferredType> {
- let mut hs = HashSet::new();
-
- for v in array {
- match v {
- Value::Null => {}
- Value::Number(n) => {
- if n.is_i64() {
- hs.insert(DataType::Int64);
- } else {
- hs.insert(DataType::Float64);
- }
- }
- Value::Bool(_) => {
- hs.insert(DataType::Boolean);
- }
- Value::String(_) => {
- hs.insert(DataType::Utf8);
- }
- Value::Array(_) | Value::Object(_) => {
- return Err(ArrowError::JsonError(format!(
- "Expected scalar value for scalar array, got: {:?}",
- v
- )));
- }
- }
- }
-
- Ok(InferredType::Scalar(hs))
-}
-
-fn infer_nested_array_type(array: &[Value]) -> Result<InferredType> {
- let mut inner_ele_type = InferredType::Any;
-
- for v in array {
- match v {
- Value::Array(inner_array) => {
- inner_ele_type.merge(infer_array_element_type(inner_array)?)?;
- }
- x => {
- return Err(ArrowError::JsonError(format!(
- "Got non array element in nested array: {:?}",
- x
- )));
- }
- }
- }
-
- Ok(InferredType::Array(Box::new(inner_ele_type)))
-}
-
-fn infer_struct_array_type(array: &[Value]) -> Result<InferredType> {
- let mut field_types = HashMap::new();
-
- for v in array {
- match v {
- Value::Object(map) => {
- collect_field_types_from_object(&mut field_types, map)?;
- }
- _ => {
- return Err(ArrowError::JsonError(format!(
- "Expected struct value for struct array, got: {:?}",
- v
- )));
- }
- }
- }
-
- Ok(InferredType::Object(field_types))
-}
-
-fn infer_array_element_type(array: &[Value]) -> Result<InferredType> {
- match array.iter().take(1).next() {
- None => Ok(InferredType::Any), // empty array, return any type that can be updated later
- Some(a) => match a {
- Value::Array(_) => infer_nested_array_type(array),
- Value::Object(_) => infer_struct_array_type(array),
- _ => infer_scalar_array_type(array),
- },
- }
-}
-
-fn collect_field_types_from_object(
- field_types: &mut HashMap<String, InferredType>,
- map: &JsonMap<String, Value>,
-) -> Result<()> {
- for (k, v) in map {
- match v {
- Value::Array(array) => {
- let ele_type = infer_array_element_type(array)?;
-
- if !field_types.contains_key(k) {
- match ele_type {
- InferredType::Scalar(_) => {
- field_types.insert(
- k.to_string(),
- InferredType::Array(Box::new(InferredType::Scalar(
- HashSet::new(),
- ))),
- );
- }
- InferredType::Object(_) => {
- field_types.insert(
- k.to_string(),
- InferredType::Array(Box::new(InferredType::Object(
- HashMap::new(),
- ))),
- );
- }
- InferredType::Any | InferredType::Array(_) => {
- // set inner type to any for nested array as well
- // so it can be updated properly from subsequent type merges
- field_types.insert(
- k.to_string(),
- InferredType::Array(Box::new(InferredType::Any)),
- );
- }
- }
- }
-
- match field_types.get_mut(k).unwrap() {
- InferredType::Array(inner_type) => {
- inner_type.merge(ele_type)?;
- }
- // in case of column contains both scalar type and scalar array type, we
- // convert type of this column to scalar array.
- field_type @ InferredType::Scalar(_) => {
- field_type.merge(ele_type)?;
- *field_type = InferredType::Array(Box::new(field_type.clone()));
- }
- t => {
- return Err(ArrowError::JsonError(format!(
- "Expected array json type, found: {:?}",
- t,
- )));
- }
- }
- }
- Value::Bool(_) => {
- set_object_scalar_field_type(field_types, k, DataType::Boolean)?;
- }
- Value::Null => {
- // do nothing, we treat json as nullable by default when
- // inferring
- }
- Value::Number(n) => {
- if n.is_f64() {
- set_object_scalar_field_type(field_types, k, DataType::Float64)?;
- } else {
- // default to i64
- set_object_scalar_field_type(field_types, k, DataType::Int64)?;
- }
- }
- Value::String(_) => {
- set_object_scalar_field_type(field_types, k, DataType::Utf8)?;
- }
- Value::Object(inner_map) => {
- if !field_types.contains_key(k) {
- field_types
- .insert(k.to_string(), InferredType::Object(HashMap::new()));
- }
- match field_types.get_mut(k).unwrap() {
- InferredType::Object(inner_field_types) => {
- collect_field_types_from_object(inner_field_types, inner_map)?;
- }
- t => {
- return Err(ArrowError::JsonError(format!(
- "Expected object json type, found: {:?}",
- t,
- )));
- }
- }
- }
- }
- }
-
- Ok(())
-}
-
-/// Infer the fields of a JSON file by reading all items from the JSON Value Iterator.
-///
-/// The following type coercion logic is implemented:
-/// * `Int64` and `Float64` are converted to `Float64`
-/// * Lists and scalars are coerced to a list of a compatible scalar
-/// * All other cases are coerced to `Utf8` (String)
-///
-/// Note that the above coercion logic is different from what Spark has, where it would default to
-/// String type in case of List and Scalar values appeared in the same field.
-///
-/// The reason we diverge here is because we don't have utilities to deal with JSON data once it's
-/// interpreted as Strings. We should match Spark's behavior once we added more JSON parsing
-/// kernels in the future.
-pub fn infer_json_schema_from_iterator<I>(value_iter: I) -> Result<Schema>
-where
- I: Iterator<Item = Result<Value>>,
-{
- let mut field_types: HashMap<String, InferredType> = HashMap::new();
-
- for record in value_iter {
- match record? {
- Value::Object(map) => {
- collect_field_types_from_object(&mut field_types, &map)?;
- }
- value => {
- return Err(ArrowError::JsonError(format!(
- "Expected JSON record to be an object, found {:?}",
- value
- )));
- }
- };
- }
-
- generate_schema(field_types)
-}
-
-/// JSON values to Arrow record batch decoder. Decoder's next_batch method takes a JSON Value
-/// iterator as input and outputs Arrow record batch.
-///
-/// # Examples
-/// ```
-/// use arrow::json::reader::{Decoder, ValueIter, infer_json_schema};
-/// use std::fs::File;
-/// use std::io::{BufReader, Seek, SeekFrom};
-/// use std::sync::Arc;
-///
-/// let mut reader =
-/// BufReader::new(File::open("test/data/mixed_arrays.json").unwrap());
-/// let inferred_schema = infer_json_schema(&mut reader, None).unwrap();
-/// let batch_size = 1024;
-/// let decoder = Decoder::new(Arc::new(inferred_schema), batch_size, None);
-///
-/// // seek back to start so that the original file is usable again
-/// reader.seek(SeekFrom::Start(0)).unwrap();
-/// let mut value_reader = ValueIter::new(&mut reader, None);
-/// let batch = decoder.next_batch(&mut value_reader).unwrap().unwrap();
-/// assert_eq!(4, batch.num_rows());
-/// assert_eq!(4, batch.num_columns());
-/// ```
-#[derive(Debug)]
-pub struct Decoder {
- /// Explicit schema for the JSON file
- schema: SchemaRef,
- /// Optional projection for which columns to load (case-sensitive names)
- projection: Option<Vec<String>>,
- /// Batch size (number of records to load each time)
- batch_size: usize,
-}
-
-impl Decoder {
- /// Create a new JSON decoder from any value that implements the `Iterator<Item=Result<Value>>`
- /// trait.
- pub fn new(
- schema: SchemaRef,
- batch_size: usize,
- projection: Option<Vec<String>>,
- ) -> Self {
- Self {
- schema,
- projection,
- batch_size,
- }
- }
-
- /// Returns the schema of the reader, useful for getting the schema without reading
- /// record batches
- pub fn schema(&self) -> SchemaRef {
- match &self.projection {
- Some(projection) => {
- let fields = self.schema.fields();
- let projected_fields: Vec<Field> = fields
- .iter()
- .filter_map(|field| {
- if projection.contains(field.name()) {
- Some(field.clone())
- } else {
- None
- }
- })
- .collect();
-
- Arc::new(Schema::new(projected_fields))
- }
- None => self.schema.clone(),
- }
- }
-
- /// Read the next batch of records
- pub fn next_batch<I>(&self, value_iter: &mut I) -> Result<Option<RecordBatch>>
- where
- I: Iterator<Item = Result<Value>>,
- {
- let mut rows: Vec<Value> = Vec::with_capacity(self.batch_size);
-
- for value in value_iter.by_ref().take(self.batch_size) {
- let v = value?;
- match v {
- Value::Object(_) => rows.push(v),
- _ => {
- return Err(ArrowError::JsonError(format!(
- "Row needs to be of type object, got: {:?}",
- v
- )));
- }
- }
- }
- if rows.is_empty() {
- // reached end of file
- return Ok(None);
- }
-
- let rows = &rows[..];
- let projection = self.projection.clone().unwrap_or_else(Vec::new);
- let arrays = self.build_struct_array(rows, self.schema.fields(), &projection);
-
- let projected_fields: Vec<Field> = if projection.is_empty() {
- self.schema.fields().to_vec()
- } else {
- projection
- .iter()
- .map(|name| self.schema.column_with_name(name))
- .flatten()
- .map(|(_, field)| field.clone())
- .collect()
- };
-
- let projected_schema = Arc::new(Schema::new(projected_fields));
-
- arrays.and_then(|arr| RecordBatch::try_new(projected_schema, arr).map(Some))
- }
-
- fn build_wrapped_list_array(
- &self,
- rows: &[Value],
- col_name: &str,
- key_type: &DataType,
- ) -> Result<ArrayRef> {
- match *key_type {
- DataType::Int8 => {
- let dtype = DataType::Dictionary(
- Box::new(DataType::Int8),
- Box::new(DataType::Utf8),
- );
- self.list_array_string_array_builder::<Int8Type>(&dtype, col_name, rows)
- }
- DataType::Int16 => {
- let dtype = DataType::Dictionary(
- Box::new(DataType::Int16),
- Box::new(DataType::Utf8),
- );
- self.list_array_string_array_builder::<Int16Type>(&dtype, col_name, rows)
- }
- DataType::Int32 => {
- let dtype = DataType::Dictionary(
- Box::new(DataType::Int32),
- Box::new(DataType::Utf8),
- );
- self.list_array_string_array_builder::<Int32Type>(&dtype, col_name, rows)
- }
- DataType::Int64 => {
- let dtype = DataType::Dictionary(
- Box::new(DataType::Int64),
- Box::new(DataType::Utf8),
- );
- self.list_array_string_array_builder::<Int64Type>(&dtype, col_name, rows)
- }
- DataType::UInt8 => {
- let dtype = DataType::Dictionary(
- Box::new(DataType::UInt8),
- Box::new(DataType::Utf8),
- );
- self.list_array_string_array_builder::<UInt8Type>(&dtype, col_name, rows)
- }
- DataType::UInt16 => {
- let dtype = DataType::Dictionary(
- Box::new(DataType::UInt16),
- Box::new(DataType::Utf8),
- );
- self.list_array_string_array_builder::<UInt16Type>(&dtype, col_name, rows)
- }
- DataType::UInt32 => {
- let dtype = DataType::Dictionary(
- Box::new(DataType::UInt32),
- Box::new(DataType::Utf8),
- );
- self.list_array_string_array_builder::<UInt32Type>(&dtype, col_name, rows)
- }
- DataType::UInt64 => {
- let dtype = DataType::Dictionary(
- Box::new(DataType::UInt64),
- Box::new(DataType::Utf8),
- );
- self.list_array_string_array_builder::<UInt64Type>(&dtype, col_name, rows)
- }
- ref e => Err(ArrowError::JsonError(format!(
- "Data type is currently not supported for dictionaries in list : {:?}",
- e
- ))),
- }
- }
-
- #[inline(always)]
- fn list_array_string_array_builder<DICT_TY>(
- &self,
- data_type: &DataType,
- col_name: &str,
- rows: &[Value],
- ) -> Result<ArrayRef>
- where
- DICT_TY: ArrowPrimitiveType + ArrowDictionaryKeyType,
- {
- let mut builder: Box<dyn ArrayBuilder> = match data_type {
- DataType::Utf8 => {
- let values_builder = StringBuilder::new(rows.len() * 5);
- Box::new(ListBuilder::new(values_builder))
- }
- DataType::Dictionary(_, _) => {
- let values_builder =
- self.build_string_dictionary_builder::<DICT_TY>(rows.len() * 5)?;
- Box::new(ListBuilder::new(values_builder))
- }
- e => {
- return Err(ArrowError::JsonError(format!(
- "Nested list data builder type is not supported: {:?}",
- e
- )))
- }
- };
-
- for row in rows {
- if let Some(value) = row.get(col_name) {
- // value can be an array or a scalar
- let vals: Vec<Option<String>> = if let Value::String(v) = value {
- vec![Some(v.to_string())]
- } else if let Value::Array(n) = value {
- n.iter()
- .map(|v: &Value| {
- if v.is_string() {
- Some(v.as_str().unwrap().to_string())
- } else if v.is_array() || v.is_object() || v.is_null() {
- // implicitly drop nested values
- // TODO support deep-nesting
- None
- } else {
- Some(v.to_string())
- }
- })
- .collect()
- } else if let Value::Null = value {
- vec![None]
- } else if !value.is_object() {
- vec![Some(value.to_string())]
- } else {
- return Err(ArrowError::JsonError(
- "Only scalars are currently supported in JSON arrays".to_string(),
- ));
- };
-
- // TODO: ARROW-10335: APIs of dictionary arrays and others are different. Unify
- // them.
- match data_type {
- DataType::Utf8 => {
- let builder = builder
- .as_any_mut()
- .downcast_mut::<ListBuilder<StringBuilder>>()
- .ok_or_else(||ArrowError::JsonError(
- "Cast failed for ListBuilder<StringBuilder> during nested data parsing".to_string(),
- ))?;
- for val in vals {
- if let Some(v) = val {
- builder.values().append_value(&v)?
- } else {
- builder.values().append_null()?
- };
- }
-
- // Append to the list
- builder.append(true)?;
- }
- DataType::Dictionary(_, _) => {
- let builder = builder.as_any_mut().downcast_mut::<ListBuilder<StringDictionaryBuilder<DICT_TY>>>().ok_or_else(||ArrowError::JsonError(
- "Cast failed for ListBuilder<StringDictionaryBuilder> during nested data parsing".to_string(),
- ))?;
- for val in vals {
- if let Some(v) = val {
- let _ = builder.values().append(&v)?;
- } else {
- builder.values().append_null()?
- };
- }
-
- // Append to the list
- builder.append(true)?;
- }
- e => {
- return Err(ArrowError::JsonError(format!(
- "Nested list data builder type is not supported: {:?}",
- e
- )))
- }
- }
- }
- }
-
- Ok(builder.finish() as ArrayRef)
- }
-
- #[inline(always)]
- #[allow(clippy::unnecessary_wraps)]
- fn build_string_dictionary_builder<T>(
- &self,
- row_len: usize,
- ) -> Result<StringDictionaryBuilder<T>>
- where
- T: ArrowPrimitiveType + ArrowDictionaryKeyType,
- {
- let key_builder = PrimitiveBuilder::<T>::new(row_len);
- let values_builder = StringBuilder::new(row_len * 5);
- Ok(StringDictionaryBuilder::new(key_builder, values_builder))
- }
-
- #[inline(always)]
- fn build_string_dictionary_array(
- &self,
- rows: &[Value],
- col_name: &str,
- key_type: &DataType,
- value_type: &DataType,
- ) -> Result<ArrayRef> {
- if let DataType::Utf8 = *value_type {
- match *key_type {
- DataType::Int8 => self.build_dictionary_array::<Int8Type>(rows, col_name),
- DataType::Int16 => {
- self.build_dictionary_array::<Int16Type>(rows, col_name)
- }
- DataType::Int32 => {
- self.build_dictionary_array::<Int32Type>(rows, col_name)
- }
- DataType::Int64 => {
- self.build_dictionary_array::<Int64Type>(rows, col_name)
- }
- DataType::UInt8 => {
- self.build_dictionary_array::<UInt8Type>(rows, col_name)
- }
- DataType::UInt16 => {
- self.build_dictionary_array::<UInt16Type>(rows, col_name)
- }
- DataType::UInt32 => {
- self.build_dictionary_array::<UInt32Type>(rows, col_name)
- }
- DataType::UInt64 => {
- self.build_dictionary_array::<UInt64Type>(rows, col_name)
- }
- _ => Err(ArrowError::JsonError(
- "unsupported dictionary key type".to_string(),
- )),
- }
- } else {
- Err(ArrowError::JsonError(
- "dictionary types other than UTF-8 not yet supported".to_string(),
- ))
- }
- }
-
- fn build_boolean_array(&self, rows: &[Value], col_name: &str) -> Result<ArrayRef> {
- let mut builder = BooleanBuilder::new(rows.len());
- for row in rows {
- if let Some(value) = row.get(&col_name) {
- if let Some(boolean) = value.as_bool() {
- builder.append_value(boolean)?
- } else {
- builder.append_null()?;
- }
- } else {
- builder.append_null()?;
- }
- }
- Ok(Arc::new(builder.finish()))
- }
-
- #[allow(clippy::unnecessary_wraps)]
- fn build_primitive_array<T: ArrowPrimitiveType>(
- &self,
- rows: &[Value],
- col_name: &str,
- ) -> Result<ArrayRef>
- where
- T: ArrowNumericType,
- T::Native: num::NumCast,
- {
- Ok(Arc::new(
- rows.iter()
- .map(|row| {
- row.get(&col_name)
- .and_then(|value| value.as_f64())
- .and_then(num::cast::cast)
- })
- .collect::<PrimitiveArray<T>>(),
- ))
- }
-
- /// Build a nested GenericListArray from a list of unnested `Value`s
- fn build_nested_list_array<OffsetSize: OffsetSizeTrait>(
- &self,
- rows: &[Value],
- list_field: &Field,
- ) -> Result<ArrayRef> {
- // build list offsets
- let mut cur_offset = OffsetSize::zero();
- let list_len = rows.len();
- let num_list_bytes = bit_util::ceil(list_len, 8);
- let mut offsets = Vec::with_capacity(list_len + 1);
- let mut list_nulls = MutableBuffer::from_len_zeroed(num_list_bytes);
- let list_nulls = list_nulls.as_slice_mut();
- offsets.push(cur_offset);
- rows.iter().enumerate().for_each(|(i, v)| {
- if let Value::Array(a) = v {
- cur_offset += OffsetSize::from_usize(a.len()).unwrap();
- bit_util::set_bit(list_nulls, i);
- } else if let Value::Null = v {
- // value is null, not incremented
- } else {
- cur_offset += OffsetSize::one();
- }
- offsets.push(cur_offset);
- });
- let valid_len = cur_offset.to_usize().unwrap();
- let array_data = match list_field.data_type() {
- DataType::Null => NullArray::new(valid_len).data().clone(),
- DataType::Boolean => {
- let num_bytes = bit_util::ceil(valid_len, 8);
- let mut bool_values = MutableBuffer::from_len_zeroed(num_bytes);
- let mut bool_nulls =
- MutableBuffer::new(num_bytes).with_bitset(num_bytes, true);
- let mut curr_index = 0;
- rows.iter().for_each(|v| {
- if let Value::Array(vs) = v {
- vs.iter().for_each(|value| {
- if let Value::Bool(child) = value {
- // if valid boolean, append value
- if *child {
- bit_util::set_bit(
- bool_values.as_slice_mut(),
- curr_index,
- );
- }
- } else {
- // null slot
- bit_util::unset_bit(
- bool_nulls.as_slice_mut(),
- curr_index,
- );
- }
- curr_index += 1;
- });
- }
- });
- ArrayData::builder(list_field.data_type().clone())
- .len(valid_len)
- .add_buffer(bool_values.into())
- .null_bit_buffer(bool_nulls.into())
- .build()
- }
- DataType::Int8 => self.read_primitive_list_values::<Int8Type>(rows),
- DataType::Int16 => self.read_primitive_list_values::<Int16Type>(rows),
- DataType::Int32 => self.read_primitive_list_values::<Int32Type>(rows),
- DataType::Int64 => self.read_primitive_list_values::<Int64Type>(rows),
- DataType::UInt8 => self.read_primitive_list_values::<UInt8Type>(rows),
- DataType::UInt16 => self.read_primitive_list_values::<UInt16Type>(rows),
- DataType::UInt32 => self.read_primitive_list_values::<UInt32Type>(rows),
- DataType::UInt64 => self.read_primitive_list_values::<UInt64Type>(rows),
- DataType::Float16 => {
- return Err(ArrowError::JsonError("Float16 not supported".to_string()))
- }
- DataType::Float32 => self.read_primitive_list_values::<Float32Type>(rows),
- DataType::Float64 => self.read_primitive_list_values::<Float64Type>(rows),
- DataType::Timestamp(_, _)
- | DataType::Date32
- | DataType::Date64
- | DataType::Time32(_)
- | DataType::Time64(_) => {
- return Err(ArrowError::JsonError(
- "Temporal types are not yet supported, see ARROW-4803".to_string(),
- ))
- }
- DataType::Utf8 => {
- StringArray::from_iter(flatten_json_string_values(rows).into_iter())
- .data()
- .clone()
- }
- DataType::LargeUtf8 => {
- LargeStringArray::from_iter(flatten_json_string_values(rows).into_iter())
- .data()
- .clone()
- }
- DataType::List(field) => {
- let child = self
- .build_nested_list_array::<i32>(&flatten_json_values(rows), field)?;
- child.data().clone()
- }
- DataType::LargeList(field) => {
- let child = self
- .build_nested_list_array::<i64>(&flatten_json_values(rows), field)?;
- child.data().clone()
- }
- DataType::Struct(fields) => {
- // extract list values, with non-lists converted to Value::Null
- let array_item_count = rows
- .iter()
- .map(|row| match row {
- Value::Array(values) => values.len(),
- _ => 1,
- })
- .sum();
- let num_bytes = bit_util::ceil(array_item_count, 8);
- let mut null_buffer = MutableBuffer::from_len_zeroed(num_bytes);
- let mut struct_index = 0;
- let rows: Vec<Value> = rows
- .iter()
- .flat_map(|row| {
- if let Value::Array(values) = row {
- values.iter().for_each(|_| {
- bit_util::set_bit(
- null_buffer.as_slice_mut(),
- struct_index,
- );
- struct_index += 1;
- });
- values.clone()
- } else {
- struct_index += 1;
- vec![Value::Null]
- }
- })
- .collect();
- let arrays =
- self.build_struct_array(rows.as_slice(), fields.as_slice(), &[])?;
- let data_type = DataType::Struct(fields.clone());
- let buf = null_buffer.into();
- ArrayDataBuilder::new(data_type)
- .len(rows.len())
- .null_bit_buffer(buf)
- .child_data(arrays.into_iter().map(|a| a.data().clone()).collect())
- .build()
- }
- datatype => {
- return Err(ArrowError::JsonError(format!(
- "Nested list of {:?} not supported",
- datatype
- )));
- }
- };
- // build list
- let list_data = ArrayData::builder(DataType::List(Box::new(list_field.clone())))
- .len(list_len)
- .add_buffer(Buffer::from_slice_ref(&offsets))
- .add_child_data(array_data)
- .null_bit_buffer(list_nulls.into())
- .build();
- Ok(Arc::new(GenericListArray::<OffsetSize>::from(list_data)))
- }
-
- /// Builds the child values of a `StructArray`, falling short of constructing the StructArray.
- /// The function does not construct the StructArray as some callers would want the child arrays.
- ///
- /// *Note*: The function is recursive, and will read nested structs.
- ///
- /// If `projection` is not empty, then all values are returned. The first level of projection
- /// occurs at the `RecordBatch` level. No further projection currently occurs, but would be
- /// useful if plucking values from a struct, e.g. getting `a.b.c.e` from `a.b.c.{d, e}`.
- fn build_struct_array(
- &self,
- rows: &[Value],
- struct_fields: &[Field],
- projection: &[String],
- ) -> Result<Vec<ArrayRef>> {
- let arrays: Result<Vec<ArrayRef>> = struct_fields
- .iter()
- .filter(|field| projection.is_empty() || projection.contains(field.name()))
- .map(|field| {
- match field.data_type() {
- DataType::Null => {
- Ok(Arc::new(NullArray::new(rows.len())) as ArrayRef)
- }
- DataType::Boolean => self.build_boolean_array(rows, field.name()),
- DataType::Float64 => {
- self.build_primitive_array::<Float64Type>(rows, field.name())
- }
- DataType::Float32 => {
- self.build_primitive_array::<Float32Type>(rows, field.name())
- }
- DataType::Int64 => {
- self.build_primitive_array::<Int64Type>(rows, field.name())
- }
- DataType::Int32 => {
- self.build_primitive_array::<Int32Type>(rows, field.name())
- }
- DataType::Int16 => {
- self.build_primitive_array::<Int16Type>(rows, field.name())
- }
- DataType::Int8 => {
- self.build_primitive_array::<Int8Type>(rows, field.name())
- }
- DataType::UInt64 => {
- self.build_primitive_array::<UInt64Type>(rows, field.name())
- }
- DataType::UInt32 => {
- self.build_primitive_array::<UInt32Type>(rows, field.name())
- }
- DataType::UInt16 => {
- self.build_primitive_array::<UInt16Type>(rows, field.name())
- }
- DataType::UInt8 => {
- self.build_primitive_array::<UInt8Type>(rows, field.name())
- }
- // TODO: this is incomplete
- DataType::Timestamp(unit, _) => match unit {
- TimeUnit::Second => self
- .build_primitive_array::<TimestampSecondType>(
- rows,
- field.name(),
- ),
- TimeUnit::Microsecond => self
- .build_primitive_array::<TimestampMicrosecondType>(
- rows,
- field.name(),
- ),
- TimeUnit::Millisecond => self
- .build_primitive_array::<TimestampMillisecondType>(
- rows,
- field.name(),
- ),
- TimeUnit::Nanosecond => self
- .build_primitive_array::<TimestampNanosecondType>(
- rows,
- field.name(),
- ),
- },
- DataType::Date64 => {
- self.build_primitive_array::<Date64Type>(rows, field.name())
- }
- DataType::Date32 => {
- self.build_primitive_array::<Date32Type>(rows, field.name())
- }
- DataType::Time64(unit) => match unit {
- TimeUnit::Microsecond => self
- .build_primitive_array::<Time64MicrosecondType>(
- rows,
- field.name(),
- ),
- TimeUnit::Nanosecond => self
- .build_primitive_array::<Time64NanosecondType>(
- rows,
- field.name(),
- ),
- t => Err(ArrowError::JsonError(format!(
- "TimeUnit {:?} not supported with Time64",
- t
- ))),
- },
- DataType::Time32(unit) => match unit {
- TimeUnit::Second => self
- .build_primitive_array::<Time32SecondType>(
- rows,
- field.name(),
- ),
- TimeUnit::Millisecond => self
- .build_primitive_array::<Time32MillisecondType>(
- rows,
- field.name(),
- ),
- t => Err(ArrowError::JsonError(format!(
- "TimeUnit {:?} not supported with Time32",
- t
- ))),
- },
- DataType::Utf8 => Ok(Arc::new(
- rows.iter()
- .map(|row| {
- let maybe_value = row.get(field.name());
- maybe_value.and_then(|value| value.as_str())
- })
- .collect::<StringArray>(),
- ) as ArrayRef),
- DataType::List(ref list_field) => {
- match list_field.data_type() {
- DataType::Dictionary(ref key_ty, _) => {
- self.build_wrapped_list_array(rows, field.name(), key_ty)
- }
- _ => {
- // extract rows by name
- let extracted_rows = rows
- .iter()
- .map(|row| {
- row.get(field.name())
- .cloned()
- .unwrap_or(Value::Null)
- })
- .collect::<Vec<Value>>();
- self.build_nested_list_array::<i32>(
- extracted_rows.as_slice(),
- list_field,
- )
- }
- }
- }
- DataType::Dictionary(ref key_ty, ref val_ty) => self
- .build_string_dictionary_array(
- rows,
- field.name(),
- key_ty,
- val_ty,
- ),
- DataType::Struct(fields) => {
- let len = rows.len();
- let num_bytes = bit_util::ceil(len, 8);
- let mut null_buffer = MutableBuffer::from_len_zeroed(num_bytes);
- let struct_rows = rows
- .iter()
- .enumerate()
- .map(|(i, row)| {
- (
- i,
- row.as_object()
- .map(|v| v.get(field.name()))
- .flatten(),
- )
- })
- .map(|(i, v)| match v {
- // we want the field as an object, if it's not, we treat as null
- Some(Value::Object(value)) => {
- bit_util::set_bit(null_buffer.as_slice_mut(), i);
- Value::Object(value.clone())
- }
- _ => Value::Object(Default::default()),
- })
- .collect::<Vec<Value>>();
- let arrays =
- self.build_struct_array(&struct_rows, fields, &[])?;
- // construct a struct array's data in order to set null buffer
- let data_type = DataType::Struct(fields.clone());
- let data = ArrayDataBuilder::new(data_type)
- .len(len)
- .null_bit_buffer(null_buffer.into())
- .child_data(
- arrays.into_iter().map(|a| a.data().clone()).collect(),
- )
- .build();
- Ok(make_array(data))
- }
- _ => Err(ArrowError::JsonError(format!(
- "{:?} type is not supported",
- field.data_type()
- ))),
- }
- })
- .collect();
- arrays
- }
-
- #[inline(always)]
- fn build_dictionary_array<T>(
- &self,
- rows: &[Value],
- col_name: &str,
- ) -> Result<ArrayRef>
- where
- T::Native: num::NumCast,
- T: ArrowPrimitiveType + ArrowDictionaryKeyType,
- {
- let mut builder: StringDictionaryBuilder<T> =
- self.build_string_dictionary_builder(rows.len())?;
- for row in rows {
- if let Some(value) = row.get(&col_name) {
- if let Some(str_v) = value.as_str() {
- builder.append(str_v).map(drop)?
- } else {
- builder.append_null()?
- }
- } else {
- builder.append_null()?
- }
- }
- Ok(Arc::new(builder.finish()) as ArrayRef)
- }
-
- /// Read the primitive list's values into ArrayData
- fn read_primitive_list_values<T>(&self, rows: &[Value]) -> ArrayData
- where
- T: ArrowPrimitiveType + ArrowNumericType,
- T::Native: num::NumCast,
- {
- let values = rows
- .iter()
- .flat_map(|row| {
- // read values from list
- if let Value::Array(values) = row {
- values
- .iter()
- .map(|value| {
- let v: Option<T::Native> =
- value.as_f64().and_then(num::cast::cast);
- v
- })
- .collect::<Vec<Option<T::Native>>>()
- } else if let Value::Number(value) = row {
- // handle the scalar number case
- let v: Option<T::Native> = value.as_f64().and_then(num::cast::cast);
- v.map(|v| vec![Some(v)]).unwrap_or_default()
- } else {
- vec![]
- }
- })
- .collect::<Vec<Option<T::Native>>>();
- let array = PrimitiveArray::<T>::from_iter(values.iter());
- array.data().clone()
- }
-}
-
-/// Reads a JSON value as a string, regardless of its type.
-/// This is useful if the expected datatype is a string, in which case we preserve
-/// all the values regardless of they type.
-///
-/// Applying `value.to_string()` unfortunately results in an escaped string, which
-/// is not what we want.
-#[inline(always)]
-fn json_value_as_string(value: &Value) -> Option<String> {
- match value {
- Value::Null => None,
- Value::String(string) => Some(string.clone()),
- _ => Some(value.to_string()),
- }
-}
-
-/// Flattens a list of JSON values, by flattening lists, and treating all other values as
-/// single-value lists.
-/// This is used to read into nested lists (list of list, list of struct) and non-dictionary lists.
-#[inline]
-fn flatten_json_values(values: &[Value]) -> Vec<Value> {
- values
- .iter()
- .flat_map(|row| {
- if let Value::Array(values) = row {
- values.clone()
- } else if let Value::Null = row {
- vec![Value::Null]
- } else {
- // we interpret a scalar as a single-value list to minimise data loss
- vec![row.clone()]
- }
- })
- .collect()
-}
-
-/// Flattens a list into string values, dropping Value::Null in the process.
-/// This is useful for interpreting any JSON array as string, dropping nulls.
-/// See `json_value_as_string`.
-#[inline]
-fn flatten_json_string_values(values: &[Value]) -> Vec<Option<String>> {
- values
- .iter()
- .flat_map(|row| {
- if let Value::Array(values) = row {
- values
- .iter()
- .map(json_value_as_string)
- .collect::<Vec<Option<_>>>()
- } else if let Value::Null = row {
- vec![]
- } else {
- vec![json_value_as_string(row)]
- }
- })
- .collect::<Vec<Option<_>>>()
-}
-/// JSON file reader
-#[derive(Debug)]
-pub struct Reader<R: Read> {
- reader: BufReader<R>,
- /// JSON value decoder
- decoder: Decoder,
-}
-
-impl<R: Read> Reader<R> {
- /// Create a new JSON Reader from any value that implements the `Read` trait.
- ///
- /// If reading a `File`, you can customise the Reader, such as to enable schema
- /// inference, use `ReaderBuilder`.
- pub fn new(
- reader: R,
- schema: SchemaRef,
- batch_size: usize,
- projection: Option<Vec<String>>,
- ) -> Self {
- Self::from_buf_reader(BufReader::new(reader), schema, batch_size, projection)
- }
-
- /// Create a new JSON Reader from a `BufReader<R: Read>`
- ///
- /// To customize the schema, such as to enable schema inference, use `ReaderBuilder`
- pub fn from_buf_reader(
- reader: BufReader<R>,
- schema: SchemaRef,
- batch_size: usize,
- projection: Option<Vec<String>>,
- ) -> Self {
- Self {
- reader,
- decoder: Decoder::new(schema, batch_size, projection),
- }
- }
-
- /// Returns the schema of the reader, useful for getting the schema without reading
- /// record batches
- pub fn schema(&self) -> SchemaRef {
- self.decoder.schema()
- }
-
- /// Read the next batch of records
- #[allow(clippy::should_implement_trait)]
- pub fn next(&mut self) -> Result<Option<RecordBatch>> {
- self.decoder
- .next_batch(&mut ValueIter::new(&mut self.reader, None))
- }
-}
-
-/// JSON file reader builder
-#[derive(Debug)]
-pub struct ReaderBuilder {
- /// Optional schema for the JSON file
- ///
- /// If the schema is not supplied, the reader will try to infer the schema
- /// based on the JSON structure.
- schema: Option<SchemaRef>,
- /// Optional maximum number of records to read during schema inference
- ///
- /// If a number is not provided, all the records are read.
- max_records: Option<usize>,
- /// Batch size (number of records to load each time)
- ///
- /// The default batch size when using the `ReaderBuilder` is 1024 records
- batch_size: usize,
- /// Optional projection for which columns to load (zero-based column indices)
- projection: Option<Vec<String>>,
-}
-
-impl Default for ReaderBuilder {
- fn default() -> Self {
- Self {
- schema: None,
- max_records: None,
- batch_size: 1024,
- projection: None,
- }
- }
-}
-
-impl ReaderBuilder {
- /// Create a new builder for configuring JSON parsing options.
- ///
- /// To convert a builder into a reader, call `Reader::from_builder`
- ///
- /// # Example
- ///
- /// ```
- /// extern crate arrow;
- ///
- /// use arrow::json;
- /// use std::fs::File;
- ///
- /// fn example() -> json::Reader<File> {
- /// let file = File::open("test/data/basic.json").unwrap();
- ///
- /// // create a builder, inferring the schema with the first 100 records
- /// let builder = json::ReaderBuilder::new().infer_schema(Some(100));
- ///
- /// let reader = builder.build::<File>(file).unwrap();
- ///
- /// reader
- /// }
- /// ```
- pub fn new() -> Self {
- Self::default()
- }
-
- /// Set the JSON file's schema
- pub fn with_schema(mut self, schema: SchemaRef) -> Self {
- self.schema = Some(schema);
- self
- }
-
- /// Set the JSON reader to infer the schema of the file
- pub fn infer_schema(mut self, max_records: Option<usize>) -> Self {
- // remove any schema that is set
- self.schema = None;
- self.max_records = max_records;
- self
- }
-
- /// Set the batch size (number of records to load at one time)
- pub fn with_batch_size(mut self, batch_size: usize) -> Self {
- self.batch_size = batch_size;
- self
- }
-
- /// Set the reader's column projection
- pub fn with_projection(mut self, projection: Vec<String>) -> Self {
- self.projection = Some(projection);
- self
- }
-
- /// Create a new `Reader` from the `ReaderBuilder`
- pub fn build<R>(self, source: R) -> Result<Reader<R>>
- where
- R: Read + Seek,
- {
- let mut buf_reader = BufReader::new(source);
-
- // check if schema should be inferred
- let schema = match self.schema {
- Some(schema) => schema,
- None => Arc::new(infer_json_schema_from_seekable(
- &mut buf_reader,
- self.max_records,
- )?),
- };
-
- Ok(Reader::from_buf_reader(
- buf_reader,
- schema,
- self.batch_size,
- self.projection,
- ))
- }
-}
-
-impl<R: Read> Iterator for Reader<R> {
- type Item = Result<RecordBatch>;
-
- fn next(&mut self) -> Option<Self::Item> {
- self.next().transpose()
- }
-}
-
-#[cfg(test)]
-mod tests {
- use crate::{
- buffer::Buffer,
- datatypes::DataType::{Dictionary, List},
- };
-
- use super::*;
- use flate2::read::GzDecoder;
- use std::fs::File;
- use std::io::Cursor;
-
- #[test]
- fn test_json_basic() {
- let builder = ReaderBuilder::new().infer_schema(None).with_batch_size(64);
- let mut reader: Reader<File> = builder
- .build::<File>(File::open("test/data/basic.json").unwrap())
- .unwrap();
- let batch = reader.next().unwrap().unwrap();
-
- assert_eq!(4, batch.num_columns());
- assert_eq!(12, batch.num_rows());
-
- let schema = reader.schema();
- let batch_schema = batch.schema();
- assert_eq!(schema, batch_schema);
-
- let a = schema.column_with_name("a").unwrap();
- assert_eq!(0, a.0);
- assert_eq!(&DataType::Int64, a.1.data_type());
- let b = schema.column_with_name("b").unwrap();
- assert_eq!(1, b.0);
- assert_eq!(&DataType::Float64, b.1.data_type());
- let c = schema.column_with_name("c").unwrap();
- assert_eq!(2, c.0);
- assert_eq!(&DataType::Boolean, c.1.data_type());
- let d = schema.column_with_name("d").unwrap();
- assert_eq!(3, d.0);
- assert_eq!(&DataType::Utf8, d.1.data_type());
-
- let aa = batch
- .column(a.0)
- .as_any()
- .downcast_ref::<Int64Array>()
- .unwrap();
- assert_eq!(1, aa.value(0));
- assert_eq!(-10, aa.value(1));
- let bb = batch
- .column(b.0)
- .as_any()
- .downcast_ref::<Float64Array>()
- .unwrap();
- assert!(2.0 - bb.value(0) < f64::EPSILON);
- assert!(-3.5 - bb.value(1) < f64::EPSILON);
- let cc = batch
- .column(c.0)
- .as_any()
- .downcast_ref::<BooleanArray>()
- .unwrap();
- assert_eq!(false, cc.value(0));
- assert_eq!(true, cc.value(10));
- let dd = batch
- .column(d.0)
- .as_any()
- .downcast_ref::<StringArray>()
- .unwrap();
- assert_eq!("4", dd.value(0));
- assert_eq!("text", dd.value(8));
- }
-
- #[test]
- fn test_json_basic_with_nulls() {
- let builder = ReaderBuilder::new().infer_schema(None).with_batch_size(64);
- let mut reader: Reader<File> = builder
- .build::<File>(File::open("test/data/basic_nulls.json").unwrap())
- .unwrap();
- let batch = reader.next().unwrap().unwrap();
-
- assert_eq!(4, batch.num_columns());
- assert_eq!(12, batch.num_rows());
-
- let schema = reader.schema();
- let batch_schema = batch.schema();
- assert_eq!(schema, batch_schema);
-
- let a = schema.column_with_name("a").unwrap();
- assert_eq!(&DataType::Int64, a.1.data_type());
- let b = schema.column_with_name("b").unwrap();
- assert_eq!(&DataType::Float64, b.1.data_type());
- let c = schema.column_with_name("c").unwrap();
- assert_eq!(&DataType::Boolean, c.1.data_type());
- let d = schema.column_with_name("d").unwrap();
- assert_eq!(&DataType::Utf8, d.1.data_type());
-
- let aa = batch
- .column(a.0)
- .as_any()
- .downcast_ref::<Int64Array>()
- .unwrap();
- assert_eq!(true, aa.is_valid(0));
- assert_eq!(false, aa.is_valid(1));
- assert_eq!(false, aa.is_valid(11));
- let bb = batch
- .column(b.0)
- .as_any()
- .downcast_ref::<Float64Array>()
- .unwrap();
- assert_eq!(true, bb.is_valid(0));
- assert_eq!(false, bb.is_valid(2));
- assert_eq!(false, bb.is_valid(11));
- let cc = batch
- .column(c.0)
- .as_any()
- .downcast_ref::<BooleanArray>()
- .unwrap();
- assert_eq!(true, cc.is_valid(0));
- assert_eq!(false, cc.is_valid(4));
- assert_eq!(false, cc.is_valid(11));
- let dd = batch
- .column(d.0)
- .as_any()
- .downcast_ref::<StringArray>()
- .unwrap();
- assert_eq!(false, dd.is_valid(0));
- assert_eq!(true, dd.is_valid(1));
- assert_eq!(false, dd.is_valid(4));
- assert_eq!(false, dd.is_valid(11));
- }
-
- #[test]
- fn test_json_basic_schema() {
- let schema = Schema::new(vec![
- Field::new("a", DataType::Int32, false),
- Field::new("b", DataType::Float32, false),
- Field::new("c", DataType::Boolean, false),
- Field::new("d", DataType::Utf8, false),
- ]);
-
- let mut reader: Reader<File> = Reader::new(
- File::open("test/data/basic.json").unwrap(),
- Arc::new(schema.clone()),
- 1024,
- None,
- );
- let reader_schema = reader.schema();
- assert_eq!(reader_schema, Arc::new(schema));
- let batch = reader.next().unwrap().unwrap();
-
- assert_eq!(4, batch.num_columns());
- assert_eq!(12, batch.num_rows());
-
- let schema = batch.schema();
-
- let a = schema.column_with_name("a").unwrap();
- assert_eq!(&DataType::Int32, a.1.data_type());
- let b = schema.column_with_name("b").unwrap();
- assert_eq!(&DataType::Float32, b.1.data_type());
- let c = schema.column_with_name("c").unwrap();
- assert_eq!(&DataType::Boolean, c.1.data_type());
- let d = schema.column_with_name("d").unwrap();
- assert_eq!(&DataType::Utf8, d.1.data_type());
-
- let aa = batch
- .column(a.0)
- .as_any()
- .downcast_ref::<Int32Array>()
- .unwrap();
- assert_eq!(1, aa.value(0));
- // test that a 64bit value is returned as null due to overflowing
- assert_eq!(false, aa.is_valid(11));
- let bb = batch
- .column(b.0)
- .as_any()
- .downcast_ref::<Float32Array>()
- .unwrap();
- assert!(2.0 - bb.value(0) < f32::EPSILON);
- assert!(-3.5 - bb.value(1) < f32::EPSILON);
- }
-
- #[test]
- fn test_json_basic_schema_projection() {
- // We test implicit and explicit projection:
- // Implicit: omitting fields from a schema
- // Explicit: supplying a vec of fields to take
- let schema = Schema::new(vec![
- Field::new("a", DataType::Int32, false),
- Field::new("b", DataType::Float32, false),
- Field::new("c", DataType::Boolean, false),
- ]);
-
- let mut reader: Reader<File> = Reader::new(
- File::open("test/data/basic.json").unwrap(),
- Arc::new(schema),
- 1024,
- Some(vec!["a".to_string(), "c".to_string()]),
- );
- let reader_schema = reader.schema();
- let expected_schema = Arc::new(Schema::new(vec![
- Field::new("a", DataType::Int32, false),
- Field::new("c", DataType::Boolean, false),
- ]));
- assert_eq!(reader_schema, expected_schema);
-
- let batch = reader.next().unwrap().unwrap();
-
- assert_eq!(2, batch.num_columns());
- assert_eq!(2, batch.schema().fields().len());
- assert_eq!(12, batch.num_rows());
-
- let schema = batch.schema();
- assert_eq!(reader_schema, schema);
-
- let a = schema.column_with_name("a").unwrap();
- assert_eq!(0, a.0);
- assert_eq!(&DataType::Int32, a.1.data_type());
- let c = schema.column_with_name("c").unwrap();
- assert_eq!(1, c.0);
- assert_eq!(&DataType::Boolean, c.1.data_type());
- }
-
- #[test]
- fn test_json_arrays() {
- let builder = ReaderBuilder::new().infer_schema(None).with_batch_size(64);
- let mut reader: Reader<File> = builder
- .build::<File>(File::open("test/data/arrays.json").unwrap())
- .unwrap();
- let batch = reader.next().unwrap().unwrap();
-
- assert_eq!(4, batch.num_columns());
- assert_eq!(3, batch.num_rows());
-
- let schema = batch.schema();
-
- let a = schema.column_with_name("a").unwrap();
- assert_eq!(&DataType::Int64, a.1.data_type());
- let b = schema.column_with_name("b").unwrap();
- assert_eq!(
- &DataType::List(Box::new(Field::new("item", DataType::Float64, true))),
- b.1.data_type()
- );
- let c = schema.column_with_name("c").unwrap();
- assert_eq!(
- &DataType::List(Box::new(Field::new("item", DataType::Boolean, true))),
- c.1.data_type()
- );
- let d = schema.column_with_name("d").unwrap();
- assert_eq!(&DataType::Utf8, d.1.data_type());
-
- let aa = batch
- .column(a.0)
- .as_any()
- .downcast_ref::<Int64Array>()
- .unwrap();
- assert_eq!(1, aa.value(0));
- assert_eq!(-10, aa.value(1));
- let bb = batch
- .column(b.0)
- .as_any()
- .downcast_ref::<ListArray>()
- .unwrap();
- let bb = bb.values();
- let bb = bb.as_any().downcast_ref::<Float64Array>().unwrap();
- assert_eq!(9, bb.len());
- assert!(2.0 - bb.value(0) < f64::EPSILON);
- assert!(-6.1 - bb.value(5) < f64::EPSILON);
- assert_eq!(false, bb.is_valid(7));
-
- let cc = batch
- .column(c.0)
- .as_any()
- .downcast_ref::<ListArray>()
- .unwrap();
- let cc = cc.values();
- let cc = cc.as_any().downcast_ref::<BooleanArray>().unwrap();
- assert_eq!(6, cc.len());
- assert_eq!(false, cc.value(0));
- assert_eq!(false, cc.value(4));
- assert_eq!(false, cc.is_valid(5));
- }
-
- #[test]
- fn test_invalid_json_infer_schema() {
- let re = infer_json_schema_from_seekable(
- &mut BufReader::new(
- File::open("test/data/uk_cities_with_headers.csv").unwrap(),
- ),
- None,
- );
- assert_eq!(
- re.err().unwrap().to_string(),
- "Json error: Not valid JSON: expected value at line 1 column 1",
- );
- }
-
- #[test]
- fn test_invalid_json_read_record() {
- let schema = Arc::new(Schema::new(vec![Field::new(
- "a",
- DataType::Struct(vec![Field::new("a", DataType::Utf8, true)]),
- true,
- )]));
- let builder = ReaderBuilder::new().with_schema(schema).with_batch_size(64);
- let mut reader: Reader<File> = builder
- .build::<File>(File::open("test/data/uk_cities_with_headers.csv").unwrap())
- .unwrap();
- assert_eq!(
- reader.next().err().unwrap().to_string(),
- "Json error: Not valid JSON: expected value at line 1 column 1",
- );
- }
-
- #[test]
- fn test_coersion_scalar_and_list() {
- use crate::datatypes::DataType::*;
-
- assert_eq!(
- List(Box::new(Field::new("item", Float64, true))),
- coerce_data_type(vec![
- &Float64,
- &List(Box::new(Field::new("item", Float64, true)))
- ])
- );
- assert_eq!(
- List(Box::new(Field::new("item", Float64, true))),
- coerce_data_type(vec![
- &Float64,
- &List(Box::new(Field::new("item", Int64, true)))
- ])
- );
- assert_eq!(
- List(Box::new(Field::new("item", Int64, true))),
- coerce_data_type(vec![
- &Int64,
- &List(Box::new(Field::new("item", Int64, true)))
- ])
- );
- // boolean and number are incompatible, return utf8
- assert_eq!(
- List(Box::new(Field::new("item", Utf8, true))),
- coerce_data_type(vec![
- &Boolean,
- &List(Box::new(Field::new("item", Float64, true)))
- ])
- );
- }
-
- #[test]
- fn test_mixed_json_arrays() {
- let builder = ReaderBuilder::new().infer_schema(None).with_batch_size(64);
- let mut reader: Reader<File> = builder
- .build::<File>(File::open("test/data/mixed_arrays.json").unwrap())
- .unwrap();
- let batch = reader.next().unwrap().unwrap();
-
- let mut file = File::open("test/data/mixed_arrays.json.gz").unwrap();
- let mut reader = BufReader::new(GzDecoder::new(&file));
- let schema = infer_json_schema(&mut reader, None).unwrap();
- file.seek(SeekFrom::Start(0)).unwrap();
-
- let reader = BufReader::new(GzDecoder::new(&file));
- let mut reader = Reader::from_buf_reader(reader, Arc::new(schema), 64, None);
- let batch_gz = reader.next().unwrap().unwrap();
-
- for batch in vec![batch, batch_gz] {
- assert_eq!(4, batch.num_columns());
- assert_eq!(4, batch.num_rows());
-
- let schema = batch.schema();
-
- let a = schema.column_with_name("a").unwrap();
- assert_eq!(&DataType::Int64, a.1.data_type());
- let b = schema.column_with_name("b").unwrap();
- assert_eq!(
- &DataType::List(Box::new(Field::new("item", DataType::Float64, true))),
- b.1.data_type()
- );
- let c = schema.column_with_name("c").unwrap();
- assert_eq!(
- &DataType::List(Box::new(Field::new("item", DataType::Boolean, true))),
- c.1.data_type()
- );
- let d = schema.column_with_name("d").unwrap();
- assert_eq!(
- &DataType::List(Box::new(Field::new("item", DataType::Utf8, true))),
- d.1.data_type()
- );
-
- let bb = batch
- .column(b.0)
- .as_any()
- .downcast_ref::<ListArray>()
- .unwrap();
- let bb = bb.values();
- let bb = bb.as_any().downcast_ref::<Float64Array>().unwrap();
- assert_eq!(10, bb.len());
- assert!(4.0 - bb.value(9) < f64::EPSILON);
-
- let cc = batch
- .column(c.0)
- .as_any()
- .downcast_ref::<ListArray>()
- .unwrap();
- // test that the list offsets are correct
- assert_eq!(
- cc.data().buffers()[0],
- Buffer::from_slice_ref(&[0i32, 2, 2, 4, 5])
- );
- let cc = cc.values();
- let cc = cc.as_any().downcast_ref::<BooleanArray>().unwrap();
- let cc_expected = BooleanArray::from(vec![
- Some(false),
- Some(true),
- Some(false),
- None,
- Some(false),
- ]);
- assert_eq!(cc.data_ref(), cc_expected.data_ref());
-
- let dd: &ListArray = batch
- .column(d.0)
- .as_any()
- .downcast_ref::<ListArray>()
- .unwrap();
- // test that the list offsets are correct
- assert_eq!(
- dd.data().buffers()[0],
- Buffer::from_slice_ref(&[0i32, 1, 1, 2, 6])
- );
- let dd = dd.values();
- let dd = dd.as_any().downcast_ref::<StringArray>().unwrap();
- // values are 6 because a `d: null` is treated as a null slot
- // and a list's null slot can be omitted from the child (i.e. same offset)
- assert_eq!(6, dd.len());
- assert_eq!("text", dd.value(1));
- assert_eq!("1", dd.value(2));
- assert_eq!("false", dd.value(3));
- assert_eq!("array", dd.value(4));
- assert_eq!("2.4", dd.value(5));
- }
- }
-
- #[test]
- fn test_nested_struct_json_arrays() {
- let c_field = Field::new(
- "c",
- DataType::Struct(vec![Field::new("d", DataType::Utf8, true)]),
- true,
- );
- let a_field = Field::new(
- "a",
- DataType::Struct(vec![
- Field::new("b", DataType::Boolean, true),
- c_field.clone(),
- ]),
- true,
- );
- let schema = Arc::new(Schema::new(vec![a_field.clone()]));
- let builder = ReaderBuilder::new().with_schema(schema).with_batch_size(64);
- let mut reader: Reader<File> = builder
- .build::<File>(File::open("test/data/nested_structs.json").unwrap())
- .unwrap();
-
- // build expected output
- let d = StringArray::from(vec![Some("text"), None, Some("text"), None]);
- let c = ArrayDataBuilder::new(c_field.data_type().clone())
- .len(4)
- .add_child_data(d.data().clone())
- .null_bit_buffer(Buffer::from(vec![0b00000101]))
- .build();
- let b = BooleanArray::from(vec![Some(true), Some(false), Some(true), None]);
- let a = ArrayDataBuilder::new(a_field.data_type().clone())
- .len(4)
- .add_child_data(b.data().clone())
- .add_child_data(c)
- .null_bit_buffer(Buffer::from(vec![0b00000111]))
- .build();
- let expected = make_array(a);
-
- // compare `a` with result from json reader
- let batch = reader.next().unwrap().unwrap();
- let read = batch.column(0);
- assert!(
- expected.data_ref() == read.data_ref(),
- "{:?} != {:?}",
- expected.data(),
- read.data(),
- );
- }
-
- #[test]
- fn test_nested_list_json_arrays() {
- let c_field = Field::new(
- "c",
- DataType::Struct(vec![Field::new("d", DataType::Utf8, true)]),
- true,
- );
- let a_struct_field = Field::new(
- "a",
- DataType::Struct(vec![
- Field::new("b", DataType::Boolean, true),
- c_field.clone(),
- ]),
- true,
- );
- let a_field =
- Field::new("a", DataType::List(Box::new(a_struct_field.clone())), true);
- let schema = Arc::new(Schema::new(vec![a_field.clone()]));
- let builder = ReaderBuilder::new().with_schema(schema).with_batch_size(64);
- let json_content = r#"
- {"a": [{"b": true, "c": {"d": "a_text"}}, {"b": false, "c": {"d": "b_text"}}]}
- {"a": [{"b": false, "c": null}]}
- {"a": [{"b": true, "c": {"d": "c_text"}}, {"b": null, "c": {"d": "d_text"}}, {"b": true, "c": {"d": null}}]}
- {"a": null}
- {"a": []}
- "#;
- let mut reader = builder.build(Cursor::new(json_content)).unwrap();
-
- // build expected output
- let d = StringArray::from(vec![
- Some("a_text"),
- Some("b_text"),
- None,
- Some("c_text"),
- Some("d_text"),
- None,
- None,
- ]);
- let c = ArrayDataBuilder::new(c_field.data_type().clone())
- .len(7)
- .add_child_data(d.data().clone())
- .null_bit_buffer(Buffer::from(vec![0b00111011]))
- .build();
- let b = BooleanArray::from(vec![
- Some(true),
- Some(false),
- Some(false),
- Some(true),
- None,
- Some(true),
- None,
- ]);
- let a = ArrayDataBuilder::new(a_struct_field.data_type().clone())
- .len(7)
- .add_child_data(b.data().clone())
- .add_child_data(c.clone())
- .null_bit_buffer(Buffer::from(vec![0b00111111]))
- .build();
- let a_list = ArrayDataBuilder::new(a_field.data_type().clone())
- .len(5)
- .add_buffer(Buffer::from_slice_ref(&[0i32, 2, 3, 6, 6, 6]))
- .add_child_data(a)
- .null_bit_buffer(Buffer::from(vec![0b00010111]))
- .build();
- let expected = make_array(a_list);
-
- // compare `a` with result from json reader
- let batch = reader.next().unwrap().unwrap();
- let read = batch.column(0);
- assert_eq!(read.len(), 5);
- // compare the arrays the long way around, to better detect differences
- let read: &ListArray = read.as_any().downcast_ref::<ListArray>().unwrap();
- let expected = expected.as_any().downcast_ref::<ListArray>().unwrap();
- assert_eq!(
- read.data().buffers()[0],
- Buffer::from_slice_ref(&[0i32, 2, 3, 6, 6, 6])
- );
- // compare list null buffers
- assert_eq!(read.data().null_buffer(), expected.data().null_buffer());
- // build struct from list
- let struct_values = read.values();
- let struct_array: &StructArray = struct_values
- .as_any()
- .downcast_ref::<StructArray>()
- .unwrap();
- let expected_struct_values = expected.values();
- let expected_struct_array = expected_struct_values
- .as_any()
- .downcast_ref::<StructArray>()
- .unwrap();
-
- assert_eq!(7, struct_array.len());
- assert_eq!(1, struct_array.null_count());
- assert_eq!(7, expected_struct_array.len());
- assert_eq!(1, expected_struct_array.null_count());
- // test struct's nulls
- assert_eq!(
- struct_array.data().null_buffer(),
- expected_struct_array.data().null_buffer()
- );
- // test struct's fields
- let read_b = struct_array.column(0);
- assert_eq!(b.data_ref(), read_b.data_ref());
- let read_c = struct_array.column(1);
- assert_eq!(&c, read_c.data_ref());
- let read_c: &StructArray = read_c.as_any().downcast_ref::<StructArray>().unwrap();
- let read_d = read_c.column(0);
- assert_eq!(d.data_ref(), read_d.data_ref());
-
- assert_eq!(read.data_ref(), expected.data_ref());
- }
-
- #[test]
- fn test_dictionary_from_json_basic_with_nulls() {
- let schema = Schema::new(vec![Field::new(
- "d",
- Dictionary(Box::new(DataType::Int16), Box::new(DataType::Utf8)),
- true,
- )]);
- let builder = ReaderBuilder::new()
- .with_schema(Arc::new(schema))
- .with_batch_size(64);
- let mut reader: Reader<File> = builder
- .build::<File>(File::open("test/data/basic_nulls.json").unwrap())
- .unwrap();
- let batch = reader.next().unwrap().unwrap();
-
- assert_eq!(1, batch.num_columns());
- assert_eq!(12, batch.num_rows());
-
- let schema = reader.schema();
- let batch_schema = batch.schema();
- assert_eq!(schema, batch_schema);
-
- let d = schema.column_with_name("d").unwrap();
- assert_eq!(
- &Dictionary(Box::new(DataType::Int16), Box::new(DataType::Utf8)),
- d.1.data_type()
- );
-
- let dd = batch
- .column(d.0)
- .as_any()
- .downcast_ref::<DictionaryArray<Int16Type>>()
- .unwrap();
- assert_eq!(false, dd.is_valid(0));
- assert_eq!(true, dd.is_valid(1));
- assert_eq!(true, dd.is_valid(2));
- assert_eq!(false, dd.is_valid(11));
-
- assert_eq!(
- dd.keys(),
- &Int16Array::from(vec![
- None,
- Some(0),
- Some(1),
- Some(0),
- None,
- None,
- Some(0),
- None,
- Some(1),
- Some(0),
- Some(0),
- None
- ])
- );
- }
-
- #[test]
- fn test_dictionary_from_json_int8() {
- let schema = Schema::new(vec![Field::new(
- "d",
- Dictionary(Box::new(DataType::Int8), Box::new(DataType::Utf8)),
- true,
- )]);
- let builder = ReaderBuilder::new()
- .with_schema(Arc::new(schema))
- .with_batch_size(64);
- let mut reader: Reader<File> = builder
- .build::<File>(File::open("test/data/basic_nulls.json").unwrap())
- .unwrap();
- let batch = reader.next().unwrap().unwrap();
-
- assert_eq!(1, batch.num_columns());
- assert_eq!(12, batch.num_rows());
-
- let schema = reader.schema();
- let batch_schema = batch.schema();
- assert_eq!(schema, batch_schema);
-
- let d = schema.column_with_name("d").unwrap();
- assert_eq!(
- &Dictionary(Box::new(DataType::Int8), Box::new(DataType::Utf8)),
- d.1.data_type()
- );
- }
-
- #[test]
- fn test_dictionary_from_json_int32() {
- let schema = Schema::new(vec![Field::new(
- "d",
- Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
- true,
- )]);
- let builder = ReaderBuilder::new()
- .with_schema(Arc::new(schema))
- .with_batch_size(64);
- let mut reader: Reader<File> = builder
- .build::<File>(File::open("test/data/basic_nulls.json").unwrap())
- .unwrap();
- let batch = reader.next().unwrap().unwrap();
-
- assert_eq!(1, batch.num_columns());
- assert_eq!(12, batch.num_rows());
-
- let schema = reader.schema();
- let batch_schema = batch.schema();
- assert_eq!(schema, batch_schema);
-
- let d = schema.column_with_name("d").unwrap();
- assert_eq!(
- &Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
- d.1.data_type()
- );
- }
-
- #[test]
- fn test_dictionary_from_json_int64() {
- let schema = Schema::new(vec![Field::new(
- "d",
- Dictionary(Box::new(DataType::Int64), Box::new(DataType::Utf8)),
- true,
- )]);
- let builder = ReaderBuilder::new()
- .with_schema(Arc::new(schema))
- .with_batch_size(64);
- let mut reader: Reader<File> = builder
- .build::<File>(File::open("test/data/basic_nulls.json").unwrap())
- .unwrap();
- let batch = reader.next().unwrap().unwrap();
-
- assert_eq!(1, batch.num_columns());
- assert_eq!(12, batch.num_rows());
-
- let schema = reader.schema();
- let batch_schema = batch.schema();
- assert_eq!(schema, batch_schema);
-
- let d = schema.column_with_name("d").unwrap();
- assert_eq!(
- &Dictionary(Box::new(DataType::Int64), Box::new(DataType::Utf8)),
- d.1.data_type()
- );
- }
-
- #[test]
- fn test_skip_empty_lines() {
- let builder = ReaderBuilder::new().infer_schema(None).with_batch_size(64);
- let json_content = "
- {\"a\": 1}
-
- {\"a\": 2}
-
- {\"a\": 3}";
- let mut reader = builder.build(Cursor::new(json_content)).unwrap();
- let batch = reader.next().unwrap().unwrap();
-
- assert_eq!(1, batch.num_columns());
- assert_eq!(3, batch.num_rows());
-
- let schema = reader.schema();
- let c = schema.column_with_name("a").unwrap();
- assert_eq!(&DataType::Int64, c.1.data_type());
- }
-
- #[test]
- fn test_row_type_validation() {
- let builder = ReaderBuilder::new().infer_schema(None).with_batch_size(64);
- let json_content = "
- [1, \"hello\"]
- \"world\"";
- let re = builder.build(Cursor::new(json_content));
- assert_eq!(
- re.err().unwrap().to_string(),
- r#"Json error: Expected JSON record to be an object, found Array([Number(1), String("hello")])"#,
- );
- }
-
- #[test]
- fn test_list_of_string_dictionary_from_json() {
- let schema = Schema::new(vec![Field::new(
- "events",
- List(Box::new(Field::new(
- "item",
- Dictionary(Box::new(DataType::UInt64), Box::new(DataType::Utf8)),
- true,
- ))),
- true,
- )]);
- let builder = ReaderBuilder::new()
- .with_schema(Arc::new(schema))
- .with_batch_size(64);
- let mut reader: Reader<File> = builder
- .build::<File>(File::open("test/data/list_string_dict_nested.json").unwrap())
- .unwrap();
- let batch = reader.next().unwrap().unwrap();
-
- assert_eq!(1, batch.num_columns());
- assert_eq!(3, batch.num_rows());
-
- let schema = reader.schema();
- let batch_schema = batch.schema();
- assert_eq!(schema, batch_schema);
-
- let events = schema.column_with_name("events").unwrap();
- assert_eq!(
- &List(Box::new(Field::new(
- "item",
- Dictionary(Box::new(DataType::UInt64), Box::new(DataType::Utf8)),
- true
- ))),
- events.1.data_type()
- );
-
- let evs_list = batch
- .column(events.0)
- .as_any()
- .downcast_ref::<ListArray>()
- .unwrap();
- let evs_list = evs_list.values();
- let evs_list = evs_list
- .as_any()
- .downcast_ref::<DictionaryArray<UInt64Type>>()
- .unwrap();
- assert_eq!(6, evs_list.len());
- assert_eq!(true, evs_list.is_valid(1));
- assert_eq!(DataType::Utf8, evs_list.value_type());
-
- // dict from the events list
- let dict_el = evs_list.values();
- let dict_el = dict_el.as_any().downcast_ref::<StringArray>().unwrap();
- assert_eq!(3, dict_el.len());
- assert_eq!("Elect Leader", dict_el.value(0));
- assert_eq!("Do Ballot", dict_el.value(1));
- assert_eq!("Send Data", dict_el.value(2));
- }
-
- #[test]
- fn test_list_of_string_dictionary_from_json_with_nulls() {
- let schema = Schema::new(vec![Field::new(
- "events",
- List(Box::new(Field::new(
- "item",
- Dictionary(Box::new(DataType::UInt64), Box::new(DataType::Utf8)),
- true,
- ))),
- true,
- )]);
- let builder = ReaderBuilder::new()
- .with_schema(Arc::new(schema))
- .with_batch_size(64);
- let mut reader: Reader<File> = builder
- .build::<File>(
- File::open("test/data/list_string_dict_nested_nulls.json").unwrap(),
- )
- .unwrap();
- let batch = reader.next().unwrap().unwrap();
-
- assert_eq!(1, batch.num_columns());
- assert_eq!(3, batch.num_rows());
-
- let schema = reader.schema();
- let batch_schema = batch.schema();
- assert_eq!(schema, batch_schema);
-
- let events = schema.column_with_name("events").unwrap();
- assert_eq!(
- &List(Box::new(Field::new(
- "item",
- Dictionary(Box::new(DataType::UInt64), Box::new(DataType::Utf8)),
- true
- ))),
- events.1.data_type()
- );
-
- let evs_list = batch
- .column(events.0)
- .as_any()
- .downcast_ref::<ListArray>()
- .unwrap();
- let evs_list = evs_list.values();
- let evs_list = evs_list
- .as_any()
- .downcast_ref::<DictionaryArray<UInt64Type>>()
- .unwrap();
- assert_eq!(8, evs_list.len());
- assert_eq!(true, evs_list.is_valid(1));
- assert_eq!(DataType::Utf8, evs_list.value_type());
-
- // dict from the events list
- let dict_el = evs_list.values();
- let dict_el = dict_el.as_any().downcast_ref::<StringArray>().unwrap();
- assert_eq!(2, evs_list.null_count());
- assert_eq!(3, dict_el.len());
- assert_eq!("Elect Leader", dict_el.value(0));
- assert_eq!("Do Ballot", dict_el.value(1));
- assert_eq!("Send Data", dict_el.value(2));
- }
-
- #[test]
- fn test_dictionary_from_json_uint8() {
- let schema = Schema::new(vec![Field::new(
- "d",
- Dictionary(Box::new(DataType::UInt8), Box::new(DataType::Utf8)),
- true,
- )]);
- let builder = ReaderBuilder::new()
- .with_schema(Arc::new(schema))
- .with_batch_size(64);
- let mut reader: Reader<File> = builder
- .build::<File>(File::open("test/data/basic_nulls.json").unwrap())
- .unwrap();
- let batch = reader.next().unwrap().unwrap();
-
- assert_eq!(1, batch.num_columns());
- assert_eq!(12, batch.num_rows());
-
- let schema = reader.schema();
- let batch_schema = batch.schema();
- assert_eq!(schema, batch_schema);
-
- let d = schema.column_with_name("d").unwrap();
- assert_eq!(
- &Dictionary(Box::new(DataType::UInt8), Box::new(DataType::Utf8)),
- d.1.data_type()
- );
- }
-
- #[test]
- fn test_dictionary_from_json_uint32() {
- let schema = Schema::new(vec![Field::new(
- "d",
- Dictionary(Box::new(DataType::UInt32), Box::new(DataType::Utf8)),
- true,
- )]);
- let builder = ReaderBuilder::new()
- .with_schema(Arc::new(schema))
- .with_batch_size(64);
- let mut reader: Reader<File> = builder
- .build::<File>(File::open("test/data/basic_nulls.json").unwrap())
- .unwrap();
- let batch = reader.next().unwrap().unwrap();
-
- assert_eq!(1, batch.num_columns());
- assert_eq!(12, batch.num_rows());
-
- let schema = reader.schema();
- let batch_schema = batch.schema();
- assert_eq!(schema, batch_schema);
-
- let d = schema.column_with_name("d").unwrap();
- assert_eq!(
- &Dictionary(Box::new(DataType::UInt32), Box::new(DataType::Utf8)),
- d.1.data_type()
- );
- }
-
- #[test]
- fn test_dictionary_from_json_uint64() {
- let schema = Schema::new(vec![Field::new(
- "d",
- Dictionary(Box::new(DataType::UInt64), Box::new(DataType::Utf8)),
- true,
- )]);
- let builder = ReaderBuilder::new()
- .with_schema(Arc::new(schema))
- .with_batch_size(64);
- let mut reader: Reader<File> = builder
- .build::<File>(File::open("test/data/basic_nulls.json").unwrap())
- .unwrap();
- let batch = reader.next().unwrap().unwrap();
-
- assert_eq!(1, batch.num_columns());
- assert_eq!(12, batch.num_rows());
-
- let schema = reader.schema();
- let batch_schema = batch.schema();
- assert_eq!(schema, batch_schema);
-
- let d = schema.column_with_name("d").unwrap();
- assert_eq!(
- &Dictionary(Box::new(DataType::UInt64), Box::new(DataType::Utf8)),
- d.1.data_type()
- );
- }
-
- #[test]
- fn test_with_multiple_batches() {
- let builder = ReaderBuilder::new()
- .infer_schema(Some(4))
- .with_batch_size(5);
- let mut reader: Reader<File> = builder
- .build::<File>(File::open("test/data/basic_nulls.json").unwrap())
- .unwrap();
-
- let mut num_records = Vec::new();
- while let Some(rb) = reader.next().unwrap() {
- num_records.push(rb.num_rows());
- }
-
- assert_eq!(vec![5, 5, 2], num_records);
- }
-
- #[test]
- fn test_json_infer_schema() {
- let schema = Schema::new(vec![
- Field::new("a", DataType::Int64, true),
- Field::new(
- "b",
- DataType::List(Box::new(Field::new("item", DataType::Float64, true))),
- true,
- ),
- Field::new(
- "c",
- DataType::List(Box::new(Field::new("item", DataType::Boolean, true))),
- true,
- ),
- Field::new(
- "d",
- DataType::List(Box::new(Field::new("item", DataType::Utf8, true))),
- true,
- ),
- ]);
-
- let mut reader =
- BufReader::new(File::open("test/data/mixed_arrays.json").unwrap());
- let inferred_schema = infer_json_schema_from_seekable(&mut reader, None).unwrap();
-
- assert_eq!(inferred_schema, schema);
-
- let file = File::open("test/data/mixed_arrays.json.gz").unwrap();
- let mut reader = BufReader::new(GzDecoder::new(&file));
- let inferred_schema = infer_json_schema(&mut reader, None).unwrap();
-
- assert_eq!(inferred_schema, schema);
- }
-
- #[test]
- fn test_json_infer_schema_nested_structs() {
- let schema = Schema::new(vec![
- Field::new(
- "c1",
- DataType::Struct(vec![
- Field::new("a", DataType::Boolean, true),
- Field::new(
- "b",
- DataType::Struct(vec![Field::new("c", DataType::Utf8, true)]),
- true,
- ),
- ]),
- true,
- ),
- Field::new("c2", DataType::Int64, true),
- Field::new("c3", DataType::Utf8, true),
- ]);
-
- let inferred_schema = infer_json_schema_from_iterator(
- vec![
- Ok(serde_json::json!({"c1": {"a": true, "b": {"c": "text"}}, "c2": 1})),
- Ok(serde_json::json!({"c1": {"a": false, "b": null}, "c2": 0})),
- Ok(serde_json::json!({"c1": {"a": true, "b": {"c": "text"}}, "c3": "ok"})),
- ]
- .into_iter(),
- )
- .unwrap();
-
- assert_eq!(inferred_schema, schema);
- }
-
- #[test]
- fn test_json_infer_schema_struct_in_list() {
- let schema = Schema::new(vec![
- Field::new(
- "c1",
- DataType::List(Box::new(Field::new(
- "item",
- DataType::Struct(vec![
- Field::new("a", DataType::Utf8, true),
- Field::new("b", DataType::Int64, true),
- Field::new("c", DataType::Boolean, true),
- ]),
- true,
- ))),
- true,
- ),
- Field::new("c2", DataType::Float64, true),
- Field::new(
- "c3",
- // empty json array's inner types are inferred as null
- DataType::List(Box::new(Field::new("item", DataType::Null, true))),
- true,
- ),
- ]);
-
- let inferred_schema = infer_json_schema_from_iterator(
- vec![
- Ok(serde_json::json!({
- "c1": [{"a": "foo", "b": 100}], "c2": 1, "c3": [],
- })),
- Ok(serde_json::json!({
- "c1": [{"a": "bar", "b": 2}, {"a": "foo", "c": true}], "c2": 0, "c3": [],
- })),
- Ok(serde_json::json!({"c1": [], "c2": 0.5, "c3": []})),
- ]
- .into_iter(),
- )
- .unwrap();
-
- assert_eq!(inferred_schema, schema);
- }
-
- #[test]
- fn test_json_infer_schema_nested_list() {
- let schema = Schema::new(vec![
- Field::new(
- "c1",
- DataType::List(Box::new(Field::new(
- "item",
- DataType::List(Box::new(Field::new("item", DataType::Utf8, true))),
- true,
- ))),
- true,
- ),
- Field::new("c2", DataType::Float64, true),
- ]);
-
- let inferred_schema = infer_json_schema_from_iterator(
- vec![
- Ok(serde_json::json!({
- "c1": [],
- "c2": 12,
- })),
- Ok(serde_json::json!({
- "c1": [["a", "b"], ["c"]],
- })),
- Ok(serde_json::json!({
- "c1": [["foo"]],
- "c2": 0.11,
- })),
- ]
- .into_iter(),
- )
- .unwrap();
-
- assert_eq!(inferred_schema, schema);
- }
-
- #[test]
- fn test_timestamp_from_json_seconds() {
- let schema = Schema::new(vec![Field::new(
- "a",
- DataType::Timestamp(TimeUnit::Second, None),
- true,
- )]);
-
- let builder = ReaderBuilder::new()
- .with_schema(Arc::new(schema))
- .with_batch_size(64);
- let mut reader: Reader<File> = builder
- .build::<File>(File::open("test/data/basic_nulls.json").unwrap())
- .unwrap();
- let batch = reader.next().unwrap().unwrap();
-
- assert_eq!(1, batch.num_columns());
- assert_eq!(12, batch.num_rows());
-
- let schema = reader.schema();
- let batch_schema = batch.schema();
- assert_eq!(schema, batch_schema);
-
- let a = schema.column_with_name("a").unwrap();
- assert_eq!(
- &DataType::Timestamp(TimeUnit::Second, None),
- a.1.data_type()
- );
-
- let aa = batch
- .column(a.0)
- .as_any()
- .downcast_ref::<TimestampSecondArray>()
- .unwrap();
- assert_eq!(true, aa.is_valid(0));
- assert_eq!(false, aa.is_valid(1));
- assert_eq!(false, aa.is_valid(2));
- assert_eq!(1, aa.value(0));
- assert_eq!(1, aa.value(3));
- assert_eq!(5, aa.value(7));
- }
-
- #[test]
- fn test_timestamp_from_json_milliseconds() {
- let schema = Schema::new(vec![Field::new(
- "a",
- DataType::Timestamp(TimeUnit::Millisecond, None),
- true,
- )]);
-
- let builder = ReaderBuilder::new()
- .with_schema(Arc::new(schema))
- .with_batch_size(64);
- let mut reader: Reader<File> = builder
- .build::<File>(File::open("test/data/basic_nulls.json").unwrap())
- .unwrap();
- let batch = reader.next().unwrap().unwrap();
-
- assert_eq!(1, batch.num_columns());
- assert_eq!(12, batch.num_rows());
-
- let schema = reader.schema();
- let batch_schema = batch.schema();
- assert_eq!(schema, batch_schema);
-
- let a = schema.column_with_name("a").unwrap();
- assert_eq!(
- &DataType::Timestamp(TimeUnit::Millisecond, None),
- a.1.data_type()
- );
-
- let aa = batch
- .column(a.0)
- .as_any()
- .downcast_ref::<TimestampMillisecondArray>()
- .unwrap();
- assert_eq!(true, aa.is_valid(0));
- assert_eq!(false, aa.is_valid(1));
- assert_eq!(false, aa.is_valid(2));
- assert_eq!(1, aa.value(0));
- assert_eq!(1, aa.value(3));
- assert_eq!(5, aa.value(7));
- }
-
- #[test]
- fn test_date_from_json_milliseconds() {
- let schema = Schema::new(vec![Field::new("a", DataType::Date64, true)]);
-
- let builder = ReaderBuilder::new()
- .with_schema(Arc::new(schema))
- .with_batch_size(64);
- let mut reader: Reader<File> = builder
- .build::<File>(File::open("test/data/basic_nulls.json").unwrap())
- .unwrap();
- let batch = reader.next().unwrap().unwrap();
-
- assert_eq!(1, batch.num_columns());
- assert_eq!(12, batch.num_rows());
-
- let schema = reader.schema();
- let batch_schema = batch.schema();
- assert_eq!(schema, batch_schema);
-
- let a = schema.column_with_name("a").unwrap();
- assert_eq!(&DataType::Date64, a.1.data_type());
-
- let aa = batch
- .column(a.0)
- .as_any()
- .downcast_ref::<Date64Array>()
- .unwrap();
- assert_eq!(true, aa.is_valid(0));
- assert_eq!(false, aa.is_valid(1));
- assert_eq!(false, aa.is_valid(2));
- assert_eq!(1, aa.value(0));
- assert_eq!(1, aa.value(3));
- assert_eq!(5, aa.value(7));
- }
-
- #[test]
- fn test_time_from_json_nanoseconds() {
- let schema = Schema::new(vec![Field::new(
- "a",
- DataType::Time64(TimeUnit::Nanosecond),
- true,
- )]);
-
- let builder = ReaderBuilder::new()
- .with_schema(Arc::new(schema))
- .with_batch_size(64);
- let mut reader: Reader<File> = builder
- .build::<File>(File::open("test/data/basic_nulls.json").unwrap())
- .unwrap();
- let batch = reader.next().unwrap().unwrap();
-
- assert_eq!(1, batch.num_columns());
- assert_eq!(12, batch.num_rows());
-
- let schema = reader.schema();
- let batch_schema = batch.schema();
- assert_eq!(schema, batch_schema);
-
- let a = schema.column_with_name("a").unwrap();
- assert_eq!(&DataType::Time64(TimeUnit::Nanosecond), a.1.data_type());
-
- let aa = batch
- .column(a.0)
- .as_any()
- .downcast_ref::<Time64NanosecondArray>()
- .unwrap();
- assert_eq!(true, aa.is_valid(0));
- assert_eq!(false, aa.is_valid(1));
- assert_eq!(false, aa.is_valid(2));
- assert_eq!(1, aa.value(0));
- assert_eq!(1, aa.value(3));
- assert_eq!(5, aa.value(7));
- }
-
- #[test]
- fn test_json_read_nested_list() {
- let schema = Schema::new(vec![Field::new(
- "c1",
- DataType::List(Box::new(Field::new(
- "item",
- DataType::List(Box::new(Field::new("item", DataType::Utf8, true))),
- true,
- ))),
- true,
- )]);
-
- let decoder = Decoder::new(Arc::new(schema), 1024, None);
- let batch = decoder
- .next_batch(
- &mut vec![
- Ok(serde_json::json!({
- "c1": [],
- })),
- Ok(serde_json::json!({
- "c1": [["a", "b"], ["c"], ["e", "f"], ["g"], ["h"], ["i"], ["j"], ["k"]],
- })),
- Ok(serde_json::json!({
- "c1": [["foo"], ["bar"]],
- })),
- ]
- .into_iter(),
- )
- .unwrap()
- .unwrap();
-
- assert_eq!(batch.num_columns(), 1);
- assert_eq!(batch.num_rows(), 3);
- }
-
- #[test]
- fn test_json_read_list_of_structs() {
- let schema = Schema::new(vec![Field::new(
- "c1",
- DataType::List(Box::new(Field::new(
- "item",
- DataType::Struct(vec![Field::new("a", DataType::Int64, true)]),
- true,
- ))),
- true,
- )]);
-
- let decoder = Decoder::new(Arc::new(schema), 1024, None);
- let batch = decoder
- .next_batch(
- // NOTE: total struct element count needs to be greater than
- // bit_util::ceil(array_count, 8) to test validity bit buffer length calculation
- // logic
- &mut vec![
- Ok(serde_json::json!({
- "c1": [{"a": 1}],
- })),
- Ok(serde_json::json!({
- "c1": [{"a": 2}, {"a": 3}, {"a": 4}, {"a": 5}, {"a": 6}, {"a": 7}],
- })),
- Ok(serde_json::json!({
- "c1": [{"a": 10}, {"a": 11}],
- })),
- ]
- .into_iter(),
- )
- .unwrap()
- .unwrap();
-
- assert_eq!(batch.num_columns(), 1);
- assert_eq!(batch.num_rows(), 3);
- }
-
- #[test]
- fn test_json_iterator() {
- let builder = ReaderBuilder::new().infer_schema(None).with_batch_size(5);
- let reader: Reader<File> = builder
- .build::<File>(File::open("test/data/basic.json").unwrap())
- .unwrap();
- let schema = reader.schema();
- let (col_a_index, _) = schema.column_with_name("a").unwrap();
-
- let mut sum_num_rows = 0;
- let mut num_batches = 0;
- let mut sum_a = 0;
- for batch in reader {
- let batch = batch.unwrap();
- assert_eq!(4, batch.num_columns());
- sum_num_rows += batch.num_rows();
- num_batches += 1;
- let batch_schema = batch.schema();
- assert_eq!(schema, batch_schema);
- let a_array = batch
- .column(col_a_index)
- .as_any()
- .downcast_ref::<Int64Array>()
- .unwrap();
- sum_a += (0..a_array.len()).map(|i| a_array.value(i)).sum::<i64>();
- }
- assert_eq!(12, sum_num_rows);
- assert_eq!(3, num_batches);
- assert_eq!(100000000000011, sum_a);
- }
-}
diff --git a/arrow/src/json/writer.rs b/arrow/src/json/writer.rs
deleted file mode 100644
index 8587c1d..0000000
--- a/arrow/src/json/writer.rs
+++ /dev/null
@@ -1,1262 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-//! # JSON Writer
-//!
-//! This JSON writer converts Arrow [`RecordBatch`]es into arrays of
-//! JSON objects or JSON formatted byte streams.
-//!
-//! ## Writing JSON Objects
-//!
-//! To serialize [`RecordBatch`]es into array of
-//! [JSON](https://docs.serde.rs/serde_json/) objects, use
-//! [`record_batches_to_json_rows`]:
-//!
-//! ```
-//! use std::sync::Arc;
-//!
-//! use arrow::array::Int32Array;
-//! use arrow::datatypes::{DataType, Field, Schema};
-//! use arrow::json;
-//! use arrow::record_batch::RecordBatch;
-//!
-//! let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]);
-//! let a = Int32Array::from(vec![1, 2, 3]);
-//! let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a)]).unwrap();
-//!
-//! let json_rows = json::writer::record_batches_to_json_rows(&[batch]);
-//! assert_eq!(
-//! serde_json::Value::Object(json_rows[1].clone()),
-//! serde_json::json!({"a": 2}),
-//! );
-//! ```
-//!
-//! ## Writing JSON formatted byte streams
-//!
-//! To serialize [`RecordBatch`]es into line-delimited JSON bytes, use
-//! [`LineDelimitedWriter`]:
-//!
-//! ```
-//! use std::sync::Arc;
-//!
-//! use arrow::array::Int32Array;
-//! use arrow::datatypes::{DataType, Field, Schema};
-//! use arrow::json;
-//! use arrow::record_batch::RecordBatch;
-//!
-//! let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]);
-//! let a = Int32Array::from(vec![1, 2, 3]);
-//! let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a)]).unwrap();
-//!
-//! // Write the record batch out as JSON
-//! let buf = Vec::new();
-//! let mut writer = json::LineDelimitedWriter::new(buf);
-//! writer.write_batches(&vec![batch]).unwrap();
-//! writer.finish().unwrap();
-//!
-//! // Get the underlying buffer back,
-//! let buf = writer.into_inner();
-//! assert_eq!(r#"{"a":1}
-//! {"a":2}
-//! {"a":3}
-//!"#, String::from_utf8(buf).unwrap())
-//! ```
-//!
-//! To serialize [`RecordBatch`]es into a well formed JSON array, use
-//! [`ArrayWriter`]:
-//!
-//! ```
-//! use std::sync::Arc;
-//!
-//! use arrow::array::Int32Array;
-//! use arrow::datatypes::{DataType, Field, Schema};
-//! use arrow::json;
-//! use arrow::record_batch::RecordBatch;
-//!
-//! let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]);
-//! let a = Int32Array::from(vec![1, 2, 3]);
-//! let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a)]).unwrap();
-//!
-//! // Write the record batch out as a JSON array
-//! let buf = Vec::new();
-//! let mut writer = json::ArrayWriter::new(buf);
-//! writer.write_batches(&vec![batch]).unwrap();
-//! writer.finish().unwrap();
-//!
-//! // Get the underlying buffer back,
-//! let buf = writer.into_inner();
-//! assert_eq!(r#"[{"a":1},{"a":2},{"a":3}]"#, String::from_utf8(buf).unwrap())
-//! ```
-
-use std::iter;
-use std::{fmt::Debug, io::Write};
-
-use serde_json::map::Map as JsonMap;
-use serde_json::Value;
-
-use crate::array::*;
-use crate::datatypes::*;
-use crate::error::Result;
-use crate::record_batch::RecordBatch;
-
-fn primitive_array_to_json<T: ArrowPrimitiveType>(array: &ArrayRef) -> Vec<Value> {
- as_primitive_array::<T>(array)
- .iter()
- .map(|maybe_value| match maybe_value {
- Some(v) => v.into_json_value().unwrap_or(Value::Null),
- None => Value::Null,
- })
- .collect()
-}
-
-fn struct_array_to_jsonmap_array(
- array: &StructArray,
- row_count: usize,
-) -> Vec<JsonMap<String, Value>> {
- let inner_col_names = array.column_names();
-
- let mut inner_objs = iter::repeat(JsonMap::new())
- .take(row_count)
- .collect::<Vec<JsonMap<String, Value>>>();
-
- array
- .columns()
- .iter()
- .enumerate()
- .for_each(|(j, struct_col)| {
- set_column_for_json_rows(
- &mut inner_objs,
- row_count,
- struct_col,
- inner_col_names[j],
- );
- });
-
- inner_objs
-}
-
-/// Converts an arrow [`ArrayRef`] into a `Vec` of Serde JSON [`serde_json::Value`]'s
-pub fn array_to_json_array(array: &ArrayRef) -> Vec<Value> {
- match array.data_type() {
- DataType::Null => iter::repeat(Value::Null).take(array.len()).collect(),
- DataType::Boolean => as_boolean_array(array)
- .iter()
- .map(|maybe_value| match maybe_value {
- Some(v) => v.into(),
- None => Value::Null,
- })
- .collect(),
-
- DataType::Utf8 => as_string_array(array)
- .iter()
- .map(|maybe_value| match maybe_value {
- Some(v) => v.into(),
- None => Value::Null,
- })
- .collect(),
- DataType::Int8 => primitive_array_to_json::<Int8Type>(array),
- DataType::Int16 => primitive_array_to_json::<Int16Type>(array),
- DataType::Int32 => primitive_array_to_json::<Int32Type>(array),
- DataType::Int64 => primitive_array_to_json::<Int64Type>(array),
- DataType::UInt8 => primitive_array_to_json::<UInt8Type>(array),
- DataType::UInt16 => primitive_array_to_json::<UInt16Type>(array),
- DataType::UInt32 => primitive_array_to_json::<UInt32Type>(array),
- DataType::UInt64 => primitive_array_to_json::<UInt64Type>(array),
- DataType::Float32 => primitive_array_to_json::<Float32Type>(array),
- DataType::Float64 => primitive_array_to_json::<Float64Type>(array),
- DataType::List(_) => as_list_array(array)
- .iter()
- .map(|maybe_value| match maybe_value {
- Some(v) => Value::Array(array_to_json_array(&v)),
- None => Value::Null,
- })
- .collect(),
- DataType::LargeList(_) => as_large_list_array(array)
- .iter()
- .map(|maybe_value| match maybe_value {
- Some(v) => Value::Array(array_to_json_array(&v)),
- None => Value::Null,
- })
- .collect(),
- DataType::Struct(_) => {
- let jsonmaps =
- struct_array_to_jsonmap_array(as_struct_array(array), array.len());
- jsonmaps.into_iter().map(Value::Object).collect()
- }
- _ => {
- panic!(
- "Unsupported datatype for array conversion: {:#?}",
- array.data_type()
- );
- }
- }
-}
-
-macro_rules! set_column_by_array_type {
- ($cast_fn:ident, $col_name:ident, $rows:ident, $array:ident, $row_count:ident) => {
- let arr = $cast_fn($array);
- $rows.iter_mut().zip(arr.iter()).take($row_count).for_each(
- |(row, maybe_value)| {
- if let Some(v) = maybe_value {
- row.insert($col_name.to_string(), v.into());
- }
- },
- );
- };
-}
-
-macro_rules! set_temporal_column_by_array_type {
- ($array_type:ident, $col_name:ident, $rows:ident, $array:ident, $row_count:ident, $cast_fn:ident) => {
- let arr = $array.as_any().downcast_ref::<$array_type>().unwrap();
-
- $rows
- .iter_mut()
- .enumerate()
- .take($row_count)
- .for_each(|(i, row)| {
- if !arr.is_null(i) {
- if let Some(v) = arr.$cast_fn(i) {
- row.insert($col_name.to_string(), v.to_string().into());
- }
- }
- });
- };
-}
-
-fn set_column_by_primitive_type<T: ArrowPrimitiveType>(
- rows: &mut [JsonMap<String, Value>],
- row_count: usize,
- array: &ArrayRef,
- col_name: &str,
-) {
- let primitive_arr = as_primitive_array::<T>(array);
-
- rows.iter_mut()
- .zip(primitive_arr.iter())
- .take(row_count)
- .for_each(|(row, maybe_value)| {
- // when value is null, we simply skip setting the key
- if let Some(j) = maybe_value.and_then(|v| v.into_json_value()) {
- row.insert(col_name.to_string(), j);
- }
- });
-}
-
-fn set_column_for_json_rows(
- rows: &mut [JsonMap<String, Value>],
- row_count: usize,
- array: &ArrayRef,
- col_name: &str,
-) {
- match array.data_type() {
- DataType::Int8 => {
- set_column_by_primitive_type::<Int8Type>(rows, row_count, array, col_name)
- }
- DataType::Int16 => {
- set_column_by_primitive_type::<Int16Type>(rows, row_count, array, col_name)
- }
- DataType::Int32 => {
- set_column_by_primitive_type::<Int32Type>(rows, row_count, array, col_name)
- }
- DataType::Int64 => {
- set_column_by_primitive_type::<Int64Type>(rows, row_count, array, col_name)
- }
- DataType::UInt8 => {
- set_column_by_primitive_type::<UInt8Type>(rows, row_count, array, col_name)
- }
- DataType::UInt16 => {
- set_column_by_primitive_type::<UInt16Type>(rows, row_count, array, col_name)
- }
- DataType::UInt32 => {
- set_column_by_primitive_type::<UInt32Type>(rows, row_count, array, col_name)
- }
- DataType::UInt64 => {
- set_column_by_primitive_type::<UInt64Type>(rows, row_count, array, col_name)
- }
- DataType::Float32 => {
- set_column_by_primitive_type::<Float32Type>(rows, row_count, array, col_name)
- }
- DataType::Float64 => {
- set_column_by_primitive_type::<Float64Type>(rows, row_count, array, col_name)
- }
- DataType::Null => {
- // when value is null, we simply skip setting the key
- }
- DataType::Boolean => {
- set_column_by_array_type!(as_boolean_array, col_name, rows, array, row_count);
- }
- DataType::Utf8 => {
- set_column_by_array_type!(as_string_array, col_name, rows, array, row_count);
- }
- DataType::Date32 => {
- set_temporal_column_by_array_type!(
- Date32Array,
- col_name,
- rows,
- array,
- row_count,
- value_as_date
- );
- }
- DataType::Date64 => {
- set_temporal_column_by_array_type!(
- Date64Array,
- col_name,
- rows,
- array,
- row_count,
- value_as_date
- );
- }
- DataType::Timestamp(TimeUnit::Second, _) => {
- set_temporal_column_by_array_type!(
- TimestampSecondArray,
- col_name,
- rows,
- array,
- row_count,
- value_as_datetime
- );
- }
- DataType::Timestamp(TimeUnit::Millisecond, _) => {
- set_temporal_column_by_array_type!(
- TimestampMillisecondArray,
- col_name,
- rows,
- array,
- row_count,
- value_as_datetime
- );
- }
- DataType::Timestamp(TimeUnit::Microsecond, _) => {
- set_temporal_column_by_array_type!(
- TimestampMicrosecondArray,
- col_name,
- rows,
- array,
- row_count,
- value_as_datetime
- );
- }
- DataType::Timestamp(TimeUnit::Nanosecond, _) => {
- set_temporal_column_by_array_type!(
- TimestampNanosecondArray,
- col_name,
- rows,
- array,
- row_count,
- value_as_datetime
- );
- }
- DataType::Time32(TimeUnit::Second) => {
- set_temporal_column_by_array_type!(
- Time32SecondArray,
- col_name,
- rows,
- array,
- row_count,
- value_as_time
- );
- }
- DataType::Time32(TimeUnit::Millisecond) => {
- set_temporal_column_by_array_type!(
- Time32MillisecondArray,
- col_name,
- rows,
- array,
- row_count,
- value_as_time
- );
- }
- DataType::Time64(TimeUnit::Microsecond) => {
- set_temporal_column_by_array_type!(
- Time64MicrosecondArray,
- col_name,
- rows,
- array,
- row_count,
- value_as_time
- );
- }
- DataType::Time64(TimeUnit::Nanosecond) => {
- set_temporal_column_by_array_type!(
- Time64NanosecondArray,
- col_name,
- rows,
- array,
- row_count,
- value_as_time
- );
- }
- DataType::Duration(TimeUnit::Second) => {
- set_temporal_column_by_array_type!(
- DurationSecondArray,
- col_name,
- rows,
- array,
- row_count,
- value_as_duration
- );
- }
- DataType::Duration(TimeUnit::Millisecond) => {
- set_temporal_column_by_array_type!(
- DurationMillisecondArray,
- col_name,
- rows,
- array,
- row_count,
- value_as_duration
- );
- }
- DataType::Duration(TimeUnit::Microsecond) => {
- set_temporal_column_by_array_type!(
- DurationMicrosecondArray,
- col_name,
- rows,
- array,
- row_count,
- value_as_duration
- );
- }
- DataType::Duration(TimeUnit::Nanosecond) => {
- set_temporal_column_by_array_type!(
- DurationNanosecondArray,
- col_name,
- rows,
- array,
- row_count,
- value_as_duration
- );
- }
- DataType::Struct(_) => {
- let inner_objs =
- struct_array_to_jsonmap_array(as_struct_array(array), row_count);
- rows.iter_mut()
- .take(row_count)
- .zip(inner_objs.into_iter())
- .for_each(|(row, obj)| {
- row.insert(col_name.to_string(), Value::Object(obj));
- });
- }
- DataType::List(_) => {
- let listarr = as_list_array(array);
- rows.iter_mut()
- .zip(listarr.iter())
- .take(row_count)
- .for_each(|(row, maybe_value)| {
- if let Some(v) = maybe_value {
- row.insert(
- col_name.to_string(),
- Value::Array(array_to_json_array(&v)),
- );
- }
- });
- }
- DataType::LargeList(_) => {
- let listarr = as_large_list_array(array);
- rows.iter_mut()
- .zip(listarr.iter())
- .take(row_count)
- .for_each(|(row, maybe_value)| {
- if let Some(v) = maybe_value {
- row.insert(
- col_name.to_string(),
- Value::Array(array_to_json_array(&v)),
- );
- }
- });
- }
- DataType::Dictionary(_, value_type) => {
- let slice = array.slice(0, row_count);
- let hydrated = crate::compute::kernels::cast::cast(&slice, &value_type)
- .expect("cannot cast dictionary to underlying values");
- set_column_for_json_rows(rows, row_count, &hydrated, col_name)
- }
- _ => {
- panic!("Unsupported datatype: {:#?}", array.data_type());
- }
- }
-}
-
-/// Converts an arrow [`RecordBatch`] into a `Vec` of Serde JSON
-/// [`JsonMap`]s (objects)
-pub fn record_batches_to_json_rows(
- batches: &[RecordBatch],
-) -> Vec<JsonMap<String, Value>> {
- let mut rows: Vec<JsonMap<String, Value>> = iter::repeat(JsonMap::new())
- .take(batches.iter().map(|b| b.num_rows()).sum())
- .collect();
-
- if !rows.is_empty() {
- let schema = batches[0].schema();
- let mut base = 0;
- batches.iter().for_each(|batch| {
- let row_count = batch.num_rows();
- batch.columns().iter().enumerate().for_each(|(j, col)| {
- let col_name = schema.field(j).name();
- set_column_for_json_rows(&mut rows[base..], row_count, col, col_name);
- });
- base += row_count;
- });
- }
-
- rows
-}
-
-/// This trait defines how to format a sequence of JSON objects to a
-/// byte stream.
-pub trait JsonFormat: Debug + Default {
- #[inline]
- /// write any bytes needed at the start of the file to the writer
- fn start_stream<W: Write>(&self, _writer: &mut W) -> Result<()> {
- Ok(())
- }
-
- #[inline]
- /// write any bytes needed for the start of each row
- fn start_row<W: Write>(&self, _writer: &mut W, _is_first_row: bool) -> Result<()> {
- Ok(())
- }
-
- #[inline]
- /// write any bytes needed for the end of each row
- fn end_row<W: Write>(&self, _writer: &mut W) -> Result<()> {
- Ok(())
- }
-
- /// write any bytes needed for the start of each row
- fn end_stream<W: Write>(&self, _writer: &mut W) -> Result<()> {
- Ok(())
- }
-}
-
-/// Produces JSON output with one record per line. For example
-///
-/// ```json
-/// {"foo":1}
-/// {"bar":1}
-///
-/// ```
-#[derive(Debug, Default)]
-pub struct LineDelimited {}
-
-impl JsonFormat for LineDelimited {
- fn end_row<W: Write>(&self, writer: &mut W) -> Result<()> {
- writer.write_all(b"\n")?;
- Ok(())
- }
-}
-
-/// Produces JSON output as a single JSON array. For example
-///
-/// ```json
-/// [{"foo":1},{"bar":1}]
-/// ```
-#[derive(Debug, Default)]
-pub struct JsonArray {}
-
-impl JsonFormat for JsonArray {
- fn start_stream<W: Write>(&self, writer: &mut W) -> Result<()> {
- writer.write_all(b"[")?;
- Ok(())
- }
-
- fn start_row<W: Write>(&self, writer: &mut W, is_first_row: bool) -> Result<()> {
- if !is_first_row {
- writer.write_all(b",")?;
- }
- Ok(())
- }
-
- fn end_stream<W: Write>(&self, writer: &mut W) -> Result<()> {
- writer.write_all(b"]")?;
- Ok(())
- }
-}
-
-/// A JSON writer which serializes [`RecordBatch`]es to newline delimited JSON objects
-pub type LineDelimitedWriter<W> = Writer<W, LineDelimited>;
-
-/// A JSON writer which serializes [`RecordBatch`]es to JSON arrays
-pub type ArrayWriter<W> = Writer<W, JsonArray>;
-
-/// A JSON writer which serializes [`RecordBatch`]es to a stream of
-/// `u8` encoded JSON objects. See the module level documentation for
-/// detailed usage and examples. The specific format of the stream is
-/// controlled by the [`JsonFormat`] type parameter.
-#[derive(Debug)]
-pub struct Writer<W, F>
-where
- W: Write,
- F: JsonFormat,
-{
- /// Underlying writer to use to write bytes
- writer: W,
-
- /// Has the writer output any records yet?
- started: bool,
-
- /// Is the writer finished?
- finished: bool,
-
- /// Determines how the byte stream is formatted
- format: F,
-}
-
-impl<W, F> Writer<W, F>
-where
- W: Write,
- F: JsonFormat,
-{
- /// Construct a new writer
- pub fn new(writer: W) -> Self {
- Self {
- writer,
- started: false,
- finished: false,
- format: F::default(),
- }
- }
-
- /// Write a single JSON row to the output writer
- pub fn write_row(&mut self, row: &Value) -> Result<()> {
- let is_first_row = !self.started;
- if !self.started {
- self.format.start_stream(&mut self.writer)?;
- self.started = true;
- }
-
- self.format.start_row(&mut self.writer, is_first_row)?;
- self.writer.write_all(&serde_json::to_vec(row)?)?;
- self.format.end_row(&mut self.writer)?;
- Ok(())
- }
-
- /// Convert the [`RecordBatch`] into JSON rows, and write them to the output
- pub fn write_batches(&mut self, batches: &[RecordBatch]) -> Result<()> {
- for row in record_batches_to_json_rows(batches) {
- self.write_row(&Value::Object(row))?;
- }
- Ok(())
- }
-
- /// Finishes the output stream. This function must be called after
- /// all record batches have been produced. (e.g. producing the final `']'` if writing
- /// arrays.
- pub fn finish(&mut self) -> Result<()> {
- if self.started && !self.finished {
- self.format.end_stream(&mut self.writer)?;
- self.finished = true;
- }
- Ok(())
- }
-
- /// Unwraps this `Writer<W>`, returning the underlying writer
- pub fn into_inner(self) -> W {
- self.writer
- }
-}
-
-#[cfg(test)]
-mod tests {
- use std::convert::TryFrom;
- use std::fs::{read_to_string, File};
- use std::sync::Arc;
-
- use serde_json::json;
-
- use crate::buffer::*;
- use crate::json::reader::*;
-
- use super::*;
-
- #[test]
- fn write_simple_rows() {
- let schema = Schema::new(vec![
- Field::new("c1", DataType::Int32, true),
- Field::new("c2", DataType::Utf8, true),
- ]);
-
- let a = Int32Array::from(vec![Some(1), Some(2), Some(3), None, Some(5)]);
- let b = StringArray::from(vec![Some("a"), Some("b"), Some("c"), Some("d"), None]);
-
- let batch =
- RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a), Arc::new(b)])
- .unwrap();
-
- let mut buf = Vec::new();
- {
- let mut writer = LineDelimitedWriter::new(&mut buf);
- writer.write_batches(&[batch]).unwrap();
- }
-
- assert_eq!(
- String::from_utf8(buf).unwrap(),
- r#"{"c1":1,"c2":"a"}
-{"c1":2,"c2":"b"}
-{"c1":3,"c2":"c"}
-{"c2":"d"}
-{"c1":5}
-"#
- );
- }
-
- #[test]
- fn write_dictionary() {
- let schema = Schema::new(vec![
- Field::new(
- "c1",
- DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
- true,
- ),
- Field::new(
- "c2",
- DataType::Dictionary(Box::new(DataType::Int8), Box::new(DataType::Utf8)),
- true,
- ),
- ]);
-
- let a: DictionaryArray<Int32Type> = vec![
- Some("cupcakes"),
- Some("foo"),
- Some("foo"),
- None,
- Some("cupcakes"),
- ]
- .into_iter()
- .collect();
- let b: DictionaryArray<Int8Type> =
- vec![Some("sdsd"), Some("sdsd"), None, Some("sd"), Some("sdsd")]
- .into_iter()
- .collect();
-
- let batch =
- RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a), Arc::new(b)])
- .unwrap();
-
- let mut buf = Vec::new();
- {
- let mut writer = LineDelimitedWriter::new(&mut buf);
- writer.write_batches(&[batch]).unwrap();
- }
-
- assert_eq!(
- String::from_utf8(buf).unwrap(),
- r#"{"c1":"cupcakes","c2":"sdsd"}
-{"c1":"foo","c2":"sdsd"}
-{"c1":"foo"}
-{"c2":"sd"}
-{"c1":"cupcakes","c2":"sdsd"}
-"#
- );
- }
-
- #[test]
- fn write_timestamps() {
- let ts_string = "2018-11-13T17:11:10.011375885995";
- let ts_nanos = ts_string
- .parse::<chrono::NaiveDateTime>()
- .unwrap()
- .timestamp_nanos();
- let ts_micros = ts_nanos / 1000;
- let ts_millis = ts_micros / 1000;
- let ts_secs = ts_millis / 1000;
-
- let arr_nanos =
- TimestampNanosecondArray::from_opt_vec(vec![Some(ts_nanos), None], None);
- let arr_micros =
- TimestampMicrosecondArray::from_opt_vec(vec![Some(ts_micros), None], None);
- let arr_millis =
- TimestampMillisecondArray::from_opt_vec(vec![Some(ts_millis), None], None);
- let arr_secs =
- TimestampSecondArray::from_opt_vec(vec![Some(ts_secs), None], None);
- let arr_names = StringArray::from(vec![Some("a"), Some("b")]);
-
- let schema = Schema::new(vec![
- Field::new("nanos", arr_nanos.data_type().clone(), false),
- Field::new("micros", arr_micros.data_type().clone(), false),
- Field::new("millis", arr_millis.data_type().clone(), false),
- Field::new("secs", arr_secs.data_type().clone(), false),
- Field::new("name", arr_names.data_type().clone(), false),
- ]);
- let schema = Arc::new(schema);
-
- let batch = RecordBatch::try_new(
- schema,
- vec![
- Arc::new(arr_nanos),
- Arc::new(arr_micros),
- Arc::new(arr_millis),
- Arc::new(arr_secs),
- Arc::new(arr_names),
- ],
- )
- .unwrap();
-
- let mut buf = Vec::new();
- {
- let mut writer = LineDelimitedWriter::new(&mut buf);
- writer.write_batches(&[batch]).unwrap();
- }
-
- assert_eq!(
- String::from_utf8(buf).unwrap(),
- r#"{"nanos":"2018-11-13 17:11:10.011375885","micros":"2018-11-13 17:11:10.011375","millis":"2018-11-13 17:11:10.011","secs":"2018-11-13 17:11:10","name":"a"}
-{"name":"b"}
-"#
- );
- }
-
- #[test]
- fn write_dates() {
- let ts_string = "2018-11-13T17:11:10.011375885995";
- let ts_millis = ts_string
- .parse::<chrono::NaiveDateTime>()
- .unwrap()
- .timestamp_millis();
-
- let arr_date32 = Date32Array::from(vec![
- Some(i32::try_from(ts_millis / 1000 / (60 * 60 * 24)).unwrap()),
- None,
- ]);
- let arr_date64 = Date64Array::from(vec![Some(ts_millis), None]);
- let arr_names = StringArray::from(vec![Some("a"), Some("b")]);
-
- let schema = Schema::new(vec![
- Field::new("date32", arr_date32.data_type().clone(), false),
- Field::new("date64", arr_date64.data_type().clone(), false),
- Field::new("name", arr_names.data_type().clone(), false),
- ]);
- let schema = Arc::new(schema);
-
- let batch = RecordBatch::try_new(
- schema,
- vec![
- Arc::new(arr_date32),
- Arc::new(arr_date64),
- Arc::new(arr_names),
- ],
- )
- .unwrap();
-
- let mut buf = Vec::new();
- {
- let mut writer = LineDelimitedWriter::new(&mut buf);
- writer.write_batches(&[batch]).unwrap();
- }
-
- assert_eq!(
- String::from_utf8(buf).unwrap(),
- r#"{"date32":"2018-11-13","date64":"2018-11-13","name":"a"}
-{"name":"b"}
-"#
- );
- }
-
- #[test]
- fn write_times() {
- let arr_time32sec = Time32SecondArray::from(vec![Some(120), None]);
- let arr_time32msec = Time32MillisecondArray::from(vec![Some(120), None]);
- let arr_time64usec = Time64MicrosecondArray::from(vec![Some(120), None]);
- let arr_time64nsec = Time64NanosecondArray::from(vec![Some(120), None]);
- let arr_names = StringArray::from(vec![Some("a"), Some("b")]);
-
- let schema = Schema::new(vec![
- Field::new("time32sec", arr_time32sec.data_type().clone(), false),
- Field::new("time32msec", arr_time32msec.data_type().clone(), false),
- Field::new("time64usec", arr_time64usec.data_type().clone(), false),
- Field::new("time64nsec", arr_time64nsec.data_type().clone(), false),
- Field::new("name", arr_names.data_type().clone(), false),
- ]);
- let schema = Arc::new(schema);
-
- let batch = RecordBatch::try_new(
- schema,
- vec![
- Arc::new(arr_time32sec),
- Arc::new(arr_time32msec),
- Arc::new(arr_time64usec),
- Arc::new(arr_time64nsec),
- Arc::new(arr_names),
- ],
- )
- .unwrap();
-
- let mut buf = Vec::new();
- {
- let mut writer = LineDelimitedWriter::new(&mut buf);
- writer.write_batches(&[batch]).unwrap();
- }
-
- assert_eq!(
- String::from_utf8(buf).unwrap(),
- r#"{"time32sec":"00:02:00","time32msec":"00:00:00.120","time64usec":"00:00:00.000120","time64nsec":"00:00:00.000000120","name":"a"}
-{"name":"b"}
-"#
- );
- }
-
- #[test]
- fn write_durations() {
- let arr_durationsec = DurationSecondArray::from(vec![Some(120), None]);
- let arr_durationmsec = DurationMillisecondArray::from(vec![Some(120), None]);
- let arr_durationusec = DurationMicrosecondArray::from(vec![Some(120), None]);
- let arr_durationnsec = DurationNanosecondArray::from(vec![Some(120), None]);
- let arr_names = StringArray::from(vec![Some("a"), Some("b")]);
-
- let schema = Schema::new(vec![
- Field::new("duration_sec", arr_durationsec.data_type().clone(), false),
- Field::new("duration_msec", arr_durationmsec.data_type().clone(), false),
- Field::new("duration_usec", arr_durationusec.data_type().clone(), false),
- Field::new("duration_nsec", arr_durationnsec.data_type().clone(), false),
- Field::new("name", arr_names.data_type().clone(), false),
- ]);
- let schema = Arc::new(schema);
-
- let batch = RecordBatch::try_new(
- schema,
- vec![
- Arc::new(arr_durationsec),
- Arc::new(arr_durationmsec),
- Arc::new(arr_durationusec),
- Arc::new(arr_durationnsec),
- Arc::new(arr_names),
- ],
- )
- .unwrap();
-
- let mut buf = Vec::new();
- {
- let mut writer = LineDelimitedWriter::new(&mut buf);
- writer.write_batches(&[batch]).unwrap();
- }
-
- assert_eq!(
- String::from_utf8(buf).unwrap(),
- r#"{"duration_sec":"PT120S","duration_msec":"PT0.120S","duration_usec":"PT0.000120S","duration_nsec":"PT0.000000120S","name":"a"}
-{"name":"b"}
-"#
- );
- }
-
- #[test]
- fn write_nested_structs() {
- let schema = Schema::new(vec![
- Field::new(
- "c1",
- DataType::Struct(vec![
- Field::new("c11", DataType::Int32, false),
- Field::new(
- "c12",
- DataType::Struct(vec![Field::new("c121", DataType::Utf8, false)]),
- false,
- ),
- ]),
- false,
- ),
- Field::new("c2", DataType::Utf8, false),
- ]);
-
- let c1 = StructArray::from(vec![
- (
- Field::new("c11", DataType::Int32, false),
- Arc::new(Int32Array::from(vec![Some(1), None, Some(5)])) as ArrayRef,
- ),
- (
- Field::new(
- "c12",
- DataType::Struct(vec![Field::new("c121", DataType::Utf8, false)]),
- false,
- ),
- Arc::new(StructArray::from(vec![(
- Field::new("c121", DataType::Utf8, false),
- Arc::new(StringArray::from(vec![Some("e"), Some("f"), Some("g")]))
- as ArrayRef,
- )])) as ArrayRef,
- ),
- ]);
- let c2 = StringArray::from(vec![Some("a"), Some("b"), Some("c")]);
-
- let batch =
- RecordBatch::try_new(Arc::new(schema), vec![Arc::new(c1), Arc::new(c2)])
- .unwrap();
-
- let mut buf = Vec::new();
- {
- let mut writer = LineDelimitedWriter::new(&mut buf);
- writer.write_batches(&[batch]).unwrap();
- }
-
- assert_eq!(
- String::from_utf8(buf).unwrap(),
- r#"{"c1":{"c11":1,"c12":{"c121":"e"}},"c2":"a"}
-{"c1":{"c12":{"c121":"f"}},"c2":"b"}
-{"c1":{"c11":5,"c12":{"c121":"g"}},"c2":"c"}
-"#
- );
- }
-
- #[test]
- fn write_struct_with_list_field() {
- let field_c1 = Field::new(
- "c1",
- DataType::List(Box::new(Field::new("c_list", DataType::Utf8, false))),
- false,
- );
- let field_c2 = Field::new("c2", DataType::Int32, false);
- let schema = Schema::new(vec![field_c1.clone(), field_c2]);
-
- let a_values = StringArray::from(vec!["a", "a1", "b", "c", "d", "e"]);
- // list column rows: ["a", "a1"], ["b"], ["c"], ["d"], ["e"]
- let a_value_offsets = Buffer::from(&[0, 2, 3, 4, 5, 6].to_byte_slice());
- let a_list_data = ArrayData::builder(field_c1.data_type().clone())
- .len(5)
- .add_buffer(a_value_offsets)
- .add_child_data(a_values.data().clone())
- .null_bit_buffer(Buffer::from(vec![0b00011111]))
- .build();
- let a = ListArray::from(a_list_data);
-
- let b = Int32Array::from(vec![1, 2, 3, 4, 5]);
-
- let batch =
- RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a), Arc::new(b)])
- .unwrap();
-
- let mut buf = Vec::new();
- {
- let mut writer = LineDelimitedWriter::new(&mut buf);
- writer.write_batches(&[batch]).unwrap();
- }
-
- assert_eq!(
- String::from_utf8(buf).unwrap(),
- r#"{"c1":["a","a1"],"c2":1}
-{"c1":["b"],"c2":2}
-{"c1":["c"],"c2":3}
-{"c1":["d"],"c2":4}
-{"c1":["e"],"c2":5}
-"#
- );
- }
-
- #[test]
- fn write_nested_list() {
- let list_inner_type = Field::new(
- "a",
- DataType::List(Box::new(Field::new("b", DataType::Int32, false))),
- false,
- );
- let field_c1 = Field::new(
- "c1",
- DataType::List(Box::new(list_inner_type.clone())),
- false,
- );
- let field_c2 = Field::new("c2", DataType::Utf8, false);
- let schema = Schema::new(vec![field_c1.clone(), field_c2]);
-
- // list column rows: [[1, 2], [3]], [], [[4, 5, 6]]
- let a_values = Int32Array::from(vec![1, 2, 3, 4, 5, 6]);
-
- let a_value_offsets = Buffer::from(&[0, 2, 3, 6].to_byte_slice());
- // Construct a list array from the above two
- let a_list_data = ArrayData::builder(list_inner_type.data_type().clone())
- .len(3)
- .add_buffer(a_value_offsets)
- .null_bit_buffer(Buffer::from(vec![0b00000111]))
- .add_child_data(a_values.data().clone())
- .build();
-
- let c1_value_offsets = Buffer::from(&[0, 2, 2, 3].to_byte_slice());
- let c1_list_data = ArrayData::builder(field_c1.data_type().clone())
- .len(3)
- .add_buffer(c1_value_offsets)
- .add_child_data(a_list_data)
- .build();
-
- let c1 = ListArray::from(c1_list_data);
- let c2 = StringArray::from(vec![Some("foo"), Some("bar"), None]);
-
- let batch =
- RecordBatch::try_new(Arc::new(schema), vec![Arc::new(c1), Arc::new(c2)])
- .unwrap();
-
- let mut buf = Vec::new();
- {
- let mut writer = LineDelimitedWriter::new(&mut buf);
- writer.write_batches(&[batch]).unwrap();
- }
-
- assert_eq!(
- String::from_utf8(buf).unwrap(),
- r#"{"c1":[[1,2],[3]],"c2":"foo"}
-{"c1":[],"c2":"bar"}
-{"c1":[[4,5,6]]}
-"#
- );
- }
-
- #[test]
- fn write_list_of_struct() {
- let field_c1 = Field::new(
- "c1",
- DataType::List(Box::new(Field::new(
- "s",
- DataType::Struct(vec![
- Field::new("c11", DataType::Int32, false),
- Field::new(
- "c12",
- DataType::Struct(vec![Field::new("c121", DataType::Utf8, false)]),
- false,
- ),
- ]),
- false,
- ))),
- true,
- );
- let field_c2 = Field::new("c2", DataType::Int32, false);
- let schema = Schema::new(vec![field_c1.clone(), field_c2]);
-
- let struct_values = StructArray::from(vec![
- (
- Field::new("c11", DataType::Int32, false),
- Arc::new(Int32Array::from(vec![Some(1), None, Some(5)])) as ArrayRef,
- ),
- (
- Field::new(
- "c12",
- DataType::Struct(vec![Field::new("c121", DataType::Utf8, false)]),
- false,
- ),
- Arc::new(StructArray::from(vec![(
- Field::new("c121", DataType::Utf8, false),
- Arc::new(StringArray::from(vec![Some("e"), Some("f"), Some("g")]))
- as ArrayRef,
- )])) as ArrayRef,
- ),
- ]);
-
- // list column rows (c1):
- // [{"c11": 1, "c12": {"c121": "e"}}, {"c12": {"c121": "f"}}],
- // null,
- // [{"c11": 5, "c12": {"c121": "g"}}]
- let c1_value_offsets = Buffer::from(&[0, 2, 2, 3].to_byte_slice());
- let c1_list_data = ArrayData::builder(field_c1.data_type().clone())
- .len(3)
- .add_buffer(c1_value_offsets)
- .add_child_data(struct_values.data().clone())
- .null_bit_buffer(Buffer::from(vec![0b00000101]))
- .build();
- let c1 = ListArray::from(c1_list_data);
-
- let c2 = Int32Array::from(vec![1, 2, 3]);
-
- let batch =
- RecordBatch::try_new(Arc::new(schema), vec![Arc::new(c1), Arc::new(c2)])
- .unwrap();
-
- let mut buf = Vec::new();
- {
- let mut writer = LineDelimitedWriter::new(&mut buf);
- writer.write_batches(&[batch]).unwrap();
- }
-
- assert_eq!(
- String::from_utf8(buf).unwrap(),
- r#"{"c1":[{"c11":1,"c12":{"c121":"e"}},{"c12":{"c121":"f"}}],"c2":1}
-{"c2":2}
-{"c1":[{"c11":5,"c12":{"c121":"g"}}],"c2":3}
-"#
- );
- }
-
- fn test_write_for_file(test_file: &str) {
- let builder = ReaderBuilder::new()
- .infer_schema(None)
- .with_batch_size(1024);
- let mut reader: Reader<File> = builder
- .build::<File>(File::open(test_file).unwrap())
- .unwrap();
- let batch = reader.next().unwrap().unwrap();
-
- let mut buf = Vec::new();
- {
- let mut writer = LineDelimitedWriter::new(&mut buf);
- writer.write_batches(&[batch]).unwrap();
- }
-
- let result = String::from_utf8(buf).unwrap();
- let expected = read_to_string(test_file).unwrap();
- for (r, e) in result.lines().zip(expected.lines()) {
- let mut expected_json = serde_json::from_str::<Value>(e).unwrap();
- // remove null value from object to make comparision consistent:
- if let Value::Object(obj) = expected_json {
- expected_json = Value::Object(
- obj.into_iter().filter(|(_, v)| *v != Value::Null).collect(),
- );
- }
- assert_eq!(serde_json::from_str::<Value>(r).unwrap(), expected_json,);
- }
- }
-
- #[test]
- fn write_basic_rows() {
- test_write_for_file("test/data/basic.json");
- }
-
- #[test]
- fn write_arrays() {
- test_write_for_file("test/data/arrays.json");
- }
-
- #[test]
- fn write_basic_nulls() {
- test_write_for_file("test/data/basic_nulls.json");
- }
-
- #[test]
- fn json_writer_empty() {
- let mut writer = ArrayWriter::new(vec![] as Vec<u8>);
- writer.finish().unwrap();
- assert_eq!(String::from_utf8(writer.into_inner()).unwrap(), "");
- }
- #[test]
- fn json_writer_one_row() {
- let mut writer = ArrayWriter::new(vec![] as Vec<u8>);
- let v = json!({ "an": "object" });
- writer.write_row(&v).unwrap();
- writer.finish().unwrap();
- assert_eq!(
- String::from_utf8(writer.into_inner()).unwrap(),
- r#"[{"an":"object"}]"#
- );
- }
-
- #[test]
- fn json_writer_two_rows() {
- let mut writer = ArrayWriter::new(vec![] as Vec<u8>);
- let v = json!({ "an": "object" });
- writer.write_row(&v).unwrap();
- let v = json!({ "another": "object" });
- writer.write_row(&v).unwrap();
- writer.finish().unwrap();
- assert_eq!(
- String::from_utf8(writer.into_inner()).unwrap(),
- r#"[{"an":"object"},{"another":"object"}]"#
- );
- }
-}
diff --git a/arrow/src/lib.rs b/arrow/src/lib.rs
deleted file mode 100644
index 02f6418..0000000
--- a/arrow/src/lib.rs
+++ /dev/null
@@ -1,164 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-//! A native Rust implementation of [Apache Arrow](https://arrow.apache.org), a cross-language
-//! development platform for in-memory data.
-//!
-//! ### DataType
-//!
-//! Every [`Array`](array::Array) in this crate has an associated [`DataType`](datatypes::DataType),
-//! that specifies how its data is layed in memory and represented.
-//! Thus, a central enum of this crate is [`DataType`](datatypes::DataType), that contains the set of valid
-//! DataTypes in the specification. For example, [`DataType::Utf8`](datatypes::DataType::Utf8).
-//!
-//! ## Array
-//!
-//! The central trait of this package is the dynamically-typed [`Array`](array::Array) that
-//! represents a fixed-sized, immutable, Send + Sync Array of nullable elements. An example of such an array is [`UInt32Array`](array::UInt32Array).
-//! One way to think about an arrow [`Array`](array::Array) is a `Arc<[Option<T>; len]>` where T can be anything ranging from an integer to a string, or even
-//! another [`Array`](array::Array).
-//!
-//! [`Arrays`](array::Array) have [`len()`](array::Array::len), [`data_type()`](array::Array::data_type), and the nullability of each of its elements,
-//! can be obtained via [`is_null(index)`](array::Array::is_null). To downcast an [`Array`](array::Array) to a specific implementation, you can use
-//!
-//! ```rust
-//! use arrow::array::{Array, UInt32Array};
-//! let array = UInt32Array::from(vec![Some(1), None, Some(3)]);
-//! assert_eq!(array.len(), 3);
-//! assert_eq!(array.value(0), 1);
-//! assert_eq!(array.is_null(1), true);
-//! ```
-//!
-//! To make the array dynamically typed, we wrap it in an [`Arc`](std::sync::Arc):
-//!
-//! ```rust
-//! # use std::sync::Arc;
-//! use arrow::datatypes::DataType;
-//! use arrow::array::{UInt32Array, ArrayRef};
-//! # let array = UInt32Array::from(vec![Some(1), None, Some(3)]);
-//! let array: ArrayRef = Arc::new(array);
-//! assert_eq!(array.len(), 3);
-//! // array.value() is not available in the dynamically-typed version
-//! assert_eq!(array.is_null(1), true);
-//! assert_eq!(array.data_type(), &DataType::UInt32);
-//! ```
-//!
-//! to downcast, use `as_any()`:
-//!
-//! ```rust
-//! # use std::sync::Arc;
-//! # use arrow::array::{UInt32Array, ArrayRef};
-//! # let array = UInt32Array::from(vec![Some(1), None, Some(3)]);
-//! # let array: ArrayRef = Arc::new(array);
-//! let array = array.as_any().downcast_ref::<UInt32Array>().unwrap();
-//! assert_eq!(array.value(0), 1);
-//! ```
-//!
-//! ## Memory and Buffers
-//!
-//! Data in [`Array`](array::Array) is stored in [`ArrayData`](array::ArrayData), that in turn
-//! is a collection of other [`ArrayData`](array::ArrayData) and [`Buffers`](buffer::Buffer).
-//! [`Buffers`](buffer::Buffer) is the central struct that array implementations use keep allocated memory and pointers.
-//! The [`MutableBuffer`](buffer::MutableBuffer) is the mutable counter-part of[`Buffer`](buffer::Buffer).
-//! These are the lowest abstractions of this crate, and are used throughout the crate to
-//! efficiently allocate, write, read and deallocate memory.
-//!
-//! ## Field, Schema and RecordBatch
-//!
-//! [`Field`](datatypes::Field) is a struct that contains an array's metadata (datatype and whether its values
-//! can be null), and a name. [`Schema`](datatypes::Schema) is a vector of fields with optional metadata.
-//! Together, they form the basis of a schematic representation of a group of [`Arrays`](array::Array).
-//!
-//! In fact, [`RecordBatch`](record_batch::RecordBatch) is a struct with a [`Schema`](datatypes::Schema) and a vector of
-//! [`Array`](array::Array)s, all with the same `len`. A record batch is the highest order struct that this crate currently offers
-//! and is broadly used to represent a table where each column in an `Array`.
-//!
-//! ## Compute
-//!
-//! This crate offers many operations (called kernels) to operate on `Array`s, that you can find at [compute::kernels].
-//! It has both vertical and horizontal operations, and some of them have an SIMD implementation.
-//!
-//! ## Status
-//!
-//! This crate has most of the implementation of the arrow specification. Specifically, it supports the following types:
-//!
-//! * All arrow primitive types, such as [`Int32Array`](array::UInt8Array), [`BooleanArray`](array::BooleanArray) and [`Float64Array`](array::Float64Array).
-//! * All arrow variable length types, such as [`StringArray`](array::StringArray) and [`BinaryArray`](array::BinaryArray)
-//! * All composite types such as [`StructArray`](array::StructArray) and [`ListArray`](array::ListArray)
-//! * Dictionary types [`DictionaryArray`](array::DictionaryArray)
-
-//!
-//! This crate also implements many common vertical operations:
-//! * all mathematical binary operators, such as [`subtract`](compute::kernels::arithmetic::subtract)
-//! * all boolean binary operators such as [`equality`](compute::kernels::comparison::eq)
-//! * [`cast`](compute::kernels::cast::cast)
-//! * [`filter`](compute::kernels::filter::filter)
-//! * [`take`](compute::kernels::take::take) and [`limit`](compute::kernels::limit::limit)
-//! * [`sort`](compute::kernels::sort::sort)
-//! * some string operators such as [`substring`](compute::kernels::substring::substring) and [`length`](compute::kernels::length::length)
-//!
-//! as well as some horizontal operations, such as
-//!
-//! * [`min`](compute::kernels::aggregate::min) and [`max`](compute::kernels::aggregate::max)
-//! * [`sum`](compute::kernels::aggregate::sum)
-//!
-//! Finally, this crate implements some readers and writers to different formats:
-//!
-//! * json: [reader](json::reader::Reader)
-//! * csv: [reader](csv::reader::Reader) and [writer](csv::writer::Writer)
-//! * ipc: [reader](ipc::reader::StreamReader) and [writer](ipc::writer::FileWriter)
-//!
-//! The parquet implementation is on a [separate crate](https://crates.io/crates/parquet)
-
-#![cfg_attr(feature = "avx512", feature(stdsimd))]
-#![cfg_attr(feature = "avx512", feature(repr_simd))]
-#![cfg_attr(feature = "avx512", feature(avx512_target_feature))]
-#![allow(dead_code)]
-#![allow(non_camel_case_types)]
-#![deny(clippy::redundant_clone)]
-#![allow(
- // introduced to ignore lint errors when upgrading from 2020-04-22 to 2020-11-14
- clippy::float_equality_without_abs,
- clippy::type_complexity,
- // upper_case_acronyms lint was introduced in Rust 1.51.
- // It is triggered in the ffi module, and ipc::gen, which we have no control over
- clippy::upper_case_acronyms,
- clippy::vec_init_then_push
-)]
-#![allow(bare_trait_objects)]
-#![warn(missing_debug_implementations)]
-
-pub mod alloc;
-mod arch;
-pub mod array;
-pub mod bitmap;
-pub mod buffer;
-mod bytes;
-pub mod compute;
-#[cfg(feature = "csv")]
-pub mod csv;
-pub mod datatypes;
-pub mod error;
-pub mod ffi;
-#[cfg(feature = "ipc")]
-pub mod ipc;
-pub mod json;
-pub mod record_batch;
-pub mod temporal_conversions;
-pub mod tensor;
-pub mod util;
-mod zz_memory_check;
diff --git a/arrow/src/record_batch.rs b/arrow/src/record_batch.rs
deleted file mode 100644
index f1fd867..0000000
--- a/arrow/src/record_batch.rs
+++ /dev/null
@@ -1,565 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-//! A two-dimensional batch of column-oriented data with a defined
-//! [schema](crate::datatypes::Schema).
-
-use std::sync::Arc;
-
-use crate::array::*;
-use crate::datatypes::*;
-use crate::error::{ArrowError, Result};
-
-/// A two-dimensional batch of column-oriented data with a defined
-/// [schema](crate::datatypes::Schema).
-///
-/// A `RecordBatch` is a two-dimensional dataset of a number of
-/// contiguous arrays, each the same length.
-/// A record batch has a schema which must match its arrays’
-/// datatypes.
-///
-/// Record batches are a convenient unit of work for various
-/// serialization and computation functions, possibly incremental.
-/// See also [CSV reader](crate::csv::Reader) and
-/// [JSON reader](crate::json::Reader).
-#[derive(Clone, Debug)]
-pub struct RecordBatch {
- schema: SchemaRef,
- columns: Vec<Arc<Array>>,
-}
-
-impl RecordBatch {
- /// Creates a `RecordBatch` from a schema and columns.
- ///
- /// Expects the following:
- /// * the vec of columns to not be empty
- /// * the schema and column data types to have equal lengths
- /// and match
- /// * each array in columns to have the same length
- ///
- /// If the conditions are not met, an error is returned.
- ///
- /// # Example
- ///
- /// ```
- /// use std::sync::Arc;
- /// use arrow::array::Int32Array;
- /// use arrow::datatypes::{Schema, Field, DataType};
- /// use arrow::record_batch::RecordBatch;
- ///
- /// # fn main() -> arrow::error::Result<()> {
- /// let id_array = Int32Array::from(vec![1, 2, 3, 4, 5]);
- /// let schema = Schema::new(vec![
- /// Field::new("id", DataType::Int32, false)
- /// ]);
- ///
- /// let batch = RecordBatch::try_new(
- /// Arc::new(schema),
- /// vec![Arc::new(id_array)]
- /// )?;
- /// # Ok(())
- /// # }
- /// ```
- pub fn try_new(schema: SchemaRef, columns: Vec<ArrayRef>) -> Result<Self> {
- let options = RecordBatchOptions::default();
- Self::validate_new_batch(&schema, columns.as_slice(), &options)?;
- Ok(RecordBatch { schema, columns })
- }
-
- /// Creates a `RecordBatch` from a schema and columns, with additional options,
- /// such as whether to strictly validate field names.
- ///
- /// See [`RecordBatch::try_new`] for the expected conditions.
- pub fn try_new_with_options(
- schema: SchemaRef,
- columns: Vec<ArrayRef>,
- options: &RecordBatchOptions,
- ) -> Result<Self> {
- Self::validate_new_batch(&schema, columns.as_slice(), options)?;
- Ok(RecordBatch { schema, columns })
- }
-
- /// Creates a new empty [`RecordBatch`].
- pub fn new_empty(schema: SchemaRef) -> Self {
- let columns = schema
- .fields()
- .iter()
- .map(|field| new_empty_array(field.data_type()))
- .collect();
- RecordBatch { schema, columns }
- }
-
- /// Validate the schema and columns using [`RecordBatchOptions`]. Returns an error
- /// if any validation check fails.
- fn validate_new_batch(
- schema: &SchemaRef,
- columns: &[ArrayRef],
- options: &RecordBatchOptions,
- ) -> Result<()> {
- // check that there are some columns
- if columns.is_empty() {
- return Err(ArrowError::InvalidArgumentError(
- "at least one column must be defined to create a record batch"
- .to_string(),
- ));
- }
- // check that number of fields in schema match column length
- if schema.fields().len() != columns.len() {
- return Err(ArrowError::InvalidArgumentError(format!(
- "number of columns({}) must match number of fields({}) in schema",
- columns.len(),
- schema.fields().len(),
- )));
- }
- // check that all columns have the same row count, and match the schema
- let len = columns[0].data().len();
-
- // This is a bit repetitive, but it is better to check the condition outside the loop
- if options.match_field_names {
- for (i, column) in columns.iter().enumerate() {
- if column.len() != len {
- return Err(ArrowError::InvalidArgumentError(
- "all columns in a record batch must have the same length"
- .to_string(),
- ));
- }
- if column.data_type() != schema.field(i).data_type() {
- return Err(ArrowError::InvalidArgumentError(format!(
- "column types must match schema types, expected {:?} but found {:?} at column index {}",
- schema.field(i).data_type(),
- column.data_type(),
- i)));
- }
- }
- } else {
- for (i, column) in columns.iter().enumerate() {
- if column.len() != len {
- return Err(ArrowError::InvalidArgumentError(
- "all columns in a record batch must have the same length"
- .to_string(),
- ));
- }
- if !column
- .data_type()
- .equals_datatype(schema.field(i).data_type())
- {
- return Err(ArrowError::InvalidArgumentError(format!(
- "column types must match schema types, expected {:?} but found {:?} at column index {}",
- schema.field(i).data_type(),
- column.data_type(),
- i)));
- }
- }
- }
-
- Ok(())
- }
-
- /// Returns the [`Schema`](crate::datatypes::Schema) of the record batch.
- pub fn schema(&self) -> SchemaRef {
- self.schema.clone()
- }
-
- /// Returns the number of columns in the record batch.
- ///
- /// # Example
- ///
- /// ```
- /// use std::sync::Arc;
- /// use arrow::array::Int32Array;
- /// use arrow::datatypes::{Schema, Field, DataType};
- /// use arrow::record_batch::RecordBatch;
- ///
- /// # fn main() -> arrow::error::Result<()> {
- /// let id_array = Int32Array::from(vec![1, 2, 3, 4, 5]);
- /// let schema = Schema::new(vec![
- /// Field::new("id", DataType::Int32, false)
- /// ]);
- ///
- /// let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(id_array)])?;
- ///
- /// assert_eq!(batch.num_columns(), 1);
- /// # Ok(())
- /// # }
- /// ```
- pub fn num_columns(&self) -> usize {
- self.columns.len()
- }
-
- /// Returns the number of rows in each column.
- ///
- /// # Panics
- ///
- /// Panics if the `RecordBatch` contains no columns.
- ///
- /// # Example
- ///
- /// ```
- /// use std::sync::Arc;
- /// use arrow::array::Int32Array;
- /// use arrow::datatypes::{Schema, Field, DataType};
- /// use arrow::record_batch::RecordBatch;
- ///
- /// # fn main() -> arrow::error::Result<()> {
- /// let id_array = Int32Array::from(vec![1, 2, 3, 4, 5]);
- /// let schema = Schema::new(vec![
- /// Field::new("id", DataType::Int32, false)
- /// ]);
- ///
- /// let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(id_array)])?;
- ///
- /// assert_eq!(batch.num_rows(), 5);
- /// # Ok(())
- /// # }
- /// ```
- pub fn num_rows(&self) -> usize {
- self.columns[0].data().len()
- }
-
- /// Get a reference to a column's array by index.
- ///
- /// # Panics
- ///
- /// Panics if `index` is outside of `0..num_columns`.
- pub fn column(&self, index: usize) -> &ArrayRef {
- &self.columns[index]
- }
-
- /// Get a reference to all columns in the record batch.
- pub fn columns(&self) -> &[ArrayRef] {
- &self.columns[..]
- }
-
- /// Create a `RecordBatch` from an iterable list of pairs of the
- /// form `(field_name, array)`, with the same requirements on
- /// fields and arrays as [`RecordBatch::try_new`]. This method is
- /// often used to create a single `RecordBatch` from arrays,
- /// e.g. for testing.
- ///
- /// The resulting schema is marked as nullable for each column if
- /// the array for that column is has any nulls. To explicitly
- /// specify nullibility, use [`RecordBatch::try_from_iter_with_nullable`]
- ///
- /// Example:
- /// ```
- /// use std::sync::Arc;
- /// use arrow::array::{ArrayRef, Int32Array, StringArray};
- /// use arrow::datatypes::{Schema, Field, DataType};
- /// use arrow::record_batch::RecordBatch;
- ///
- /// let a: ArrayRef = Arc::new(Int32Array::from(vec![1, 2]));
- /// let b: ArrayRef = Arc::new(StringArray::from(vec!["a", "b"]));
- ///
- /// let record_batch = RecordBatch::try_from_iter(vec![
- /// ("a", a),
- /// ("b", b),
- /// ]);
- /// ```
- pub fn try_from_iter<I, F>(value: I) -> Result<Self>
- where
- I: IntoIterator<Item = (F, ArrayRef)>,
- F: AsRef<str>,
- {
- // TODO: implement `TryFrom` trait, once
- // https://github.com/rust-lang/rust/issues/50133 is no longer an
- // issue
- let iter = value.into_iter().map(|(field_name, array)| {
- let nullable = array.null_count() > 0;
- (field_name, array, nullable)
- });
-
- Self::try_from_iter_with_nullable(iter)
- }
-
- /// Create a `RecordBatch` from an iterable list of tuples of the
- /// form `(field_name, array, nullable)`, with the same requirements on
- /// fields and arrays as [`RecordBatch::try_new`]. This method is often
- /// used to create a single `RecordBatch` from arrays, e.g. for
- /// testing.
- ///
- /// Example:
- /// ```
- /// use std::sync::Arc;
- /// use arrow::array::{ArrayRef, Int32Array, StringArray};
- /// use arrow::datatypes::{Schema, Field, DataType};
- /// use arrow::record_batch::RecordBatch;
- ///
- /// let a: ArrayRef = Arc::new(Int32Array::from(vec![1, 2]));
- /// let b: ArrayRef = Arc::new(StringArray::from(vec![Some("a"), Some("b")]));
- ///
- /// // Note neither `a` nor `b` has any actual nulls, but we mark
- /// // b an nullable
- /// let record_batch = RecordBatch::try_from_iter_with_nullable(vec![
- /// ("a", a, false),
- /// ("b", b, true),
- /// ]);
- /// ```
- pub fn try_from_iter_with_nullable<I, F>(value: I) -> Result<Self>
- where
- I: IntoIterator<Item = (F, ArrayRef, bool)>,
- F: AsRef<str>,
- {
- // TODO: implement `TryFrom` trait, once
- // https://github.com/rust-lang/rust/issues/50133 is no longer an
- // issue
- let (fields, columns) = value
- .into_iter()
- .map(|(field_name, array, nullable)| {
- let field_name = field_name.as_ref();
- let field = Field::new(field_name, array.data_type().clone(), nullable);
- (field, array)
- })
- .unzip();
-
- let schema = Arc::new(Schema::new(fields));
- RecordBatch::try_new(schema, columns)
- }
-}
-
-/// Options that control the behaviour used when creating a [`RecordBatch`].
-#[derive(Debug)]
-pub struct RecordBatchOptions {
- /// Match field names of structs and lists. If set to `true`, the names must match.
- pub match_field_names: bool,
-}
-
-impl Default for RecordBatchOptions {
- fn default() -> Self {
- Self {
- match_field_names: true,
- }
- }
-}
-
-impl From<&StructArray> for RecordBatch {
- /// Create a record batch from struct array, where each field of
- /// the `StructArray` becomes a `Field` in the schema.
- ///
- /// This currently does not flatten and nested struct types
- fn from(struct_array: &StructArray) -> Self {
- if let DataType::Struct(fields) = struct_array.data_type() {
- let schema = Schema::new(fields.clone());
- let columns = struct_array.boxed_fields.clone();
- RecordBatch {
- schema: Arc::new(schema),
- columns,
- }
- } else {
- unreachable!("unable to get datatype as struct")
- }
- }
-}
-
-impl From<RecordBatch> for StructArray {
- fn from(batch: RecordBatch) -> Self {
- batch
- .schema
- .fields
- .iter()
- .zip(batch.columns.iter())
- .map(|t| (t.0.clone(), t.1.clone()))
- .collect::<Vec<(Field, ArrayRef)>>()
- .into()
- }
-}
-
-/// Trait for types that can read `RecordBatch`'s.
-pub trait RecordBatchReader: Iterator<Item = Result<RecordBatch>> {
- /// Returns the schema of this `RecordBatchReader`.
- ///
- /// Implementation of this trait should guarantee that all `RecordBatch`'s returned by this
- /// reader should have the same schema as returned from this method.
- fn schema(&self) -> SchemaRef;
-
- /// Reads the next `RecordBatch`.
- #[deprecated(
- since = "2.0.0",
- note = "This method is deprecated in favour of `next` from the trait Iterator."
- )]
- fn next_batch(&mut self) -> Result<Option<RecordBatch>> {
- self.next().transpose()
- }
-}
-
-#[cfg(test)]
-mod tests {
- use super::*;
-
- use crate::buffer::Buffer;
-
- #[test]
- fn create_record_batch() {
- let schema = Schema::new(vec![
- Field::new("a", DataType::Int32, false),
- Field::new("b", DataType::Utf8, false),
- ]);
-
- let a = Int32Array::from(vec![1, 2, 3, 4, 5]);
- let b = StringArray::from(vec!["a", "b", "c", "d", "e"]);
-
- let record_batch =
- RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a), Arc::new(b)])
- .unwrap();
- check_batch(record_batch)
- }
-
- fn check_batch(record_batch: RecordBatch) {
- assert_eq!(5, record_batch.num_rows());
- assert_eq!(2, record_batch.num_columns());
- assert_eq!(&DataType::Int32, record_batch.schema().field(0).data_type());
- assert_eq!(&DataType::Utf8, record_batch.schema().field(1).data_type());
- assert_eq!(5, record_batch.column(0).data().len());
- assert_eq!(5, record_batch.column(1).data().len());
- }
-
- #[test]
- fn create_record_batch_try_from_iter() {
- let a: ArrayRef = Arc::new(Int32Array::from(vec![
- Some(1),
- Some(2),
- None,
- Some(4),
- Some(5),
- ]));
- let b: ArrayRef = Arc::new(StringArray::from(vec!["a", "b", "c", "d", "e"]));
-
- let record_batch = RecordBatch::try_from_iter(vec![("a", a), ("b", b)])
- .expect("valid conversion");
-
- let expected_schema = Schema::new(vec![
- Field::new("a", DataType::Int32, true),
- Field::new("b", DataType::Utf8, false),
- ]);
- assert_eq!(record_batch.schema().as_ref(), &expected_schema);
- check_batch(record_batch);
- }
-
- #[test]
- fn create_record_batch_try_from_iter_with_nullable() {
- let a: ArrayRef = Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5]));
- let b: ArrayRef = Arc::new(StringArray::from(vec!["a", "b", "c", "d", "e"]));
-
- // Note there are no nulls in a or b, but we specify that b is nullable
- let record_batch = RecordBatch::try_from_iter_with_nullable(vec![
- ("a", a, false),
- ("b", b, true),
- ])
- .expect("valid conversion");
-
- let expected_schema = Schema::new(vec![
- Field::new("a", DataType::Int32, false),
- Field::new("b", DataType::Utf8, true),
- ]);
- assert_eq!(record_batch.schema().as_ref(), &expected_schema);
- check_batch(record_batch);
- }
-
- #[test]
- fn create_record_batch_schema_mismatch() {
- let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]);
-
- let a = Int64Array::from(vec![1, 2, 3, 4, 5]);
-
- let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a)]);
- assert!(!batch.is_ok());
- }
-
- #[test]
- fn create_record_batch_field_name_mismatch() {
- let struct_fields = vec![
- Field::new("a1", DataType::Int32, false),
- Field::new(
- "a2",
- DataType::List(Box::new(Field::new("item", DataType::Int8, false))),
- false,
- ),
- ];
- let struct_type = DataType::Struct(struct_fields);
- let schema = Arc::new(Schema::new(vec![Field::new("a", struct_type, true)]));
-
- let a1: ArrayRef = Arc::new(Int32Array::from(vec![1, 2]));
- let a2_child = Int8Array::from(vec![1, 2, 3, 4]);
- let a2 = ArrayDataBuilder::new(DataType::List(Box::new(Field::new(
- "array",
- DataType::Int8,
- false,
- ))))
- .add_child_data(a2_child.data().clone())
- .len(2)
- .add_buffer(Buffer::from(vec![0i32, 3, 4].to_byte_slice()))
- .build();
- let a2: ArrayRef = Arc::new(ListArray::from(a2));
- let a = ArrayDataBuilder::new(DataType::Struct(vec![
- Field::new("aa1", DataType::Int32, false),
- Field::new("a2", a2.data_type().clone(), false),
- ]))
- .add_child_data(a1.data().clone())
- .add_child_data(a2.data().clone())
- .len(2)
- .build();
- let a: ArrayRef = Arc::new(StructArray::from(a));
-
- // creating the batch with field name validation should fail
- let batch = RecordBatch::try_new(schema.clone(), vec![a.clone()]);
- assert!(batch.is_err());
-
- // creating the batch without field name validation should pass
- let options = RecordBatchOptions {
- match_field_names: false,
- };
- let batch = RecordBatch::try_new_with_options(schema, vec![a], &options);
- assert!(batch.is_ok());
- }
-
- #[test]
- fn create_record_batch_record_mismatch() {
- let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]);
-
- let a = Int32Array::from(vec![1, 2, 3, 4, 5]);
- let b = Int32Array::from(vec![1, 2, 3, 4, 5]);
-
- let batch =
- RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a), Arc::new(b)]);
- assert!(!batch.is_ok());
- }
-
- #[test]
- fn create_record_batch_from_struct_array() {
- let boolean = Arc::new(BooleanArray::from(vec![false, false, true, true]));
- let int = Arc::new(Int32Array::from(vec![42, 28, 19, 31]));
- let struct_array = StructArray::from(vec![
- (
- Field::new("b", DataType::Boolean, false),
- boolean.clone() as ArrayRef,
- ),
- (
- Field::new("c", DataType::Int32, false),
- int.clone() as ArrayRef,
- ),
- ]);
-
- let batch = RecordBatch::from(&struct_array);
- assert_eq!(2, batch.num_columns());
- assert_eq!(4, batch.num_rows());
- assert_eq!(
- struct_array.data_type(),
- &DataType::Struct(batch.schema().fields().to_vec())
- );
- assert_eq!(batch.column(0).as_ref(), boolean.as_ref());
- assert_eq!(batch.column(1).as_ref(), int.as_ref());
- }
-}
diff --git a/arrow/src/temporal_conversions.rs b/arrow/src/temporal_conversions.rs
deleted file mode 100644
index 2d6d677..0000000
--- a/arrow/src/temporal_conversions.rs
+++ /dev/null
@@ -1,151 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-//! Conversion methods for dates and times.
-
-use chrono::{Duration, NaiveDateTime, NaiveTime};
-
-/// Number of seconds in a day
-const SECONDS_IN_DAY: i64 = 86_400;
-/// Number of milliseconds in a second
-const MILLISECONDS: i64 = 1_000;
-/// Number of microseconds in a second
-const MICROSECONDS: i64 = 1_000_000;
-/// Number of nanoseconds in a second
-const NANOSECONDS: i64 = 1_000_000_000;
-
-/// converts a `i32` representing a `date32` to [`NaiveDateTime`]
-#[inline]
-pub fn date32_to_datetime(v: i32) -> NaiveDateTime {
- NaiveDateTime::from_timestamp(v as i64 * SECONDS_IN_DAY, 0)
-}
-
-/// converts a `i64` representing a `date64` to [`NaiveDateTime`]
-#[inline]
-pub fn date64_to_datetime(v: i64) -> NaiveDateTime {
- NaiveDateTime::from_timestamp(
- // extract seconds from milliseconds
- v / MILLISECONDS,
- // discard extracted seconds and convert milliseconds to nanoseconds
- (v % MILLISECONDS * MICROSECONDS) as u32,
- )
-}
-
-/// converts a `i32` representing a `time32(s)` to [`NaiveDateTime`]
-#[inline]
-pub fn time32s_to_time(v: i32) -> NaiveTime {
- NaiveTime::from_num_seconds_from_midnight(v as u32, 0)
-}
-
-/// converts a `i32` representing a `time32(ms)` to [`NaiveDateTime`]
-#[inline]
-pub fn time32ms_to_time(v: i32) -> NaiveTime {
- let v = v as i64;
- NaiveTime::from_num_seconds_from_midnight(
- // extract seconds from milliseconds
- (v / MILLISECONDS) as u32,
- // discard extracted seconds and convert milliseconds to
- // nanoseconds
- (v % MILLISECONDS * MICROSECONDS) as u32,
- )
-}
-
-/// converts a `i64` representing a `time64(us)` to [`NaiveDateTime`]
-#[inline]
-pub fn time64us_to_time(v: i64) -> NaiveTime {
- NaiveTime::from_num_seconds_from_midnight(
- // extract seconds from microseconds
- (v / MICROSECONDS) as u32,
- // discard extracted seconds and convert microseconds to
- // nanoseconds
- (v % MICROSECONDS * MILLISECONDS) as u32,
- )
-}
-
-/// converts a `i64` representing a `time64(ns)` to [`NaiveDateTime`]
-#[inline]
-pub fn time64ns_to_time(v: i64) -> NaiveTime {
- NaiveTime::from_num_seconds_from_midnight(
- // extract seconds from nanoseconds
- (v / NANOSECONDS) as u32,
- // discard extracted seconds
- (v % NANOSECONDS) as u32,
- )
-}
-
-/// converts a `i64` representing a `timestamp(s)` to [`NaiveDateTime`]
-#[inline]
-pub fn timestamp_s_to_datetime(v: i64) -> NaiveDateTime {
- NaiveDateTime::from_timestamp(v, 0)
-}
-
-/// converts a `i64` representing a `timestamp(ms)` to [`NaiveDateTime`]
-#[inline]
-pub fn timestamp_ms_to_datetime(v: i64) -> NaiveDateTime {
- NaiveDateTime::from_timestamp(
- // extract seconds from milliseconds
- v / MILLISECONDS,
- // discard extracted seconds and convert milliseconds to nanoseconds
- (v % MILLISECONDS * MICROSECONDS) as u32,
- )
-}
-
-/// converts a `i64` representing a `timestamp(us)` to [`NaiveDateTime`]
-#[inline]
-pub fn timestamp_us_to_datetime(v: i64) -> NaiveDateTime {
- NaiveDateTime::from_timestamp(
- // extract seconds from microseconds
- v / MICROSECONDS,
- // discard extracted seconds and convert microseconds to nanoseconds
- (v % MICROSECONDS * MILLISECONDS) as u32,
- )
-}
-
-/// converts a `i64` representing a `timestamp(ns)` to [`NaiveDateTime`]
-#[inline]
-pub fn timestamp_ns_to_datetime(v: i64) -> NaiveDateTime {
- NaiveDateTime::from_timestamp(
- // extract seconds from nanoseconds
- v / NANOSECONDS,
- // discard extracted seconds
- (v % NANOSECONDS) as u32,
- )
-}
-
-/// converts a `i64` representing a `duration(s)` to [`Duration`]
-#[inline]
-pub fn duration_s_to_duration(v: i64) -> Duration {
- Duration::seconds(v)
-}
-
-/// converts a `i64` representing a `duration(ms)` to [`Duration`]
-#[inline]
-pub fn duration_ms_to_duration(v: i64) -> Duration {
- Duration::milliseconds(v)
-}
-
-/// converts a `i64` representing a `duration(us)` to [`Duration`]
-#[inline]
-pub fn duration_us_to_duration(v: i64) -> Duration {
- Duration::microseconds(v)
-}
-
-/// converts a `i64` representing a `duration(ns)` to [`Duration`]
-#[inline]
-pub fn duration_ns_to_duration(v: i64) -> Duration {
- Duration::nanoseconds(v)
-}
diff --git a/arrow/src/tensor.rs b/arrow/src/tensor.rs
deleted file mode 100644
index 35e45a2..0000000
--- a/arrow/src/tensor.rs
+++ /dev/null
@@ -1,495 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-//! Arrow Tensor Type, defined in
-//! [`format/Tensor.fbs`](https://github.com/apache/arrow/blob/master/format/Tensor.fbs).
-
-use std::marker::PhantomData;
-use std::mem;
-
-use crate::buffer::Buffer;
-use crate::datatypes::*;
-
-use crate::error::{ArrowError, Result};
-
-/// Computes the strides required assuming a row major memory layout
-fn compute_row_major_strides<T: ArrowPrimitiveType>(
- shape: &[usize],
-) -> Result<Vec<usize>> {
- let mut remaining_bytes = mem::size_of::<T::Native>();
-
- for i in shape {
- if let Some(val) = remaining_bytes.checked_mul(*i) {
- remaining_bytes = val;
- } else {
- return Err(ArrowError::ComputeError(
- "overflow occurred when computing row major strides.".to_string(),
- ));
- }
- }
-
- let mut strides = Vec::<usize>::new();
- for i in shape {
- remaining_bytes /= *i;
- strides.push(remaining_bytes);
- }
-
- Ok(strides)
-}
-
-/// Computes the strides required assuming a column major memory layout
-fn compute_column_major_strides<T: ArrowPrimitiveType>(
- shape: &[usize],
-) -> Result<Vec<usize>> {
- let mut remaining_bytes = mem::size_of::<T::Native>();
- let mut strides = Vec::<usize>::new();
-
- for i in shape {
- strides.push(remaining_bytes);
-
- if let Some(val) = remaining_bytes.checked_mul(*i) {
- remaining_bytes = val;
- } else {
- return Err(ArrowError::ComputeError(
- "overflow occurred when computing column major strides.".to_string(),
- ));
- }
- }
-
- Ok(strides)
-}
-
-/// Tensor of primitive types
-#[derive(Debug)]
-pub struct Tensor<'a, T: ArrowPrimitiveType> {
- data_type: DataType,
- buffer: Buffer,
- shape: Option<Vec<usize>>,
- strides: Option<Vec<usize>>,
- names: Option<Vec<&'a str>>,
- _marker: PhantomData<T>,
-}
-
-pub type BooleanTensor<'a> = Tensor<'a, BooleanType>;
-pub type Int8Tensor<'a> = Tensor<'a, Int8Type>;
-pub type Int16Tensor<'a> = Tensor<'a, Int16Type>;
-pub type Int32Tensor<'a> = Tensor<'a, Int32Type>;
-pub type Int64Tensor<'a> = Tensor<'a, Int64Type>;
-pub type UInt8Tensor<'a> = Tensor<'a, UInt8Type>;
-pub type UInt16Tensor<'a> = Tensor<'a, UInt16Type>;
-pub type UInt32Tensor<'a> = Tensor<'a, UInt32Type>;
-pub type UInt64Tensor<'a> = Tensor<'a, UInt64Type>;
-pub type Float32Tensor<'a> = Tensor<'a, Float32Type>;
-pub type Float64Tensor<'a> = Tensor<'a, Float64Type>;
-
-impl<'a, T: ArrowPrimitiveType> Tensor<'a, T> {
- /// Creates a new `Tensor`
- pub fn try_new(
- buffer: Buffer,
- shape: Option<Vec<usize>>,
- strides: Option<Vec<usize>>,
- names: Option<Vec<&'a str>>,
- ) -> Result<Self> {
- match shape {
- None => {
- if buffer.len() != mem::size_of::<T::Native>() {
- return Err(ArrowError::InvalidArgumentError(
- "underlying buffer should only contain a single tensor element"
- .to_string(),
- ));
- }
-
- if strides != None {
- return Err(ArrowError::InvalidArgumentError(
- "expected None strides for tensor with no shape".to_string(),
- ));
- }
-
- if names != None {
- return Err(ArrowError::InvalidArgumentError(
- "expected None names for tensor with no shape".to_string(),
- ));
- }
- }
-
- Some(ref s) => {
- if let Some(ref st) = strides {
- if st.len() != s.len() {
- return Err(ArrowError::InvalidArgumentError(
- "shape and stride dimensions differ".to_string(),
- ));
- }
- }
-
- if let Some(ref n) = names {
- if n.len() != s.len() {
- return Err(ArrowError::InvalidArgumentError(
- "number of dimensions and number of dimension names differ"
- .to_string(),
- ));
- }
- }
-
- let total_elements: usize = s.iter().product();
- if total_elements != (buffer.len() / mem::size_of::<T::Native>()) {
- return Err(ArrowError::InvalidArgumentError(
- "number of elements in buffer does not match dimensions"
- .to_string(),
- ));
- }
- }
- };
-
- // Checking that the tensor strides used for construction are correct
- // otherwise a row major stride is calculated and used as value for the tensor
- let tensor_strides = {
- if let Some(st) = strides {
- if let Some(ref s) = shape {
- if compute_row_major_strides::<T>(s)? == st
- || compute_column_major_strides::<T>(s)? == st
- {
- Some(st)
- } else {
- return Err(ArrowError::InvalidArgumentError(
- "the input stride does not match the selected shape"
- .to_string(),
- ));
- }
- } else {
- Some(st)
- }
- } else if let Some(ref s) = shape {
- Some(compute_row_major_strides::<T>(s)?)
- } else {
- None
- }
- };
-
- Ok(Self {
- data_type: T::DATA_TYPE,
- buffer,
- shape,
- strides: tensor_strides,
- names,
- _marker: PhantomData,
- })
- }
-
- /// Creates a new Tensor using row major memory layout
- pub fn new_row_major(
- buffer: Buffer,
- shape: Option<Vec<usize>>,
- names: Option<Vec<&'a str>>,
- ) -> Result<Self> {
- if let Some(ref s) = shape {
- let strides = Some(compute_row_major_strides::<T>(&s)?);
-
- Self::try_new(buffer, shape, strides, names)
- } else {
- Err(ArrowError::InvalidArgumentError(
- "shape required to create row major tensor".to_string(),
- ))
- }
- }
-
- /// Creates a new Tensor using column major memory layout
- pub fn new_column_major(
- buffer: Buffer,
- shape: Option<Vec<usize>>,
- names: Option<Vec<&'a str>>,
- ) -> Result<Self> {
- if let Some(ref s) = shape {
- let strides = Some(compute_column_major_strides::<T>(&s)?);
-
- Self::try_new(buffer, shape, strides, names)
- } else {
- Err(ArrowError::InvalidArgumentError(
- "shape required to create column major tensor".to_string(),
- ))
- }
- }
-
- /// The data type of the `Tensor`
- pub fn data_type(&self) -> &DataType {
- &self.data_type
- }
-
- /// The sizes of the dimensions
- pub fn shape(&self) -> Option<&Vec<usize>> {
- self.shape.as_ref()
- }
-
- /// Returns a reference to the underlying `Buffer`
- pub fn data(&self) -> &Buffer {
- &self.buffer
- }
-
- /// The number of bytes between elements in each dimension
- pub fn strides(&self) -> Option<&Vec<usize>> {
- self.strides.as_ref()
- }
-
- /// The names of the dimensions
- pub fn names(&self) -> Option<&Vec<&'a str>> {
- self.names.as_ref()
- }
-
- /// The number of dimensions
- pub fn ndim(&self) -> usize {
- match &self.shape {
- None => 0,
- Some(v) => v.len(),
- }
- }
-
- /// The name of dimension i
- pub fn dim_name(&self, i: usize) -> Option<&'a str> {
- self.names.as_ref().map(|ref names| names[i])
- }
-
- /// The total number of elements in the `Tensor`
- pub fn size(&self) -> usize {
- match self.shape {
- None => 0,
- Some(ref s) => s.iter().product(),
- }
- }
-
- /// Indicates if the data is laid out contiguously in memory
- pub fn is_contiguous(&self) -> Result<bool> {
- Ok(self.is_row_major()? || self.is_column_major()?)
- }
-
- /// Indicates if the memory layout row major
- pub fn is_row_major(&self) -> Result<bool> {
- match self.shape {
- None => Ok(false),
- Some(ref s) => Ok(Some(compute_row_major_strides::<T>(s)?) == self.strides),
- }
- }
-
- /// Indicates if the memory layout column major
- pub fn is_column_major(&self) -> Result<bool> {
- match self.shape {
- None => Ok(false),
- Some(ref s) => {
- Ok(Some(compute_column_major_strides::<T>(s)?) == self.strides)
- }
- }
- }
-}
-
-#[cfg(test)]
-mod tests {
- use super::*;
-
- use crate::array::*;
- use crate::buffer::Buffer;
-
- #[test]
- fn test_compute_row_major_strides() {
- assert_eq!(
- vec![48_usize, 8],
- compute_row_major_strides::<Int64Type>(&[4_usize, 6]).unwrap()
- );
- assert_eq!(
- vec![24_usize, 4],
- compute_row_major_strides::<Int32Type>(&[4_usize, 6]).unwrap()
- );
- assert_eq!(
- vec![6_usize, 1],
- compute_row_major_strides::<Int8Type>(&[4_usize, 6]).unwrap()
- );
- }
-
- #[test]
- fn test_compute_column_major_strides() {
- assert_eq!(
- vec![8_usize, 32],
- compute_column_major_strides::<Int64Type>(&[4_usize, 6]).unwrap()
- );
- assert_eq!(
- vec![4_usize, 16],
- compute_column_major_strides::<Int32Type>(&[4_usize, 6]).unwrap()
- );
- assert_eq!(
- vec![1_usize, 4],
- compute_column_major_strides::<Int8Type>(&[4_usize, 6]).unwrap()
- );
- }
-
- #[test]
- fn test_zero_dim() {
- let buf = Buffer::from(&[1]);
- let tensor = UInt8Tensor::try_new(buf, None, None, None).unwrap();
- assert_eq!(0, tensor.size());
- assert_eq!(None, tensor.shape());
- assert_eq!(None, tensor.names());
- assert_eq!(0, tensor.ndim());
- assert_eq!(false, tensor.is_row_major().unwrap());
- assert_eq!(false, tensor.is_column_major().unwrap());
- assert_eq!(false, tensor.is_contiguous().unwrap());
-
- let buf = Buffer::from(&[1, 2, 2, 2]);
- let tensor = Int32Tensor::try_new(buf, None, None, None).unwrap();
- assert_eq!(0, tensor.size());
- assert_eq!(None, tensor.shape());
- assert_eq!(None, tensor.names());
- assert_eq!(0, tensor.ndim());
- assert_eq!(false, tensor.is_row_major().unwrap());
- assert_eq!(false, tensor.is_column_major().unwrap());
- assert_eq!(false, tensor.is_contiguous().unwrap());
- }
-
- #[test]
- fn test_tensor() {
- let mut builder = Int32BufferBuilder::new(16);
- for i in 0..16 {
- builder.append(i);
- }
- let buf = builder.finish();
- let tensor = Int32Tensor::try_new(buf, Some(vec![2, 8]), None, None).unwrap();
- assert_eq!(16, tensor.size());
- assert_eq!(Some(vec![2_usize, 8]).as_ref(), tensor.shape());
- assert_eq!(Some(vec![32_usize, 4]).as_ref(), tensor.strides());
- assert_eq!(2, tensor.ndim());
- assert_eq!(None, tensor.names());
- }
-
- #[test]
- fn test_new_row_major() {
- let mut builder = Int32BufferBuilder::new(16);
- for i in 0..16 {
- builder.append(i);
- }
- let buf = builder.finish();
- let tensor = Int32Tensor::new_row_major(buf, Some(vec![2, 8]), None).unwrap();
- assert_eq!(16, tensor.size());
- assert_eq!(Some(vec![2_usize, 8]).as_ref(), tensor.shape());
- assert_eq!(Some(vec![32_usize, 4]).as_ref(), tensor.strides());
- assert_eq!(None, tensor.names());
- assert_eq!(2, tensor.ndim());
- assert_eq!(true, tensor.is_row_major().unwrap());
- assert_eq!(false, tensor.is_column_major().unwrap());
- assert_eq!(true, tensor.is_contiguous().unwrap());
- }
-
- #[test]
- fn test_new_column_major() {
- let mut builder = Int32BufferBuilder::new(16);
- for i in 0..16 {
- builder.append(i);
- }
- let buf = builder.finish();
- let tensor = Int32Tensor::new_column_major(buf, Some(vec![2, 8]), None).unwrap();
- assert_eq!(16, tensor.size());
- assert_eq!(Some(vec![2_usize, 8]).as_ref(), tensor.shape());
- assert_eq!(Some(vec![4_usize, 8]).as_ref(), tensor.strides());
- assert_eq!(None, tensor.names());
- assert_eq!(2, tensor.ndim());
- assert_eq!(false, tensor.is_row_major().unwrap());
- assert_eq!(true, tensor.is_column_major().unwrap());
- assert_eq!(true, tensor.is_contiguous().unwrap());
- }
-
- #[test]
- fn test_with_names() {
- let mut builder = Int64BufferBuilder::new(8);
- for i in 0..8 {
- builder.append(i);
- }
- let buf = builder.finish();
- let names = vec!["Dim 1", "Dim 2"];
- let tensor =
- Int64Tensor::new_column_major(buf, Some(vec![2, 4]), Some(names)).unwrap();
- assert_eq!(8, tensor.size());
- assert_eq!(Some(vec![2_usize, 4]).as_ref(), tensor.shape());
- assert_eq!(Some(vec![8_usize, 16]).as_ref(), tensor.strides());
- assert_eq!("Dim 1", tensor.dim_name(0).unwrap());
- assert_eq!("Dim 2", tensor.dim_name(1).unwrap());
- assert_eq!(2, tensor.ndim());
- assert_eq!(false, tensor.is_row_major().unwrap());
- assert_eq!(true, tensor.is_column_major().unwrap());
- assert_eq!(true, tensor.is_contiguous().unwrap());
- }
-
- #[test]
- fn test_inconsistent_strides() {
- let mut builder = Int32BufferBuilder::new(16);
- for i in 0..16 {
- builder.append(i);
- }
- let buf = builder.finish();
-
- let result =
- Int32Tensor::try_new(buf, Some(vec![2, 8]), Some(vec![2, 8, 1]), None);
-
- if result.is_ok() {
- panic!("shape and stride dimensions are different")
- }
- }
-
- #[test]
- fn test_inconsistent_names() {
- let mut builder = Int32BufferBuilder::new(16);
- for i in 0..16 {
- builder.append(i);
- }
- let buf = builder.finish();
-
- let result = Int32Tensor::try_new(
- buf,
- Some(vec![2, 8]),
- Some(vec![4, 8]),
- Some(vec!["1", "2", "3"]),
- );
-
- if result.is_ok() {
- panic!("dimensions and names have different shape")
- }
- }
-
- #[test]
- fn test_incorrect_shape() {
- let mut builder = Int32BufferBuilder::new(16);
- for i in 0..16 {
- builder.append(i);
- }
- let buf = builder.finish();
-
- let result = Int32Tensor::try_new(buf, Some(vec![2, 6]), None, None);
-
- if result.is_ok() {
- panic!("number of elements does not match for the shape")
- }
- }
-
- #[test]
- fn test_incorrect_stride() {
- let mut builder = Int32BufferBuilder::new(16);
- for i in 0..16 {
- builder.append(i);
- }
- let buf = builder.finish();
-
- let result = Int32Tensor::try_new(buf, Some(vec![2, 8]), Some(vec![30, 4]), None);
-
- if result.is_ok() {
- panic!("the input stride does not match the selected shape")
- }
- }
-}
diff --git a/arrow/src/util/bench_util.rs b/arrow/src/util/bench_util.rs
deleted file mode 100644
index fd0ece8..0000000
--- a/arrow/src/util/bench_util.rs
+++ /dev/null
@@ -1,155 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-//! Utils to make benchmarking easier
-
-use crate::array::*;
-use crate::datatypes::*;
-use crate::util::test_util::seedable_rng;
-use rand::Rng;
-use rand::SeedableRng;
-use rand::{
- distributions::{Alphanumeric, Distribution, Standard},
- prelude::StdRng,
-};
-
-/// Creates an random (but fixed-seeded) array of a given size and null density
-pub fn create_primitive_array<T>(size: usize, null_density: f32) -> PrimitiveArray<T>
-where
- T: ArrowPrimitiveType,
- Standard: Distribution<T::Native>,
-{
- let mut rng = seedable_rng();
-
- (0..size)
- .map(|_| {
- if rng.gen::<f32>() < null_density {
- None
- } else {
- Some(rng.gen())
- }
- })
- .collect()
-}
-
-pub fn create_primitive_array_with_seed<T>(
- size: usize,
- null_density: f32,
- seed: u64,
-) -> PrimitiveArray<T>
-where
- T: ArrowPrimitiveType,
- Standard: Distribution<T::Native>,
-{
- let mut rng = StdRng::seed_from_u64(seed);
-
- (0..size)
- .map(|_| {
- if rng.gen::<f32>() < null_density {
- None
- } else {
- Some(rng.gen())
- }
- })
- .collect()
-}
-
-/// Creates an random (but fixed-seeded) array of a given size and null density
-pub fn create_boolean_array(
- size: usize,
- null_density: f32,
- true_density: f32,
-) -> BooleanArray
-where
- Standard: Distribution<bool>,
-{
- let mut rng = seedable_rng();
- (0..size)
- .map(|_| {
- if rng.gen::<f32>() < null_density {
- None
- } else {
- let value = rng.gen::<f32>() < true_density;
- Some(value)
- }
- })
- .collect()
-}
-
-/// Creates an random (but fixed-seeded) array of a given size and null density
-pub fn create_string_array<Offset: StringOffsetSizeTrait>(
- size: usize,
- null_density: f32,
-) -> GenericStringArray<Offset> {
- let rng = &mut seedable_rng();
-
- (0..size)
- .map(|_| {
- if rng.gen::<f32>() < null_density {
- None
- } else {
- let value = rng.sample_iter(&Alphanumeric).take(4).collect::<String>();
- Some(value)
- }
- })
- .collect()
-}
-
-/// Creates an random (but fixed-seeded) binary array of a given size and null density
-pub fn create_binary_array<Offset: BinaryOffsetSizeTrait>(
- size: usize,
- null_density: f32,
-) -> GenericBinaryArray<Offset> {
- let rng = &mut seedable_rng();
- let range_rng = &mut seedable_rng();
-
- (0..size)
- .map(|_| {
- if rng.gen::<f32>() < null_density {
- None
- } else {
- let value = rng
- .sample_iter::<u8, _>(Standard)
- .take(range_rng.gen_range(0, 8))
- .collect::<Vec<u8>>();
- Some(value)
- }
- })
- .collect()
-}
-
-/// Creates an random (but fixed-seeded) array of a given size and null density
-pub fn create_fsb_array(
- size: usize,
- null_density: f32,
- value_len: usize,
-) -> FixedSizeBinaryArray {
- let rng = &mut seedable_rng();
-
- FixedSizeBinaryArray::try_from_sparse_iter((0..size).map(|_| {
- if rng.gen::<f32>() < null_density {
- None
- } else {
- let value = rng
- .sample_iter::<u8, _>(Standard)
- .take(value_len)
- .collect::<Vec<u8>>();
- Some(value)
- }
- }))
- .unwrap()
-}
diff --git a/arrow/src/util/bit_chunk_iterator.rs b/arrow/src/util/bit_chunk_iterator.rs
deleted file mode 100644
index ea9280c..0000000
--- a/arrow/src/util/bit_chunk_iterator.rs
+++ /dev/null
@@ -1,275 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-use crate::util::bit_util::ceil;
-use std::fmt::Debug;
-
-#[derive(Debug)]
-pub struct BitChunks<'a> {
- buffer: &'a [u8],
- /// offset inside a byte, guaranteed to be between 0 and 7 (inclusive)
- bit_offset: usize,
- /// number of complete u64 chunks
- chunk_len: usize,
- /// number of remaining bits, guaranteed to be between 0 and 63 (inclusive)
- remainder_len: usize,
-}
-
-impl<'a> BitChunks<'a> {
- pub fn new(buffer: &'a [u8], offset: usize, len: usize) -> Self {
- assert!(ceil(offset + len, 8) <= buffer.len() * 8);
-
- let byte_offset = offset / 8;
- let bit_offset = offset % 8;
-
- // number of complete u64 chunks
- let chunk_len = len / 64;
- // number of remaining bits
- let remainder_len = len % 64;
-
- BitChunks::<'a> {
- buffer: &buffer[byte_offset..],
- bit_offset,
- chunk_len,
- remainder_len,
- }
- }
-}
-
-#[derive(Debug)]
-pub struct BitChunkIterator<'a> {
- buffer: &'a [u8],
- bit_offset: usize,
- chunk_len: usize,
- index: usize,
-}
-
-impl<'a> BitChunks<'a> {
- /// Returns the number of remaining bits, guaranteed to be between 0 and 63 (inclusive)
- #[inline]
- pub const fn remainder_len(&self) -> usize {
- self.remainder_len
- }
-
- /// Returns the number of chunks
- #[inline]
- pub const fn chunk_len(&self) -> usize {
- self.chunk_len
- }
-
- /// Returns the bitmask of remaining bits
- #[inline]
- pub fn remainder_bits(&self) -> u64 {
- let bit_len = self.remainder_len;
- if bit_len == 0 {
- 0
- } else {
- let bit_offset = self.bit_offset;
- // number of bytes to read
- // might be one more than sizeof(u64) if the offset is in the middle of a byte
- let byte_len = ceil(bit_len + bit_offset, 8);
- // pointer to remainder bytes after all complete chunks
- let base = unsafe {
- self.buffer
- .as_ptr()
- .add(self.chunk_len * std::mem::size_of::<u64>())
- };
-
- let mut bits = unsafe { std::ptr::read(base) } as u64 >> bit_offset;
- for i in 1..byte_len {
- let byte = unsafe { std::ptr::read(base.add(i)) };
- bits |= (byte as u64) << (i * 8 - bit_offset);
- }
-
- bits & ((1 << bit_len) - 1)
- }
- }
-
- /// Returns an iterator over chunks of 64 bits represented as an u64
- #[inline]
- pub const fn iter(&self) -> BitChunkIterator<'a> {
- BitChunkIterator::<'a> {
- buffer: self.buffer,
- bit_offset: self.bit_offset,
- chunk_len: self.chunk_len,
- index: 0,
- }
- }
-}
-
-impl<'a> IntoIterator for BitChunks<'a> {
- type Item = u64;
- type IntoIter = BitChunkIterator<'a>;
-
- fn into_iter(self) -> Self::IntoIter {
- self.iter()
- }
-}
-
-impl Iterator for BitChunkIterator<'_> {
- type Item = u64;
-
- #[inline]
- fn next(&mut self) -> Option<u64> {
- let index = self.index;
- if index >= self.chunk_len {
- return None;
- }
-
- // cast to *const u64 should be fine since we are using read_unaligned below
- #[allow(clippy::cast_ptr_alignment)]
- let raw_data = self.buffer.as_ptr() as *const u64;
-
- // bit-packed buffers are stored starting with the least-significant byte first
- // so when reading as u64 on a big-endian machine, the bytes need to be swapped
- let current = unsafe { std::ptr::read_unaligned(raw_data.add(index)).to_le() };
-
- let bit_offset = self.bit_offset;
-
- let combined = if bit_offset == 0 {
- current
- } else {
- // the constructor ensures that bit_offset is in 0..8
- // that means we need to read at most one additional byte to fill in the high bits
- let next = unsafe {
- std::ptr::read_unaligned(raw_data.add(index + 1) as *const u8) as u64
- };
-
- (current >> bit_offset) | (next << (64 - bit_offset))
- };
-
- self.index = index + 1;
-
- Some(combined)
- }
-
- #[inline]
- fn size_hint(&self) -> (usize, Option<usize>) {
- (
- self.chunk_len - self.index,
- Some(self.chunk_len - self.index),
- )
- }
-}
-
-impl ExactSizeIterator for BitChunkIterator<'_> {
- #[inline]
- fn len(&self) -> usize {
- self.chunk_len - self.index
- }
-}
-
-#[cfg(test)]
-mod tests {
- use crate::buffer::Buffer;
-
- #[test]
- fn test_iter_aligned() {
- let input: &[u8] = &[0, 1, 2, 3, 4, 5, 6, 7];
- let buffer: Buffer = Buffer::from(input);
-
- let bitchunks = buffer.bit_chunks(0, 64);
- let result = bitchunks.into_iter().collect::<Vec<_>>();
-
- assert_eq!(vec![0x0706050403020100], result);
- }
-
- #[test]
- fn test_iter_unaligned() {
- let input: &[u8] = &[
- 0b00000000, 0b00000001, 0b00000010, 0b00000100, 0b00001000, 0b00010000,
- 0b00100000, 0b01000000, 0b11111111,
- ];
- let buffer: Buffer = Buffer::from(input);
-
- let bitchunks = buffer.bit_chunks(4, 64);
-
- assert_eq!(0, bitchunks.remainder_len());
- assert_eq!(0, bitchunks.remainder_bits());
-
- let result = bitchunks.into_iter().collect::<Vec<_>>();
-
- assert_eq!(
- vec![0b1111010000000010000000010000000010000000010000000010000000010000],
- result
- );
- }
-
- #[test]
- fn test_iter_unaligned_remainder_1_byte() {
- let input: &[u8] = &[
- 0b00000000, 0b00000001, 0b00000010, 0b00000100, 0b00001000, 0b00010000,
- 0b00100000, 0b01000000, 0b11111111,
- ];
- let buffer: Buffer = Buffer::from(input);
-
- let bitchunks = buffer.bit_chunks(4, 66);
-
- assert_eq!(2, bitchunks.remainder_len());
- assert_eq!(0b00000011, bitchunks.remainder_bits());
-
- let result = bitchunks.into_iter().collect::<Vec<_>>();
-
- assert_eq!(
- vec![0b1111010000000010000000010000000010000000010000000010000000010000],
- result
- );
- }
-
- #[test]
- fn test_iter_unaligned_remainder_bits_across_bytes() {
- let input: &[u8] = &[0b00111111, 0b11111100];
- let buffer: Buffer = Buffer::from(input);
-
- // remainder contains bits from both bytes
- // result should be the highest 2 bits from first byte followed by lowest 5 bits of second bytes
- let bitchunks = buffer.bit_chunks(6, 7);
-
- assert_eq!(7, bitchunks.remainder_len());
- assert_eq!(0b1110000, bitchunks.remainder_bits());
- }
-
- #[test]
- fn test_iter_unaligned_remainder_bits_large() {
- let input: &[u8] = &[
- 0b11111111, 0b00000000, 0b11111111, 0b00000000, 0b11111111, 0b00000000,
- 0b11111111, 0b00000000, 0b11111111,
- ];
- let buffer: Buffer = Buffer::from(input);
-
- let bitchunks = buffer.bit_chunks(2, 63);
-
- assert_eq!(63, bitchunks.remainder_len());
- assert_eq!(
- 0b100_0000_0011_1111_1100_0000_0011_1111_1100_0000_0011_1111_1100_0000_0011_1111,
- bitchunks.remainder_bits()
- );
- }
-
- #[test]
- fn test_iter_remainder_out_of_bounds() {
- // allocating a full page should trigger a fault when reading out of bounds
- const ALLOC_SIZE: usize = 4 * 1024;
- let input = vec![0xFF_u8; ALLOC_SIZE];
-
- let buffer: Buffer = Buffer::from(input);
-
- let bitchunks = buffer.bit_chunks(57, ALLOC_SIZE * 8 - 57);
-
- assert_eq!(u64::MAX, bitchunks.iter().last().unwrap());
- assert_eq!(0x7F, bitchunks.remainder_bits());
- }
-}
diff --git a/arrow/src/util/bit_util.rs b/arrow/src/util/bit_util.rs
deleted file mode 100644
index 7298eb3..0000000
--- a/arrow/src/util/bit_util.rs
+++ /dev/null
@@ -1,322 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-//! Utils for working with bits
-
-#[cfg(feature = "simd")]
-use packed_simd::u8x64;
-
-const BIT_MASK: [u8; 8] = [1, 2, 4, 8, 16, 32, 64, 128];
-const UNSET_BIT_MASK: [u8; 8] = [
- 255 - 1,
- 255 - 2,
- 255 - 4,
- 255 - 8,
- 255 - 16,
- 255 - 32,
- 255 - 64,
- 255 - 128,
-];
-
-/// Returns the nearest number that is `>=` than `num` and is a multiple of 64
-#[inline]
-pub fn round_upto_multiple_of_64(num: usize) -> usize {
- round_upto_power_of_2(num, 64)
-}
-
-/// Returns the nearest multiple of `factor` that is `>=` than `num`. Here `factor` must
-/// be a power of 2.
-pub fn round_upto_power_of_2(num: usize, factor: usize) -> usize {
- debug_assert!(factor > 0 && (factor & (factor - 1)) == 0);
- (num + (factor - 1)) & !(factor - 1)
-}
-
-/// Returns whether bit at position `i` in `data` is set or not
-#[inline]
-pub fn get_bit(data: &[u8], i: usize) -> bool {
- (data[i >> 3] & BIT_MASK[i & 7]) != 0
-}
-
-/// Returns whether bit at position `i` in `data` is set or not.
-///
-/// # Safety
-///
-/// Note this doesn't do any bound checking, for performance reason. The caller is
-/// responsible to guarantee that `i` is within bounds.
-#[inline]
-pub unsafe fn get_bit_raw(data: *const u8, i: usize) -> bool {
- (*data.add(i >> 3) & BIT_MASK[i & 7]) != 0
-}
-
-/// Sets bit at position `i` for `data` to 1
-#[inline]
-pub fn set_bit(data: &mut [u8], i: usize) {
- data[i >> 3] |= BIT_MASK[i & 7];
-}
-
-/// Sets bit at position `i` for `data`
-///
-/// # Safety
-///
-/// Note this doesn't do any bound checking, for performance reason. The caller is
-/// responsible to guarantee that `i` is within bounds.
-#[inline]
-pub unsafe fn set_bit_raw(data: *mut u8, i: usize) {
- *data.add(i >> 3) |= BIT_MASK[i & 7];
-}
-
-/// Sets bit at position `i` for `data` to 0
-#[inline]
-pub fn unset_bit(data: &mut [u8], i: usize) {
- data[i >> 3] &= UNSET_BIT_MASK[i & 7];
-}
-
-/// Sets bit at position `i` for `data` to 0
-///
-/// # Safety
-///
-/// Note this doesn't do any bound checking, for performance reason. The caller is
-/// responsible to guarantee that `i` is within bounds.
-#[inline]
-pub unsafe fn unset_bit_raw(data: *mut u8, i: usize) {
- *data.add(i >> 3) &= UNSET_BIT_MASK[i & 7];
-}
-
-/// Returns the ceil of `value`/`divisor`
-#[inline]
-pub fn ceil(value: usize, divisor: usize) -> usize {
- let (quot, rem) = (value / divisor, value % divisor);
- if rem > 0 && divisor > 0 {
- quot + 1
- } else {
- quot
- }
-}
-
-/// Performs SIMD bitwise binary operations.
-///
-/// # Safety
-///
-/// Note that each slice should be 64 bytes and it is the callers responsibility to ensure
-/// that this is the case. If passed slices larger than 64 bytes the operation will only
-/// be performed on the first 64 bytes. Slices less than 64 bytes will panic.
-#[cfg(feature = "simd")]
-pub unsafe fn bitwise_bin_op_simd<F>(left: &[u8], right: &[u8], result: &mut [u8], op: F)
-where
- F: Fn(u8x64, u8x64) -> u8x64,
-{
- let left_simd = u8x64::from_slice_unaligned_unchecked(left);
- let right_simd = u8x64::from_slice_unaligned_unchecked(right);
- let simd_result = op(left_simd, right_simd);
- simd_result.write_to_slice_unaligned_unchecked(result);
-}
-
-#[cfg(test)]
-mod tests {
- use std::collections::HashSet;
-
- use super::*;
- use crate::util::test_util::seedable_rng;
- use rand::Rng;
-
- #[test]
- fn test_round_upto_multiple_of_64() {
- assert_eq!(0, round_upto_multiple_of_64(0));
- assert_eq!(64, round_upto_multiple_of_64(1));
- assert_eq!(64, round_upto_multiple_of_64(63));
- assert_eq!(64, round_upto_multiple_of_64(64));
- assert_eq!(128, round_upto_multiple_of_64(65));
- assert_eq!(192, round_upto_multiple_of_64(129));
- }
-
- #[test]
- fn test_get_bit() {
- // 00001101
- assert_eq!(true, get_bit(&[0b00001101], 0));
- assert_eq!(false, get_bit(&[0b00001101], 1));
- assert_eq!(true, get_bit(&[0b00001101], 2));
- assert_eq!(true, get_bit(&[0b00001101], 3));
-
- // 01001001 01010010
- assert_eq!(true, get_bit(&[0b01001001, 0b01010010], 0));
- assert_eq!(false, get_bit(&[0b01001001, 0b01010010], 1));
- assert_eq!(false, get_bit(&[0b01001001, 0b01010010], 2));
- assert_eq!(true, get_bit(&[0b01001001, 0b01010010], 3));
- assert_eq!(false, get_bit(&[0b01001001, 0b01010010], 4));
- assert_eq!(false, get_bit(&[0b01001001, 0b01010010], 5));
- assert_eq!(true, get_bit(&[0b01001001, 0b01010010], 6));
- assert_eq!(false, get_bit(&[0b01001001, 0b01010010], 7));
- assert_eq!(false, get_bit(&[0b01001001, 0b01010010], 8));
- assert_eq!(true, get_bit(&[0b01001001, 0b01010010], 9));
- assert_eq!(false, get_bit(&[0b01001001, 0b01010010], 10));
- assert_eq!(false, get_bit(&[0b01001001, 0b01010010], 11));
- assert_eq!(true, get_bit(&[0b01001001, 0b01010010], 12));
- assert_eq!(false, get_bit(&[0b01001001, 0b01010010], 13));
- assert_eq!(true, get_bit(&[0b01001001, 0b01010010], 14));
- assert_eq!(false, get_bit(&[0b01001001, 0b01010010], 15));
- }
-
- #[test]
- fn test_get_bit_raw() {
- const NUM_BYTE: usize = 10;
- let mut buf = vec![0; NUM_BYTE];
- let mut expected = vec![];
- let mut rng = seedable_rng();
- for i in 0..8 * NUM_BYTE {
- let b = rng.gen_bool(0.5);
- expected.push(b);
- if b {
- set_bit(&mut buf[..], i)
- }
- }
-
- let raw_ptr = buf.as_ptr();
- for (i, b) in expected.iter().enumerate() {
- unsafe {
- assert_eq!(*b, get_bit_raw(raw_ptr, i));
- }
- }
- }
-
- #[test]
- fn test_set_bit() {
- let mut b = [0b00000010];
- set_bit(&mut b, 0);
- assert_eq!([0b00000011], b);
- set_bit(&mut b, 1);
- assert_eq!([0b00000011], b);
- set_bit(&mut b, 7);
- assert_eq!([0b10000011], b);
- }
-
- #[test]
- fn test_unset_bit() {
- let mut b = [0b11111101];
- unset_bit(&mut b, 0);
- assert_eq!([0b11111100], b);
- unset_bit(&mut b, 1);
- assert_eq!([0b11111100], b);
- unset_bit(&mut b, 7);
- assert_eq!([0b01111100], b);
- }
-
- #[test]
- fn test_set_bit_raw() {
- const NUM_BYTE: usize = 10;
- let mut buf = vec![0; NUM_BYTE];
- let mut expected = vec![];
- let mut rng = seedable_rng();
- for i in 0..8 * NUM_BYTE {
- let b = rng.gen_bool(0.5);
- expected.push(b);
- if b {
- unsafe {
- set_bit_raw(buf.as_mut_ptr(), i);
- }
- }
- }
-
- let raw_ptr = buf.as_ptr();
- for (i, b) in expected.iter().enumerate() {
- unsafe {
- assert_eq!(*b, get_bit_raw(raw_ptr, i));
- }
- }
- }
-
- #[test]
- fn test_unset_bit_raw() {
- const NUM_BYTE: usize = 10;
- let mut buf = vec![255; NUM_BYTE];
- let mut expected = vec![];
- let mut rng = seedable_rng();
- for i in 0..8 * NUM_BYTE {
- let b = rng.gen_bool(0.5);
- expected.push(b);
- if !b {
- unsafe {
- unset_bit_raw(buf.as_mut_ptr(), i);
- }
- }
- }
-
- let raw_ptr = buf.as_ptr();
- for (i, b) in expected.iter().enumerate() {
- unsafe {
- assert_eq!(*b, get_bit_raw(raw_ptr, i));
- }
- }
- }
-
- #[test]
- fn test_get_set_bit_roundtrip() {
- const NUM_BYTES: usize = 10;
- const NUM_SETS: usize = 10;
-
- let mut buffer: [u8; NUM_BYTES * 8] = [0; NUM_BYTES * 8];
- let mut v = HashSet::new();
- let mut rng = seedable_rng();
- for _ in 0..NUM_SETS {
- let offset = rng.gen_range(0, 8 * NUM_BYTES);
- v.insert(offset);
- set_bit(&mut buffer[..], offset);
- }
- for i in 0..NUM_BYTES * 8 {
- assert_eq!(v.contains(&i), get_bit(&buffer[..], i));
- }
- }
-
- #[test]
- #[cfg(all(any(target_arch = "x86", target_arch = "x86_64")))]
- fn test_ceil() {
- assert_eq!(ceil(0, 1), 0);
- assert_eq!(ceil(1, 1), 1);
- assert_eq!(ceil(1, 2), 1);
- assert_eq!(ceil(1, 8), 1);
- assert_eq!(ceil(7, 8), 1);
- assert_eq!(ceil(8, 8), 1);
- assert_eq!(ceil(9, 8), 2);
- assert_eq!(ceil(9, 9), 1);
- assert_eq!(ceil(10000000000, 10), 1000000000);
- assert_eq!(ceil(10, 10000000000), 1);
- assert_eq!(ceil(10000000000, 1000000000), 10);
- }
-
- #[test]
- #[cfg(feature = "simd")]
- fn test_bitwise_and_simd() {
- let buf1 = [0b00110011u8; 64];
- let buf2 = [0b11110000u8; 64];
- let mut buf3 = [0b00000000; 64];
- unsafe { bitwise_bin_op_simd(&buf1, &buf2, &mut buf3, |a, b| a & b) };
- for i in buf3.iter() {
- assert_eq!(&0b00110000u8, i);
- }
- }
-
- #[test]
- #[cfg(feature = "simd")]
- fn test_bitwise_or_simd() {
- let buf1 = [0b00110011u8; 64];
- let buf2 = [0b11110000u8; 64];
- let mut buf3 = [0b00000000; 64];
- unsafe { bitwise_bin_op_simd(&buf1, &buf2, &mut buf3, |a, b| a | b) };
- for i in buf3.iter() {
- assert_eq!(&0b11110011u8, i);
- }
- }
-}
diff --git a/arrow/src/util/data_gen.rs b/arrow/src/util/data_gen.rs
deleted file mode 100644
index cd1f25e..0000000
--- a/arrow/src/util/data_gen.rs
+++ /dev/null
@@ -1,347 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-//! Utilities to generate random arrays and batches
-
-use std::{convert::TryFrom, sync::Arc};
-
-use rand::{distributions::uniform::SampleUniform, Rng};
-
-use crate::error::{ArrowError, Result};
-use crate::record_batch::{RecordBatch, RecordBatchOptions};
-use crate::{array::*, datatypes::SchemaRef};
-use crate::{
- buffer::{Buffer, MutableBuffer},
- datatypes::*,
-};
-
-use super::{bench_util::*, bit_util, test_util::seedable_rng};
-
-/// Create a random [RecordBatch] from a schema
-pub fn create_random_batch(
- schema: SchemaRef,
- size: usize,
- null_density: f32,
- true_density: f32,
-) -> Result<RecordBatch> {
- let columns = schema
- .fields()
- .iter()
- .map(|field| create_random_array(field, size, null_density, true_density))
- .collect::<Result<Vec<ArrayRef>>>()?;
-
- RecordBatch::try_new_with_options(
- schema,
- columns,
- &RecordBatchOptions {
- match_field_names: false,
- },
- )
-}
-
-/// Create a random [ArrayRef] from a [DataType] with a length,
-/// null density and true density (for [BooleanArray]).
-pub fn create_random_array(
- field: &Field,
- size: usize,
- null_density: f32,
- true_density: f32,
-) -> Result<ArrayRef> {
- // Override null density with 0.0 if the array is non-nullable
- let null_density = match field.is_nullable() {
- true => null_density,
- false => 0.0,
- };
- use DataType::*;
- Ok(match field.data_type() {
- Null => Arc::new(NullArray::new(size)) as ArrayRef,
- Boolean => Arc::new(create_boolean_array(size, null_density, true_density)),
- Int8 => Arc::new(create_primitive_array::<Int8Type>(size, null_density)),
- Int16 => Arc::new(create_primitive_array::<Int16Type>(size, null_density)),
- Int32 => Arc::new(create_primitive_array::<Int32Type>(size, null_density)),
- Int64 => Arc::new(create_primitive_array::<Int64Type>(size, null_density)),
- UInt8 => Arc::new(create_primitive_array::<UInt8Type>(size, null_density)),
- UInt16 => Arc::new(create_primitive_array::<UInt16Type>(size, null_density)),
- UInt32 => Arc::new(create_primitive_array::<UInt32Type>(size, null_density)),
- UInt64 => Arc::new(create_primitive_array::<UInt64Type>(size, null_density)),
- Float16 => {
- return Err(ArrowError::NotYetImplemented(
- "Float16 is not implememted".to_string(),
- ))
- }
- Float32 => Arc::new(create_primitive_array::<Float32Type>(size, null_density)),
- Float64 => Arc::new(create_primitive_array::<Float64Type>(size, null_density)),
- Timestamp(_, _) => {
- let int64_array =
- Arc::new(create_primitive_array::<Int64Type>(size, null_density))
- as ArrayRef;
- return crate::compute::cast(&int64_array, field.data_type());
- }
- Date32 => Arc::new(create_primitive_array::<Date32Type>(size, null_density)),
- Date64 => Arc::new(create_primitive_array::<Date64Type>(size, null_density)),
- Time32(unit) => match unit {
- TimeUnit::Second => Arc::new(create_primitive_array::<Time32SecondType>(
- size,
- null_density,
- )) as ArrayRef,
- TimeUnit::Millisecond => Arc::new(create_primitive_array::<
- Time32MillisecondType,
- >(size, null_density)),
- _ => {
- return Err(ArrowError::InvalidArgumentError(format!(
- "Unsupported unit {:?} for Time32",
- unit
- )))
- }
- },
- Time64(unit) => match unit {
- TimeUnit::Microsecond => Arc::new(create_primitive_array::<
- Time64MicrosecondType,
- >(size, null_density)) as ArrayRef,
- TimeUnit::Nanosecond => Arc::new(create_primitive_array::<
- Time64NanosecondType,
- >(size, null_density)),
- _ => {
- return Err(ArrowError::InvalidArgumentError(format!(
- "Unsupported unit {:?} for Time64",
- unit
- )))
- }
- },
- Utf8 => Arc::new(create_string_array::<i32>(size, null_density)),
- LargeUtf8 => Arc::new(create_string_array::<i64>(size, null_density)),
- Binary => Arc::new(create_binary_array::<i32>(size, null_density)),
- LargeBinary => Arc::new(create_binary_array::<i64>(size, null_density)),
- FixedSizeBinary(len) => {
- Arc::new(create_fsb_array(size, null_density, *len as usize))
- }
- List(_) => create_random_list_array(field, size, null_density, true_density)?,
- LargeList(_) => {
- create_random_list_array(field, size, null_density, true_density)?
- }
- Struct(fields) => Arc::new(StructArray::try_from(
- fields
- .iter()
- .map(|struct_field| {
- create_random_array(struct_field, size, null_density, true_density)
- .map(|array_ref| (struct_field.name().as_str(), array_ref))
- })
- .collect::<Result<Vec<(&str, ArrayRef)>>>()?,
- )?),
- other => {
- return Err(ArrowError::NotYetImplemented(format!(
- "Generating random arrays not yet implemented for {:?}",
- other
- )))
- }
- })
-}
-
-#[inline]
-fn create_random_list_array(
- field: &Field,
- size: usize,
- null_density: f32,
- true_density: f32,
-) -> Result<ArrayRef> {
- // Override null density with 0.0 if the array is non-nullable
- let null_density = match field.is_nullable() {
- true => null_density,
- false => 0.0,
- };
- let list_field;
- let (offsets, child_len) = match field.data_type() {
- DataType::List(f) => {
- let (offsets, child_len) = create_random_offsets::<i32>(size, 0, 5);
- list_field = f;
- (Buffer::from(offsets.to_byte_slice()), child_len as usize)
- }
- DataType::LargeList(f) => {
- let (offsets, child_len) = create_random_offsets::<i64>(size, 0, 5);
- list_field = f;
- (Buffer::from(offsets.to_byte_slice()), child_len as usize)
- }
- _ => {
- return Err(ArrowError::InvalidArgumentError(format!(
- "Cannot create list array for field {:?}",
- field
- )))
- }
- };
-
- // Create list's child data
- let child_array =
- create_random_array(list_field, child_len as usize, null_density, true_density)?;
- let child_data = child_array.data();
- // Create list's null buffers, if it is nullable
- let null_buffer = match field.is_nullable() {
- true => Some(create_random_null_buffer(size, null_density)),
- false => None,
- };
- let list_data = ArrayData::new(
- field.data_type().clone(),
- size,
- None,
- null_buffer,
- 0,
- vec![offsets],
- vec![child_data.clone()],
- );
- Ok(make_array(list_data))
-}
-
-/// Generate random offsets for list arrays
-fn create_random_offsets<T: OffsetSizeTrait + SampleUniform>(
- size: usize,
- min: T,
- max: T,
-) -> (Vec<T>, T) {
- let rng = &mut seedable_rng();
-
- let mut current_offset = T::zero();
-
- let mut offsets = Vec::with_capacity(size + 1);
- offsets.push(current_offset);
-
- (0..size).for_each(|_| {
- current_offset += rng.gen_range(min, max);
- offsets.push(current_offset);
- });
-
- (offsets, current_offset)
-}
-
-fn create_random_null_buffer(size: usize, null_density: f32) -> Buffer {
- let mut rng = seedable_rng();
- let mut mut_buf = MutableBuffer::new_null(size);
- {
- let mut_slice = mut_buf.as_slice_mut();
- (0..size).for_each(|i| {
- if rng.gen::<f32>() >= null_density {
- bit_util::set_bit(mut_slice, i)
- }
- })
- };
- mut_buf.into()
-}
-
-#[cfg(test)]
-mod tests {
- use super::*;
-
- #[test]
- fn test_create_batch() {
- let size = 32;
- let fields = vec![Field::new("a", DataType::Int32, true)];
- let schema = Schema::new(fields);
- let schema_ref = Arc::new(schema);
- let batch = create_random_batch(schema_ref.clone(), size, 0.35, 0.7).unwrap();
-
- assert_eq!(batch.schema(), schema_ref);
- assert_eq!(batch.num_columns(), schema_ref.fields().len());
- for array in batch.columns() {
- assert_eq!(array.len(), size);
- }
- }
-
- #[test]
- fn test_create_batch_non_null() {
- let size = 32;
- let fields = vec![
- Field::new("a", DataType::Int32, false),
- Field::new(
- "b",
- DataType::List(Box::new(Field::new("item", DataType::LargeUtf8, true))),
- false,
- ),
- Field::new("a", DataType::Int32, false),
- ];
- let schema = Schema::new(fields);
- let schema_ref = Arc::new(schema);
- let batch = create_random_batch(schema_ref.clone(), size, 0.35, 0.7).unwrap();
-
- assert_eq!(batch.schema(), schema_ref);
- assert_eq!(batch.num_columns(), schema_ref.fields().len());
- for array in batch.columns() {
- assert_eq!(array.null_count(), 0);
- }
- // Test that the list's child values are non-null
- let b_array = batch.column(1);
- let list_array = b_array.as_any().downcast_ref::<ListArray>().unwrap();
- let child_array = make_array(list_array.data().child_data()[0].clone());
- assert_eq!(child_array.null_count(), 0);
- // There should be more values than the list, to show that it's a list
- assert!(child_array.len() > list_array.len());
- }
-
- #[test]
- fn test_create_struct_array() {
- let size = 32;
- let struct_fields = vec![
- Field::new("b", DataType::Boolean, true),
- Field::new(
- "c",
- DataType::LargeList(Box::new(Field::new(
- "item",
- DataType::List(Box::new(Field::new(
- "item",
- DataType::FixedSizeBinary(6),
- true,
- ))),
- false,
- ))),
- true,
- ),
- Field::new(
- "d",
- DataType::Struct(vec![
- Field::new("d_x", DataType::Int32, true),
- Field::new("d_y", DataType::Float32, false),
- Field::new("d_z", DataType::Binary, true),
- ]),
- true,
- ),
- ];
- let field = Field::new("struct", DataType::Struct(struct_fields), true);
- let array = create_random_array(&field, size, 0.2, 0.5).unwrap();
-
- assert_eq!(array.len(), 32);
- let struct_array = array.as_any().downcast_ref::<StructArray>().unwrap();
- assert_eq!(struct_array.columns().len(), 3);
-
- // Test that the nested list makes sense,
- // i.e. its children's values are more than the parent, to show repetition
- let col_c = struct_array.column_by_name("c").unwrap();
- let col_c = col_c.as_any().downcast_ref::<LargeListArray>().unwrap();
- assert_eq!(col_c.len(), size);
- let col_c_values = col_c.values();
- assert!(col_c_values.len() > size);
- // col_c_values should be a list
- let col_c_list = col_c_values.as_any().downcast_ref::<ListArray>().unwrap();
- // Its values should be FixedSizeBinary(6)
- let fsb = col_c_list.values();
- assert_eq!(fsb.data_type(), &DataType::FixedSizeBinary(6));
- assert!(fsb.len() > col_c_list.len());
-
- // Test nested struct
- let col_d = struct_array.column_by_name("d").unwrap();
- let col_d = col_d.as_any().downcast_ref::<StructArray>().unwrap();
- let col_d_y = col_d.column_by_name("d_y").unwrap();
- assert_eq!(col_d_y.data_type(), &DataType::Float32);
- assert_eq!(col_d_y.null_count(), 0);
- }
-}
diff --git a/arrow/src/util/display.rs b/arrow/src/util/display.rs
deleted file mode 100644
index bb75a3a..0000000
--- a/arrow/src/util/display.rs
+++ /dev/null
@@ -1,312 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-//! Functions for printing array values, as strings, for debugging
-//! purposes. See the `pretty` crate for additional functions for
-//! record batch pretty printing.
-
-use std::sync::Arc;
-
-use crate::array::Array;
-use crate::datatypes::{
- ArrowNativeType, ArrowPrimitiveType, DataType, Int16Type, Int32Type, Int64Type,
- Int8Type, TimeUnit, UInt16Type, UInt32Type, UInt64Type, UInt8Type,
-};
-use crate::{array, datatypes::IntervalUnit};
-
-use array::DictionaryArray;
-
-use crate::error::{ArrowError, Result};
-
-macro_rules! make_string {
- ($array_type:ty, $column: ident, $row: ident) => {{
- let array = $column.as_any().downcast_ref::<$array_type>().unwrap();
-
- let s = if array.is_null($row) {
- "".to_string()
- } else {
- array.value($row).to_string()
- };
-
- Ok(s)
- }};
-}
-
-macro_rules! make_string_interval_year_month {
- ($column: ident, $row: ident) => {{
- let array = $column
- .as_any()
- .downcast_ref::<array::IntervalYearMonthArray>()
- .unwrap();
-
- let s = if array.is_null($row) {
- "NULL".to_string()
- } else {
- let interval = array.value($row) as f64;
- let years = (interval / 12_f64).floor();
- let month = interval - (years * 12_f64);
-
- format!(
- "{} years {} mons 0 days 0 hours 0 mins 0.00 secs",
- years, month,
- )
- };
-
- Ok(s)
- }};
-}
-
-macro_rules! make_string_interval_day_time {
- ($column: ident, $row: ident) => {{
- let array = $column
- .as_any()
- .downcast_ref::<array::IntervalDayTimeArray>()
- .unwrap();
-
- let s = if array.is_null($row) {
- "NULL".to_string()
- } else {
- let value: u64 = array.value($row) as u64;
-
- let days_parts: i32 = ((value & 0xFFFFFFFF00000000) >> 32) as i32;
- let milliseconds_part: i32 = (value & 0xFFFFFFFF) as i32;
-
- let secs = milliseconds_part / 1000;
- let mins = secs / 60;
- let hours = mins / 60;
-
- let secs = secs - (mins * 60);
- let mins = mins - (hours * 60);
-
- format!(
- "0 years 0 mons {} days {} hours {} mins {}.{:02} secs",
- days_parts,
- hours,
- mins,
- secs,
- (milliseconds_part % 1000),
- )
- };
-
- Ok(s)
- }};
-}
-
-macro_rules! make_string_date {
- ($array_type:ty, $column: ident, $row: ident) => {{
- let array = $column.as_any().downcast_ref::<$array_type>().unwrap();
-
- let s = if array.is_null($row) {
- "".to_string()
- } else {
- array
- .value_as_date($row)
- .map(|d| d.to_string())
- .unwrap_or_else(|| "ERROR CONVERTING DATE".to_string())
- };
-
- Ok(s)
- }};
-}
-
-macro_rules! make_string_time {
- ($array_type:ty, $column: ident, $row: ident) => {{
- let array = $column.as_any().downcast_ref::<$array_type>().unwrap();
-
- let s = if array.is_null($row) {
- "".to_string()
- } else {
- array
- .value_as_time($row)
- .map(|d| d.to_string())
- .unwrap_or_else(|| "ERROR CONVERTING DATE".to_string())
- };
-
- Ok(s)
- }};
-}
-
-macro_rules! make_string_datetime {
- ($array_type:ty, $column: ident, $row: ident) => {{
- let array = $column.as_any().downcast_ref::<$array_type>().unwrap();
-
- let s = if array.is_null($row) {
- "".to_string()
- } else {
- array
- .value_as_datetime($row)
- .map(|d| d.to_string())
- .unwrap_or_else(|| "ERROR CONVERTING DATE".to_string())
- };
-
- Ok(s)
- }};
-}
-
-// It's not possible to do array.value($row).to_string() for &[u8], let's format it as hex
-macro_rules! make_string_hex {
- ($array_type:ty, $column: ident, $row: ident) => {{
- let array = $column.as_any().downcast_ref::<$array_type>().unwrap();
-
- let s = if array.is_null($row) {
- "".to_string()
- } else {
- let mut tmp = "".to_string();
-
- for character in array.value($row) {
- tmp += &format!("{:02x}", character);
- }
-
- tmp
- };
-
- Ok(s)
- }};
-}
-
-macro_rules! make_string_from_list {
- ($column: ident, $row: ident) => {{
- let list = $column
- .as_any()
- .downcast_ref::<array::ListArray>()
- .ok_or(ArrowError::InvalidArgumentError(format!(
- "Repl error: could not convert list column to list array."
- )))?
- .value($row);
- let string_values = (0..list.len())
- .map(|i| array_value_to_string(&list.clone(), i))
- .collect::<Result<Vec<String>>>()?;
- Ok(format!("[{}]", string_values.join(", ")))
- }};
-}
-
-#[inline(always)]
-pub fn make_string_from_decimal(column: &Arc<dyn Array>, row: usize) -> Result<String> {
- let array = column
- .as_any()
- .downcast_ref::<array::DecimalArray>()
- .unwrap();
-
- let formatted_decimal = array.value_as_string(row);
- Ok(formatted_decimal)
-}
-
-/// Get the value at the given row in an array as a String.
-///
-/// Note this function is quite inefficient and is unlikely to be
-/// suitable for converting large arrays or record batches.
-pub fn array_value_to_string(column: &array::ArrayRef, row: usize) -> Result<String> {
- if column.is_null(row) {
- return Ok("".to_string());
- }
- match column.data_type() {
- DataType::Utf8 => make_string!(array::StringArray, column, row),
- DataType::LargeUtf8 => make_string!(array::LargeStringArray, column, row),
- DataType::Binary => make_string_hex!(array::BinaryArray, column, row),
- DataType::LargeBinary => make_string_hex!(array::LargeBinaryArray, column, row),
- DataType::Boolean => make_string!(array::BooleanArray, column, row),
- DataType::Int8 => make_string!(array::Int8Array, column, row),
- DataType::Int16 => make_string!(array::Int16Array, column, row),
- DataType::Int32 => make_string!(array::Int32Array, column, row),
- DataType::Int64 => make_string!(array::Int64Array, column, row),
- DataType::UInt8 => make_string!(array::UInt8Array, column, row),
- DataType::UInt16 => make_string!(array::UInt16Array, column, row),
- DataType::UInt32 => make_string!(array::UInt32Array, column, row),
- DataType::UInt64 => make_string!(array::UInt64Array, column, row),
- DataType::Float16 => make_string!(array::Float32Array, column, row),
- DataType::Float32 => make_string!(array::Float32Array, column, row),
- DataType::Float64 => make_string!(array::Float64Array, column, row),
- DataType::Decimal(..) => make_string_from_decimal(column, row),
- DataType::Timestamp(unit, _) if *unit == TimeUnit::Second => {
- make_string_datetime!(array::TimestampSecondArray, column, row)
- }
- DataType::Timestamp(unit, _) if *unit == TimeUnit::Millisecond => {
- make_string_datetime!(array::TimestampMillisecondArray, column, row)
- }
- DataType::Timestamp(unit, _) if *unit == TimeUnit::Microsecond => {
- make_string_datetime!(array::TimestampMicrosecondArray, column, row)
- }
- DataType::Timestamp(unit, _) if *unit == TimeUnit::Nanosecond => {
- make_string_datetime!(array::TimestampNanosecondArray, column, row)
- }
- DataType::Date32 => make_string_date!(array::Date32Array, column, row),
- DataType::Date64 => make_string_date!(array::Date64Array, column, row),
- DataType::Time32(unit) if *unit == TimeUnit::Second => {
- make_string_time!(array::Time32SecondArray, column, row)
- }
- DataType::Time32(unit) if *unit == TimeUnit::Millisecond => {
- make_string_time!(array::Time32MillisecondArray, column, row)
- }
- DataType::Time64(unit) if *unit == TimeUnit::Microsecond => {
- make_string_time!(array::Time64MicrosecondArray, column, row)
- }
- DataType::Time64(unit) if *unit == TimeUnit::Nanosecond => {
- make_string_time!(array::Time64NanosecondArray, column, row)
- }
- DataType::Interval(unit) => match unit {
- IntervalUnit::DayTime => {
- make_string_interval_day_time!(column, row)
- }
- IntervalUnit::YearMonth => {
- make_string_interval_year_month!(column, row)
- }
- },
- DataType::List(_) => make_string_from_list!(column, row),
- DataType::Dictionary(index_type, _value_type) => match **index_type {
- DataType::Int8 => dict_array_value_to_string::<Int8Type>(column, row),
- DataType::Int16 => dict_array_value_to_string::<Int16Type>(column, row),
- DataType::Int32 => dict_array_value_to_string::<Int32Type>(column, row),
- DataType::Int64 => dict_array_value_to_string::<Int64Type>(column, row),
- DataType::UInt8 => dict_array_value_to_string::<UInt8Type>(column, row),
- DataType::UInt16 => dict_array_value_to_string::<UInt16Type>(column, row),
- DataType::UInt32 => dict_array_value_to_string::<UInt32Type>(column, row),
- DataType::UInt64 => dict_array_value_to_string::<UInt64Type>(column, row),
- _ => Err(ArrowError::InvalidArgumentError(format!(
- "Pretty printing not supported for {:?} due to index type",
- column.data_type()
- ))),
- },
- _ => Err(ArrowError::InvalidArgumentError(format!(
- "Pretty printing not implemented for {:?} type",
- column.data_type()
- ))),
- }
-}
-
-/// Converts the value of the dictionary array at `row` to a String
-fn dict_array_value_to_string<K: ArrowPrimitiveType>(
- colum: &array::ArrayRef,
- row: usize,
-) -> Result<String> {
- let dict_array = colum.as_any().downcast_ref::<DictionaryArray<K>>().unwrap();
-
- let keys_array = dict_array.keys();
-
- if keys_array.is_null(row) {
- return Ok(String::from(""));
- }
-
- let dict_index = keys_array.value(row).to_usize().ok_or_else(|| {
- ArrowError::InvalidArgumentError(format!(
- "Can not convert value {:?} at index {:?} to usize for string conversion.",
- keys_array.value(row),
- row
- ))
- })?;
-
- array_value_to_string(&dict_array.values(), dict_index)
-}
diff --git a/arrow/src/util/integration_util.rs b/arrow/src/util/integration_util.rs
deleted file mode 100644
index 2403151..0000000
--- a/arrow/src/util/integration_util.rs
+++ /dev/null
@@ -1,958 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-//! Utils for JSON integration testing
-//!
-//! These utilities define structs that read the integration JSON format for integration testing purposes.
-
-use serde_derive::{Deserialize, Serialize};
-use serde_json::{Map as SJMap, Number as VNumber, Value};
-
-use crate::array::*;
-use crate::datatypes::*;
-use crate::error::Result;
-use crate::record_batch::{RecordBatch, RecordBatchReader};
-
-/// A struct that represents an Arrow file with a schema and record batches
-#[derive(Deserialize, Serialize, Debug)]
-pub struct ArrowJson {
- pub schema: ArrowJsonSchema,
- pub batches: Vec<ArrowJsonBatch>,
- #[serde(skip_serializing_if = "Option::is_none")]
- pub dictionaries: Option<Vec<ArrowJsonDictionaryBatch>>,
-}
-
-/// A struct that partially reads the Arrow JSON schema.
-///
-/// Fields are left as JSON `Value` as they vary by `DataType`
-#[derive(Deserialize, Serialize, Debug)]
-pub struct ArrowJsonSchema {
- pub fields: Vec<ArrowJsonField>,
-}
-
-/// Fields are left as JSON `Value` as they vary by `DataType`
-#[derive(Deserialize, Serialize, Debug)]
-pub struct ArrowJsonField {
- pub name: String,
- #[serde(rename = "type")]
- pub field_type: Value,
- pub nullable: bool,
- pub children: Vec<ArrowJsonField>,
- #[serde(skip_serializing_if = "Option::is_none")]
- pub dictionary: Option<ArrowJsonFieldDictionary>,
- #[serde(skip_serializing_if = "Option::is_none")]
- pub metadata: Option<Value>,
-}
-
-impl From<&Field> for ArrowJsonField {
- fn from(field: &Field) -> Self {
- let metadata_value = match field.metadata() {
- Some(kv_list) => {
- let mut array = Vec::new();
- for (k, v) in kv_list {
- let mut kv_map = SJMap::new();
- kv_map.insert(k.clone(), Value::String(v.clone()));
- array.push(Value::Object(kv_map));
- }
- if !array.is_empty() {
- Some(Value::Array(array))
- } else {
- None
- }
- }
- _ => None,
- };
-
- Self {
- name: field.name().to_string(),
- field_type: field.data_type().to_json(),
- nullable: field.is_nullable(),
- children: vec![],
- dictionary: None, // TODO: not enough info
- metadata: metadata_value,
- }
- }
-}
-
-#[derive(Deserialize, Serialize, Debug)]
-pub struct ArrowJsonFieldDictionary {
- pub id: i64,
- #[serde(rename = "indexType")]
- pub index_type: DictionaryIndexType,
- #[serde(rename = "isOrdered")]
- pub is_ordered: bool,
-}
-
-#[derive(Deserialize, Serialize, Debug)]
-pub struct DictionaryIndexType {
- pub name: String,
- #[serde(rename = "isSigned")]
- pub is_signed: bool,
- #[serde(rename = "bitWidth")]
- pub bit_width: i64,
-}
-
-/// A struct that partially reads the Arrow JSON record batch
-#[derive(Deserialize, Serialize, Debug)]
-pub struct ArrowJsonBatch {
- count: usize,
- pub columns: Vec<ArrowJsonColumn>,
-}
-
-/// A struct that partially reads the Arrow JSON dictionary batch
-#[derive(Deserialize, Serialize, Debug)]
-#[allow(non_snake_case)]
-pub struct ArrowJsonDictionaryBatch {
- pub id: i64,
- pub data: ArrowJsonBatch,
-}
-
-/// A struct that partially reads the Arrow JSON column/array
-#[derive(Deserialize, Serialize, Clone, Debug)]
-pub struct ArrowJsonColumn {
- name: String,
- pub count: usize,
- #[serde(rename = "VALIDITY")]
- pub validity: Option<Vec<u8>>,
- #[serde(rename = "DATA")]
- pub data: Option<Vec<Value>>,
- #[serde(rename = "OFFSET")]
- pub offset: Option<Vec<Value>>, // leaving as Value as 64-bit offsets are strings
- pub children: Option<Vec<ArrowJsonColumn>>,
-}
-
-impl ArrowJson {
- /// Compare the Arrow JSON with a record batch reader
- pub fn equals_reader(&self, reader: &mut dyn RecordBatchReader) -> bool {
- if !self.schema.equals_schema(&reader.schema()) {
- return false;
- }
- self.batches.iter().all(|col| {
- let batch = reader.next();
- match batch {
- Some(Ok(batch)) => col.equals_batch(&batch),
- _ => false,
- }
- })
- }
-}
-
-impl ArrowJsonSchema {
- /// Compare the Arrow JSON schema with the Arrow `Schema`
- fn equals_schema(&self, schema: &Schema) -> bool {
- let field_len = self.fields.len();
- if field_len != schema.fields().len() {
- return false;
- }
- for i in 0..field_len {
- let json_field = &self.fields[i];
- let field = schema.field(i);
- if !json_field.equals_field(field) {
- return false;
- }
- }
- true
- }
-}
-
-impl ArrowJsonField {
- /// Compare the Arrow JSON field with the Arrow `Field`
- fn equals_field(&self, field: &Field) -> bool {
- // convert to a field
- match self.to_arrow_field() {
- Ok(self_field) => {
- assert_eq!(&self_field, field, "Arrow fields not the same");
- true
- }
- Err(e) => {
- eprintln!(
- "Encountered error while converting JSON field to Arrow field: {:?}",
- e
- );
- false
- }
- }
- }
-
- /// Convert to an Arrow Field
- /// TODO: convert to use an Into
- fn to_arrow_field(&self) -> Result<Field> {
- // a bit regressive, but we have to convert the field to JSON in order to convert it
- let field = serde_json::to_value(self)?;
- Field::from(&field)
- }
-}
-
-impl ArrowJsonBatch {
- /// Compare the Arrow JSON record batch with a `RecordBatch`
- fn equals_batch(&self, batch: &RecordBatch) -> bool {
- if self.count != batch.num_rows() {
- return false;
- }
- let num_columns = self.columns.len();
- if num_columns != batch.num_columns() {
- return false;
- }
- let schema = batch.schema();
- self.columns
- .iter()
- .zip(batch.columns())
- .zip(schema.fields())
- .all(|((col, arr), field)| {
- // compare each column based on its type
- if &col.name != field.name() {
- return false;
- }
- let json_array: Vec<Value> = json_from_col(&col, field.data_type());
- match field.data_type() {
- DataType::Null => {
- let arr: &NullArray =
- arr.as_any().downcast_ref::<NullArray>().unwrap();
- // NullArrays should have the same length, json_array is empty
- arr.len() == col.count
- }
- DataType::Boolean => {
- let arr = arr.as_any().downcast_ref::<BooleanArray>().unwrap();
- arr.equals_json(&json_array.iter().collect::<Vec<&Value>>()[..])
- }
- DataType::Int8 => {
- let arr = arr.as_any().downcast_ref::<Int8Array>().unwrap();
- arr.equals_json(&json_array.iter().collect::<Vec<&Value>>()[..])
- }
- DataType::Int16 => {
- let arr = arr.as_any().downcast_ref::<Int16Array>().unwrap();
- arr.equals_json(&json_array.iter().collect::<Vec<&Value>>()[..])
- }
- DataType::Int32 | DataType::Date32 | DataType::Time32(_) => {
- let arr = Int32Array::from(arr.data().clone());
- arr.equals_json(&json_array.iter().collect::<Vec<&Value>>()[..])
- }
- DataType::Int64
- | DataType::Date64
- | DataType::Time64(_)
- | DataType::Timestamp(_, _)
- | DataType::Duration(_) => {
- let arr = Int64Array::from(arr.data().clone());
- arr.equals_json(&json_array.iter().collect::<Vec<&Value>>()[..])
- }
- DataType::Interval(IntervalUnit::YearMonth) => {
- let arr = IntervalYearMonthArray::from(arr.data().clone());
- arr.equals_json(&json_array.iter().collect::<Vec<&Value>>()[..])
- }
- DataType::Interval(IntervalUnit::DayTime) => {
- let arr = IntervalDayTimeArray::from(arr.data().clone());
- let x = json_array
- .iter()
- .map(|v| {
- match v {
- Value::Null => Value::Null,
- Value::Object(v) => {
- // interval has days and milliseconds
- let days: i32 =
- v.get("days").unwrap().as_i64().unwrap()
- as i32;
- let milliseconds: i32 = v
- .get("milliseconds")
- .unwrap()
- .as_i64()
- .unwrap()
- as i32;
- let value: i64 = unsafe {
- std::mem::transmute::<[i32; 2], i64>([
- days,
- milliseconds,
- ])
- };
- Value::Number(VNumber::from(value))
- }
- // return null if Value is not an object
- _ => Value::Null,
- }
- })
- .collect::<Vec<Value>>();
- arr.equals_json(&x.iter().collect::<Vec<&Value>>()[..])
- }
- DataType::UInt8 => {
- let arr = arr.as_any().downcast_ref::<UInt8Array>().unwrap();
- arr.equals_json(&json_array.iter().collect::<Vec<&Value>>()[..])
- }
- DataType::UInt16 => {
- let arr = arr.as_any().downcast_ref::<UInt16Array>().unwrap();
- arr.equals_json(&json_array.iter().collect::<Vec<&Value>>()[..])
- }
- DataType::UInt32 => {
- let arr = arr.as_any().downcast_ref::<UInt32Array>().unwrap();
- arr.equals_json(&json_array.iter().collect::<Vec<&Value>>()[..])
- }
- DataType::UInt64 => {
- let arr = arr.as_any().downcast_ref::<UInt64Array>().unwrap();
- arr.equals_json(&json_array.iter().collect::<Vec<&Value>>()[..])
- }
- DataType::Float32 => {
- let arr = arr.as_any().downcast_ref::<Float32Array>().unwrap();
- arr.equals_json(&json_array.iter().collect::<Vec<&Value>>()[..])
- }
- DataType::Float64 => {
- let arr = arr.as_any().downcast_ref::<Float64Array>().unwrap();
- arr.equals_json(&json_array.iter().collect::<Vec<&Value>>()[..])
- }
- DataType::Binary => {
- let arr = arr.as_any().downcast_ref::<BinaryArray>().unwrap();
- arr.equals_json(&json_array.iter().collect::<Vec<&Value>>()[..])
- }
- DataType::LargeBinary => {
- let arr =
- arr.as_any().downcast_ref::<LargeBinaryArray>().unwrap();
- arr.equals_json(&json_array.iter().collect::<Vec<&Value>>()[..])
- }
- DataType::FixedSizeBinary(_) => {
- let arr =
- arr.as_any().downcast_ref::<FixedSizeBinaryArray>().unwrap();
- arr.equals_json(&json_array.iter().collect::<Vec<&Value>>()[..])
- }
- DataType::Utf8 => {
- let arr = arr.as_any().downcast_ref::<StringArray>().unwrap();
- arr.equals_json(&json_array.iter().collect::<Vec<&Value>>()[..])
- }
- DataType::LargeUtf8 => {
- let arr =
- arr.as_any().downcast_ref::<LargeStringArray>().unwrap();
- arr.equals_json(&json_array.iter().collect::<Vec<&Value>>()[..])
- }
- DataType::List(_) => {
- let arr = arr.as_any().downcast_ref::<ListArray>().unwrap();
- arr.equals_json(&json_array.iter().collect::<Vec<&Value>>()[..])
- }
- DataType::LargeList(_) => {
- let arr = arr.as_any().downcast_ref::<LargeListArray>().unwrap();
- arr.equals_json(&json_array.iter().collect::<Vec<&Value>>()[..])
- }
- DataType::FixedSizeList(_, _) => {
- let arr =
- arr.as_any().downcast_ref::<FixedSizeListArray>().unwrap();
- arr.equals_json(&json_array.iter().collect::<Vec<&Value>>()[..])
- }
- DataType::Struct(_) => {
- let arr = arr.as_any().downcast_ref::<StructArray>().unwrap();
- arr.equals_json(&json_array.iter().collect::<Vec<&Value>>()[..])
- }
- DataType::Decimal(_, _) => {
- let arr = arr.as_any().downcast_ref::<DecimalArray>().unwrap();
- arr.equals_json(&json_array.iter().collect::<Vec<&Value>>()[..])
- }
- DataType::Dictionary(ref key_type, _) => match key_type.as_ref() {
- DataType::Int8 => {
- let arr = arr
- .as_any()
- .downcast_ref::<Int8DictionaryArray>()
- .unwrap();
- arr.equals_json(
- &json_array.iter().collect::<Vec<&Value>>()[..],
- )
- }
- DataType::Int16 => {
- let arr = arr
- .as_any()
- .downcast_ref::<Int16DictionaryArray>()
- .unwrap();
- arr.equals_json(
- &json_array.iter().collect::<Vec<&Value>>()[..],
- )
- }
- DataType::Int32 => {
- let arr = arr
- .as_any()
- .downcast_ref::<Int32DictionaryArray>()
- .unwrap();
- arr.equals_json(
- &json_array.iter().collect::<Vec<&Value>>()[..],
- )
- }
- DataType::Int64 => {
- let arr = arr
- .as_any()
- .downcast_ref::<Int64DictionaryArray>()
- .unwrap();
- arr.equals_json(
- &json_array.iter().collect::<Vec<&Value>>()[..],
- )
- }
- DataType::UInt8 => {
- let arr = arr
- .as_any()
- .downcast_ref::<UInt8DictionaryArray>()
- .unwrap();
- arr.equals_json(
- &json_array.iter().collect::<Vec<&Value>>()[..],
- )
- }
- DataType::UInt16 => {
- let arr = arr
- .as_any()
- .downcast_ref::<UInt16DictionaryArray>()
- .unwrap();
- arr.equals_json(
- &json_array.iter().collect::<Vec<&Value>>()[..],
- )
- }
- DataType::UInt32 => {
- let arr = arr
- .as_any()
- .downcast_ref::<UInt32DictionaryArray>()
- .unwrap();
- arr.equals_json(
- &json_array.iter().collect::<Vec<&Value>>()[..],
- )
- }
- DataType::UInt64 => {
- let arr = arr
- .as_any()
- .downcast_ref::<UInt64DictionaryArray>()
- .unwrap();
- arr.equals_json(
- &json_array.iter().collect::<Vec<&Value>>()[..],
- )
- }
- t => panic!("Unsupported dictionary comparison for {:?}", t),
- },
- t => panic!("Unsupported comparison for {:?}", t),
- }
- })
- }
-
- pub fn from_batch(batch: &RecordBatch) -> ArrowJsonBatch {
- let mut json_batch = ArrowJsonBatch {
- count: batch.num_rows(),
- columns: Vec::with_capacity(batch.num_columns()),
- };
-
- for (col, field) in batch.columns().iter().zip(batch.schema().fields.iter()) {
- let json_col = match field.data_type() {
- DataType::Int8 => {
- let col = col.as_any().downcast_ref::<Int8Array>().unwrap();
-
- let mut validity: Vec<u8> = Vec::with_capacity(col.len());
- let mut data: Vec<Value> = Vec::with_capacity(col.len());
-
- for i in 0..col.len() {
- if col.is_null(i) {
- validity.push(1);
- data.push(0i8.into());
- } else {
- validity.push(0);
- data.push(col.value(i).into());
- }
- }
-
- ArrowJsonColumn {
- name: field.name().clone(),
- count: col.len(),
- validity: Some(validity),
- data: Some(data),
- offset: None,
- children: None,
- }
- }
- _ => ArrowJsonColumn {
- name: field.name().clone(),
- count: col.len(),
- validity: None,
- data: None,
- offset: None,
- children: None,
- },
- };
-
- json_batch.columns.push(json_col);
- }
-
- json_batch
- }
-}
-
-/// Convert an Arrow JSON column/array into a vector of `Value`
-fn json_from_col(col: &ArrowJsonColumn, data_type: &DataType) -> Vec<Value> {
- match data_type {
- DataType::List(field) => json_from_list_col(col, field.data_type()),
- DataType::FixedSizeList(field, list_size) => {
- json_from_fixed_size_list_col(col, field.data_type(), *list_size as usize)
- }
- DataType::Struct(fields) => json_from_struct_col(col, fields),
- DataType::Int64
- | DataType::UInt64
- | DataType::Date64
- | DataType::Time64(_)
- | DataType::Timestamp(_, _)
- | DataType::Duration(_) => {
- // convert int64 data from strings to numbers
- let converted_col: Vec<Value> = col
- .data
- .clone()
- .unwrap()
- .iter()
- .map(|v| {
- Value::Number(match v {
- Value::Number(number) => number.clone(),
- Value::String(string) => VNumber::from(
- string
- .parse::<i64>()
- .expect("Unable to parse string as i64"),
- ),
- t => panic!("Cannot convert {} to number", t),
- })
- })
- .collect();
- merge_json_array(
- col.validity.as_ref().unwrap().as_slice(),
- converted_col.as_slice(),
- )
- }
- DataType::Null => vec![],
- _ => merge_json_array(
- col.validity.as_ref().unwrap().as_slice(),
- &col.data.clone().unwrap(),
- ),
- }
-}
-
-/// Merge VALIDITY and DATA vectors from a primitive data type into a `Value` vector with nulls
-fn merge_json_array(validity: &[u8], data: &[Value]) -> Vec<Value> {
- validity
- .iter()
- .zip(data)
- .map(|(v, d)| match v {
- 0 => Value::Null,
- 1 => d.clone(),
- _ => panic!("Validity data should be 0 or 1"),
- })
- .collect()
-}
-
-/// Convert an Arrow JSON column/array of a `DataType::Struct` into a vector of `Value`
-fn json_from_struct_col(col: &ArrowJsonColumn, fields: &[Field]) -> Vec<Value> {
- let mut values = Vec::with_capacity(col.count);
-
- let children: Vec<Vec<Value>> = col
- .children
- .clone()
- .unwrap()
- .iter()
- .zip(fields)
- .map(|(child, field)| json_from_col(child, field.data_type()))
- .collect();
-
- // create a struct from children
- for j in 0..col.count {
- let mut map = serde_json::map::Map::new();
- for i in 0..children.len() {
- map.insert(fields[i].name().to_string(), children[i][j].clone());
- }
- values.push(Value::Object(map));
- }
-
- values
-}
-
-/// Convert an Arrow JSON column/array of a `DataType::List` into a vector of `Value`
-fn json_from_list_col(col: &ArrowJsonColumn, data_type: &DataType) -> Vec<Value> {
- let mut values = Vec::with_capacity(col.count);
-
- // get the inner array
- let child = &col.children.clone().expect("list type must have children")[0];
- let offsets: Vec<usize> = col
- .offset
- .clone()
- .unwrap()
- .iter()
- .map(|o| match o {
- Value::String(s) => s.parse::<usize>().unwrap(),
- Value::Number(n) => n.as_u64().unwrap() as usize,
- _ => panic!(
- "Offsets should be numbers or strings that are convertible to numbers"
- ),
- })
- .collect();
- let inner = match data_type {
- DataType::List(ref field) => json_from_col(child, field.data_type()),
- DataType::Struct(fields) => json_from_struct_col(col, fields),
- _ => merge_json_array(
- child.validity.as_ref().unwrap().as_slice(),
- &child.data.clone().unwrap(),
- ),
- };
-
- for i in 0..col.count {
- match &col.validity {
- Some(validity) => match &validity[i] {
- 0 => values.push(Value::Null),
- 1 => {
- values.push(Value::Array(inner[offsets[i]..offsets[i + 1]].to_vec()))
- }
- _ => panic!("Validity data should be 0 or 1"),
- },
- None => {
- // Null type does not have a validity vector
- }
- }
- }
-
- values
-}
-
-/// Convert an Arrow JSON column/array of a `DataType::List` into a vector of `Value`
-fn json_from_fixed_size_list_col(
- col: &ArrowJsonColumn,
- data_type: &DataType,
- list_size: usize,
-) -> Vec<Value> {
- let mut values = Vec::with_capacity(col.count);
-
- // get the inner array
- let child = &col.children.clone().expect("list type must have children")[0];
- let inner = match data_type {
- DataType::List(ref field) => json_from_col(child, field.data_type()),
- DataType::FixedSizeList(ref field, _) => json_from_col(child, field.data_type()),
- DataType::Struct(fields) => json_from_struct_col(col, fields),
- _ => merge_json_array(
- child.validity.as_ref().unwrap().as_slice(),
- &child.data.clone().unwrap(),
- ),
- };
-
- for i in 0..col.count {
- match &col.validity {
- Some(validity) => match &validity[i] {
- 0 => values.push(Value::Null),
- 1 => values.push(Value::Array(
- inner[(list_size * i)..(list_size * (i + 1))].to_vec(),
- )),
- _ => panic!("Validity data should be 0 or 1"),
- },
- None => {}
- }
- }
-
- values
-}
-
-#[cfg(test)]
-mod tests {
- use super::*;
-
- use std::fs::File;
- use std::io::Read;
- use std::sync::Arc;
-
- use crate::buffer::Buffer;
-
- #[test]
- fn test_schema_equality() {
- let json = r#"
- {
- "fields": [
- {
- "name": "c1",
- "type": {"name": "int", "isSigned": true, "bitWidth": 32},
- "nullable": true,
- "children": []
- },
- {
- "name": "c2",
- "type": {"name": "floatingpoint", "precision": "DOUBLE"},
- "nullable": true,
- "children": []
- },
- {
- "name": "c3",
- "type": {"name": "utf8"},
- "nullable": true,
- "children": []
- },
- {
- "name": "c4",
- "type": {
- "name": "list"
- },
- "nullable": true,
- "children": [
- {
- "name": "custom_item",
- "type": {
- "name": "int",
- "isSigned": true,
- "bitWidth": 32
- },
- "nullable": false,
- "children": []
- }
- ]
- }
- ]
- }"#;
- let json_schema: ArrowJsonSchema = serde_json::from_str(json).unwrap();
- let schema = Schema::new(vec![
- Field::new("c1", DataType::Int32, true),
- Field::new("c2", DataType::Float64, true),
- Field::new("c3", DataType::Utf8, true),
- Field::new(
- "c4",
- DataType::List(Box::new(Field::new(
- "custom_item",
- DataType::Int32,
- false,
- ))),
- true,
- ),
- ]);
- assert!(json_schema.equals_schema(&schema));
- }
-
- #[test]
- #[cfg_attr(miri, ignore)] // running forever
- fn test_arrow_data_equality() {
- let secs_tz = Some("Europe/Budapest".to_string());
- let millis_tz = Some("America/New_York".to_string());
- let micros_tz = Some("UTC".to_string());
- let nanos_tz = Some("Africa/Johannesburg".to_string());
-
- let schema = Schema::new(vec![
- {
- let mut f =
- Field::new("bools-with-metadata-map", DataType::Boolean, true);
- f.set_metadata(Some(
- [("k".to_string(), "v".to_string())]
- .iter()
- .cloned()
- .collect(),
- ));
- f
- },
- {
- let mut f =
- Field::new("bools-with-metadata-vec", DataType::Boolean, true);
- f.set_metadata(Some(
- [("k2".to_string(), "v2".to_string())]
- .iter()
- .cloned()
- .collect(),
- ));
- f
- },
- Field::new("bools", DataType::Boolean, true),
- Field::new("int8s", DataType::Int8, true),
- Field::new("int16s", DataType::Int16, true),
- Field::new("int32s", DataType::Int32, true),
- Field::new("int64s", DataType::Int64, true),
- Field::new("uint8s", DataType::UInt8, true),
- Field::new("uint16s", DataType::UInt16, true),
- Field::new("uint32s", DataType::UInt32, true),
- Field::new("uint64s", DataType::UInt64, true),
- Field::new("float32s", DataType::Float32, true),
- Field::new("float64s", DataType::Float64, true),
- Field::new("date_days", DataType::Date32, true),
- Field::new("date_millis", DataType::Date64, true),
- Field::new("time_secs", DataType::Time32(TimeUnit::Second), true),
- Field::new("time_millis", DataType::Time32(TimeUnit::Millisecond), true),
- Field::new("time_micros", DataType::Time64(TimeUnit::Microsecond), true),
- Field::new("time_nanos", DataType::Time64(TimeUnit::Nanosecond), true),
- Field::new("ts_secs", DataType::Timestamp(TimeUnit::Second, None), true),
- Field::new(
- "ts_millis",
- DataType::Timestamp(TimeUnit::Millisecond, None),
- true,
- ),
- Field::new(
- "ts_micros",
- DataType::Timestamp(TimeUnit::Microsecond, None),
- true,
- ),
- Field::new(
- "ts_nanos",
- DataType::Timestamp(TimeUnit::Nanosecond, None),
- true,
- ),
- Field::new(
- "ts_secs_tz",
- DataType::Timestamp(TimeUnit::Second, secs_tz.clone()),
- true,
- ),
- Field::new(
- "ts_millis_tz",
- DataType::Timestamp(TimeUnit::Millisecond, millis_tz.clone()),
- true,
- ),
- Field::new(
- "ts_micros_tz",
- DataType::Timestamp(TimeUnit::Microsecond, micros_tz.clone()),
- true,
- ),
- Field::new(
- "ts_nanos_tz",
- DataType::Timestamp(TimeUnit::Nanosecond, nanos_tz.clone()),
- true,
- ),
- Field::new("utf8s", DataType::Utf8, true),
- Field::new(
- "lists",
- DataType::List(Box::new(Field::new("item", DataType::Int32, true))),
- true,
- ),
- Field::new(
- "structs",
- DataType::Struct(vec![
- Field::new("int32s", DataType::Int32, true),
- Field::new("utf8s", DataType::Utf8, true),
- ]),
- true,
- ),
- ]);
-
- let bools_with_metadata_map =
- BooleanArray::from(vec![Some(true), None, Some(false)]);
- let bools_with_metadata_vec =
- BooleanArray::from(vec![Some(true), None, Some(false)]);
- let bools = BooleanArray::from(vec![Some(true), None, Some(false)]);
- let int8s = Int8Array::from(vec![Some(1), None, Some(3)]);
- let int16s = Int16Array::from(vec![Some(1), None, Some(3)]);
- let int32s = Int32Array::from(vec![Some(1), None, Some(3)]);
- let int64s = Int64Array::from(vec![Some(1), None, Some(3)]);
- let uint8s = UInt8Array::from(vec![Some(1), None, Some(3)]);
- let uint16s = UInt16Array::from(vec![Some(1), None, Some(3)]);
- let uint32s = UInt32Array::from(vec![Some(1), None, Some(3)]);
- let uint64s = UInt64Array::from(vec![Some(1), None, Some(3)]);
- let float32s = Float32Array::from(vec![Some(1.0), None, Some(3.0)]);
- let float64s = Float64Array::from(vec![Some(1.0), None, Some(3.0)]);
- let date_days = Date32Array::from(vec![Some(1196848), None, None]);
- let date_millis = Date64Array::from(vec![
- Some(167903550396207),
- Some(29923997007884),
- Some(30612271819236),
- ]);
- let time_secs =
- Time32SecondArray::from(vec![Some(27974), Some(78592), Some(43207)]);
- let time_millis = Time32MillisecondArray::from(vec![
- Some(6613125),
- Some(74667230),
- Some(52260079),
- ]);
- let time_micros =
- Time64MicrosecondArray::from(vec![Some(62522958593), None, None]);
- let time_nanos = Time64NanosecondArray::from(vec![
- Some(73380123595985),
- None,
- Some(16584393546415),
- ]);
- let ts_secs = TimestampSecondArray::from_opt_vec(
- vec![None, Some(193438817552), None],
- None,
- );
- let ts_millis = TimestampMillisecondArray::from_opt_vec(
- vec![None, Some(38606916383008), Some(58113709376587)],
- None,
- );
- let ts_micros =
- TimestampMicrosecondArray::from_opt_vec(vec![None, None, None], None);
- let ts_nanos = TimestampNanosecondArray::from_opt_vec(
- vec![None, None, Some(-6473623571954960143)],
- None,
- );
- let ts_secs_tz = TimestampSecondArray::from_opt_vec(
- vec![None, Some(193438817552), None],
- secs_tz,
- );
- let ts_millis_tz = TimestampMillisecondArray::from_opt_vec(
- vec![None, Some(38606916383008), Some(58113709376587)],
- millis_tz,
- );
- let ts_micros_tz =
- TimestampMicrosecondArray::from_opt_vec(vec![None, None, None], micros_tz);
- let ts_nanos_tz = TimestampNanosecondArray::from_opt_vec(
- vec![None, None, Some(-6473623571954960143)],
- nanos_tz,
- );
- let utf8s = StringArray::from(vec![Some("aa"), None, Some("bbb")]);
-
- let value_data = Int32Array::from(vec![None, Some(2), None, None]);
- let value_offsets = Buffer::from_slice_ref(&[0, 3, 4, 4]);
- let list_data_type =
- DataType::List(Box::new(Field::new("item", DataType::Int32, true)));
- let list_data = ArrayData::builder(list_data_type)
- .len(3)
- .add_buffer(value_offsets)
- .add_child_data(value_data.data().clone())
- .build();
- let lists = ListArray::from(list_data);
-
- let structs_int32s = Int32Array::from(vec![None, Some(-2), None]);
- let structs_utf8s = StringArray::from(vec![None, None, Some("aaaaaa")]);
- let structs = StructArray::from(vec![
- (
- Field::new("int32s", DataType::Int32, true),
- Arc::new(structs_int32s) as ArrayRef,
- ),
- (
- Field::new("utf8s", DataType::Utf8, true),
- Arc::new(structs_utf8s) as ArrayRef,
- ),
- ]);
-
- let record_batch = RecordBatch::try_new(
- Arc::new(schema.clone()),
- vec![
- Arc::new(bools_with_metadata_map),
- Arc::new(bools_with_metadata_vec),
- Arc::new(bools),
- Arc::new(int8s),
- Arc::new(int16s),
- Arc::new(int32s),
- Arc::new(int64s),
- Arc::new(uint8s),
- Arc::new(uint16s),
- Arc::new(uint32s),
- Arc::new(uint64s),
- Arc::new(float32s),
- Arc::new(float64s),
- Arc::new(date_days),
- Arc::new(date_millis),
- Arc::new(time_secs),
- Arc::new(time_millis),
- Arc::new(time_micros),
- Arc::new(time_nanos),
- Arc::new(ts_secs),
- Arc::new(ts_millis),
- Arc::new(ts_micros),
- Arc::new(ts_nanos),
- Arc::new(ts_secs_tz),
- Arc::new(ts_millis_tz),
- Arc::new(ts_micros_tz),
- Arc::new(ts_nanos_tz),
- Arc::new(utf8s),
- Arc::new(lists),
- Arc::new(structs),
- ],
- )
- .unwrap();
- let mut file = File::open("test/data/integration.json").unwrap();
- let mut json = String::new();
- file.read_to_string(&mut json).unwrap();
- let arrow_json: ArrowJson = serde_json::from_str(&json).unwrap();
- // test schemas
- assert!(arrow_json.schema.equals_schema(&schema));
- // test record batch
- assert!(arrow_json.batches[0].equals_batch(&record_batch));
- }
-}
diff --git a/arrow/src/util/mod.rs b/arrow/src/util/mod.rs
deleted file mode 100644
index b2fd4f7..0000000
--- a/arrow/src/util/mod.rs
+++ /dev/null
@@ -1,31 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-pub mod bench_util;
-pub mod bit_chunk_iterator;
-pub mod bit_util;
-pub mod data_gen;
-pub mod display;
-pub mod integration_util;
-#[cfg(feature = "prettyprint")]
-pub mod pretty;
-pub(crate) mod serialization;
-pub mod string_writer;
-pub mod test_util;
-
-mod trusted_len;
-pub(crate) use trusted_len::trusted_len_unzip;
diff --git a/arrow/src/util/pretty.rs b/arrow/src/util/pretty.rs
deleted file mode 100644
index d307406..0000000
--- a/arrow/src/util/pretty.rs
+++ /dev/null
@@ -1,510 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-//! Utilities for printing record batches. Note this module is not
-//! available unless `feature = "prettyprint"` is enabled.
-
-use crate::{array::ArrayRef, record_batch::RecordBatch};
-
-use prettytable::format;
-use prettytable::{Cell, Row, Table};
-
-use crate::error::Result;
-
-use super::display::array_value_to_string;
-
-///! Create a visual representation of record batches
-pub fn pretty_format_batches(results: &[RecordBatch]) -> Result<String> {
- Ok(create_table(results)?.to_string())
-}
-
-///! Create a visual representation of columns
-pub fn pretty_format_columns(col_name: &str, results: &[ArrayRef]) -> Result<String> {
- Ok(create_column(col_name, results)?.to_string())
-}
-
-///! Prints a visual representation of record batches to stdout
-pub fn print_batches(results: &[RecordBatch]) -> Result<()> {
- create_table(results)?.printstd();
- Ok(())
-}
-
-///! Prints a visual representation of a list of column to stdout
-pub fn print_columns(col_name: &str, results: &[ArrayRef]) -> Result<()> {
- create_column(col_name, results)?.printstd();
- Ok(())
-}
-
-///! Convert a series of record batches into a table
-fn create_table(results: &[RecordBatch]) -> Result<Table> {
- let mut table = Table::new();
- table.set_format(*format::consts::FORMAT_NO_LINESEP_WITH_TITLE);
-
- if results.is_empty() {
- return Ok(table);
- }
-
- let schema = results[0].schema();
-
- let mut header = Vec::new();
- for field in schema.fields() {
- header.push(Cell::new(&field.name()));
- }
- table.set_titles(Row::new(header));
-
- for batch in results {
- for row in 0..batch.num_rows() {
- let mut cells = Vec::new();
- for col in 0..batch.num_columns() {
- let column = batch.column(col);
- cells.push(Cell::new(&array_value_to_string(&column, row)?));
- }
- table.add_row(Row::new(cells));
- }
- }
-
- Ok(table)
-}
-
-fn create_column(field: &str, columns: &[ArrayRef]) -> Result<Table> {
- let mut table = Table::new();
- table.set_format(*format::consts::FORMAT_NO_LINESEP_WITH_TITLE);
-
- if columns.is_empty() {
- return Ok(table);
- }
-
- let header = vec![Cell::new(field)];
- table.set_titles(Row::new(header));
-
- for col in columns {
- for row in 0..col.len() {
- let cells = vec![Cell::new(&array_value_to_string(&col, row)?)];
- table.add_row(Row::new(cells));
- }
- }
-
- Ok(table)
-}
-
-#[cfg(test)]
-mod tests {
- use crate::{
- array::{
- self, new_null_array, Array, Date32Array, Date64Array, PrimitiveBuilder,
- StringBuilder, StringDictionaryBuilder, Time32MillisecondArray,
- Time32SecondArray, Time64MicrosecondArray, Time64NanosecondArray,
- TimestampMicrosecondArray, TimestampMillisecondArray,
- TimestampNanosecondArray, TimestampSecondArray,
- },
- datatypes::{DataType, Field, Int32Type, Schema},
- };
-
- use super::*;
- use crate::array::{DecimalBuilder, Int32Array};
- use std::sync::Arc;
-
- #[test]
- fn test_pretty_format_batches() -> Result<()> {
- // define a schema.
- let schema = Arc::new(Schema::new(vec![
- Field::new("a", DataType::Utf8, true),
- Field::new("b", DataType::Int32, true),
- ]));
-
- // define data.
- let batch = RecordBatch::try_new(
- schema,
- vec![
- Arc::new(array::StringArray::from(vec![
- Some("a"),
- Some("b"),
- None,
- Some("d"),
- ])),
- Arc::new(array::Int32Array::from(vec![
- Some(1),
- None,
- Some(10),
- Some(100),
- ])),
- ],
- )?;
-
- let table = pretty_format_batches(&[batch])?;
-
- let expected = vec![
- "+---+-----+",
- "| a | b |",
- "+---+-----+",
- "| a | 1 |",
- "| b | |",
- "| | 10 |",
- "| d | 100 |",
- "+---+-----+",
- ];
-
- let actual: Vec<&str> = table.lines().collect();
-
- assert_eq!(expected, actual, "Actual result:\n{}", table);
-
- Ok(())
- }
-
- #[test]
- fn test_pretty_format_columns() -> Result<()> {
- let columns = vec![
- Arc::new(array::StringArray::from(vec![
- Some("a"),
- Some("b"),
- None,
- Some("d"),
- ])) as ArrayRef,
- Arc::new(array::StringArray::from(vec![Some("e"), None, Some("g")])),
- ];
-
- let table = pretty_format_columns("a", &columns)?;
-
- let expected = vec![
- "+---+", "| a |", "+---+", "| a |", "| b |", "| |", "| d |", "| e |",
- "| |", "| g |", "+---+",
- ];
-
- let actual: Vec<&str> = table.lines().collect();
-
- assert_eq!(expected, actual, "Actual result:\n{}", table);
-
- Ok(())
- }
-
- #[test]
- fn test_pretty_format_null() {
- let schema = Arc::new(Schema::new(vec![
- Field::new("a", DataType::Utf8, true),
- Field::new("b", DataType::Int32, true),
- Field::new("c", DataType::Null, true),
- ]));
-
- let num_rows = 4;
- let arrays = schema
- .fields()
- .iter()
- .map(|f| new_null_array(f.data_type(), num_rows))
- .collect();
-
- // define data (null)
- let batch = RecordBatch::try_new(schema, arrays).unwrap();
-
- let table = pretty_format_batches(&[batch]).unwrap();
-
- let expected = vec![
- "+---+---+---+",
- "| a | b | c |",
- "+---+---+---+",
- "| | | |",
- "| | | |",
- "| | | |",
- "| | | |",
- "+---+---+---+",
- ];
-
- let actual: Vec<&str> = table.lines().collect();
-
- assert_eq!(expected, actual, "Actual result:\n{:#?}", table);
- }
-
- #[test]
- fn test_pretty_format_dictionary() -> Result<()> {
- // define a schema.
- let field_type =
- DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8));
- let schema = Arc::new(Schema::new(vec![Field::new("d1", field_type, true)]));
-
- let keys_builder = PrimitiveBuilder::<Int32Type>::new(10);
- let values_builder = StringBuilder::new(10);
- let mut builder = StringDictionaryBuilder::new(keys_builder, values_builder);
-
- builder.append("one")?;
- builder.append_null()?;
- builder.append("three")?;
- let array = Arc::new(builder.finish());
-
- let batch = RecordBatch::try_new(schema, vec![array])?;
-
- let table = pretty_format_batches(&[batch])?;
-
- let expected = vec![
- "+-------+",
- "| d1 |",
- "+-------+",
- "| one |",
- "| |",
- "| three |",
- "+-------+",
- ];
-
- let actual: Vec<&str> = table.lines().collect();
-
- assert_eq!(expected, actual, "Actual result:\n{}", table);
-
- Ok(())
- }
-
- /// Generate an array with type $ARRAYTYPE with a numeric value of
- /// $VALUE, and compare $EXPECTED_RESULT to the output of
- /// formatting that array with `pretty_format_batches`
- macro_rules! check_datetime {
- ($ARRAYTYPE:ident, $VALUE:expr, $EXPECTED_RESULT:expr) => {
- let mut builder = $ARRAYTYPE::builder(10);
- builder.append_value($VALUE).unwrap();
- builder.append_null().unwrap();
- let array = builder.finish();
-
- let schema = Arc::new(Schema::new(vec![Field::new(
- "f",
- array.data_type().clone(),
- true,
- )]));
- let batch = RecordBatch::try_new(schema, vec![Arc::new(array)]).unwrap();
-
- let table = pretty_format_batches(&[batch]).expect("formatting batches");
-
- let expected = $EXPECTED_RESULT;
- let actual: Vec<&str> = table.lines().collect();
-
- assert_eq!(expected, actual, "Actual result:\n\n{:#?}\n\n", actual);
- };
- }
-
- #[test]
- fn test_pretty_format_timestamp_second() {
- let expected = vec![
- "+---------------------+",
- "| f |",
- "+---------------------+",
- "| 1970-05-09 14:25:11 |",
- "| |",
- "+---------------------+",
- ];
- check_datetime!(TimestampSecondArray, 11111111, expected);
- }
-
- #[test]
- fn test_pretty_format_timestamp_millisecond() {
- let expected = vec![
- "+-------------------------+",
- "| f |",
- "+-------------------------+",
- "| 1970-01-01 03:05:11.111 |",
- "| |",
- "+-------------------------+",
- ];
- check_datetime!(TimestampMillisecondArray, 11111111, expected);
- }
-
- #[test]
- fn test_pretty_format_timestamp_microsecond() {
- let expected = vec![
- "+----------------------------+",
- "| f |",
- "+----------------------------+",
- "| 1970-01-01 00:00:11.111111 |",
- "| |",
- "+----------------------------+",
- ];
- check_datetime!(TimestampMicrosecondArray, 11111111, expected);
- }
-
- #[test]
- fn test_pretty_format_timestamp_nanosecond() {
- let expected = vec![
- "+-------------------------------+",
- "| f |",
- "+-------------------------------+",
- "| 1970-01-01 00:00:00.011111111 |",
- "| |",
- "+-------------------------------+",
- ];
- check_datetime!(TimestampNanosecondArray, 11111111, expected);
- }
-
- #[test]
- fn test_pretty_format_date_32() {
- let expected = vec![
- "+------------+",
- "| f |",
- "+------------+",
- "| 1973-05-19 |",
- "| |",
- "+------------+",
- ];
- check_datetime!(Date32Array, 1234, expected);
- }
-
- #[test]
- fn test_pretty_format_date_64() {
- let expected = vec![
- "+------------+",
- "| f |",
- "+------------+",
- "| 2005-03-18 |",
- "| |",
- "+------------+",
- ];
- check_datetime!(Date64Array, 1111111100000, expected);
- }
-
- #[test]
- fn test_pretty_format_time_32_second() {
- let expected = vec![
- "+----------+",
- "| f |",
- "+----------+",
- "| 00:18:31 |",
- "| |",
- "+----------+",
- ];
- check_datetime!(Time32SecondArray, 1111, expected);
- }
-
- #[test]
- fn test_pretty_format_time_32_millisecond() {
- let expected = vec![
- "+--------------+",
- "| f |",
- "+--------------+",
- "| 03:05:11.111 |",
- "| |",
- "+--------------+",
- ];
- check_datetime!(Time32MillisecondArray, 11111111, expected);
- }
-
- #[test]
- fn test_pretty_format_time_64_microsecond() {
- let expected = vec![
- "+-----------------+",
- "| f |",
- "+-----------------+",
- "| 00:00:11.111111 |",
- "| |",
- "+-----------------+",
- ];
- check_datetime!(Time64MicrosecondArray, 11111111, expected);
- }
-
- #[test]
- fn test_pretty_format_time_64_nanosecond() {
- let expected = vec![
- "+--------------------+",
- "| f |",
- "+--------------------+",
- "| 00:00:00.011111111 |",
- "| |",
- "+--------------------+",
- ];
- check_datetime!(Time64NanosecondArray, 11111111, expected);
- }
-
- #[test]
- fn test_int_display() -> Result<()> {
- let array = Arc::new(Int32Array::from(vec![6, 3])) as ArrayRef;
- let actual_one = array_value_to_string(&array, 0).unwrap();
- let expected_one = "6";
-
- let actual_two = array_value_to_string(&array, 1).unwrap();
- let expected_two = "3";
- assert_eq!(actual_one, expected_one);
- assert_eq!(actual_two, expected_two);
- Ok(())
- }
-
- #[test]
- fn test_decimal_display() -> Result<()> {
- let capacity = 10;
- let precision = 10;
- let scale = 2;
-
- let mut builder = DecimalBuilder::new(capacity, precision, scale);
- builder.append_value(101).unwrap();
- builder.append_null().unwrap();
- builder.append_value(200).unwrap();
- builder.append_value(3040).unwrap();
-
- let dm = Arc::new(builder.finish()) as ArrayRef;
-
- let schema = Arc::new(Schema::new(vec![Field::new(
- "f",
- dm.data_type().clone(),
- true,
- )]));
-
- let batch = RecordBatch::try_new(schema, vec![dm])?;
-
- let table = pretty_format_batches(&[batch])?;
-
- let expected = vec![
- "+-------+",
- "| f |",
- "+-------+",
- "| 1.01 |",
- "| |",
- "| 2.00 |",
- "| 30.40 |",
- "+-------+",
- ];
-
- let actual: Vec<&str> = table.lines().collect();
- assert_eq!(expected, actual, "Actual result:\n{}", table);
-
- Ok(())
- }
-
- #[test]
- fn test_decimal_display_zero_scale() -> Result<()> {
- let capacity = 10;
- let precision = 5;
- let scale = 0;
-
- let mut builder = DecimalBuilder::new(capacity, precision, scale);
- builder.append_value(101).unwrap();
- builder.append_null().unwrap();
- builder.append_value(200).unwrap();
- builder.append_value(3040).unwrap();
-
- let dm = Arc::new(builder.finish()) as ArrayRef;
-
- let schema = Arc::new(Schema::new(vec![Field::new(
- "f",
- dm.data_type().clone(),
- true,
- )]));
-
- let batch = RecordBatch::try_new(schema, vec![dm])?;
-
- let table = pretty_format_batches(&[batch])?;
- let expected = vec![
- "+------+", "| f |", "+------+", "| 101 |", "| |", "| 200 |",
- "| 3040 |", "+------+",
- ];
-
- let actual: Vec<&str> = table.lines().collect();
- assert_eq!(expected, actual, "Actual result:\n{}", table);
-
- Ok(())
- }
-}
diff --git a/arrow/src/util/serialization.rs b/arrow/src/util/serialization.rs
deleted file mode 100644
index 14d67ca..0000000
--- a/arrow/src/util/serialization.rs
+++ /dev/null
@@ -1,33 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-/// Converts numeric type to a `String`
-pub fn lexical_to_string<N: lexical_core::ToLexical>(n: N) -> String {
- let mut buf = Vec::<u8>::with_capacity(N::FORMATTED_SIZE_DECIMAL);
- unsafe {
- // JUSTIFICATION
- // Benefit
- // Allows using the faster serializer lexical core and convert to string
- // Soundness
- // Length of buf is set as written length afterwards. lexical_core
- // creates a valid string, so doesn't need to be checked.
- let slice = std::slice::from_raw_parts_mut(buf.as_mut_ptr(), buf.capacity());
- let len = lexical_core::write(n, slice).len();
- buf.set_len(len);
- String::from_utf8_unchecked(buf)
- }
-}
diff --git a/arrow/src/util/string_writer.rs b/arrow/src/util/string_writer.rs
deleted file mode 100644
index 4c61f18..0000000
--- a/arrow/src/util/string_writer.rs
+++ /dev/null
@@ -1,108 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-//! String Writer
-//! This string writer encapsulates `std::string::String` and
-//! implements `std::io::Write` trait, which makes String as a
-//! writable object like File.
-//!
-//! Example:
-//!
-//! ```
-//! #[cfg(feature = "csv")]
-//! {
-//! use arrow::array::*;
-//! use arrow::csv;
-//! use arrow::datatypes::*;
-//! use arrow::record_batch::RecordBatch;
-//! use arrow::util::string_writer::StringWriter;
-//! use std::sync::Arc;
-//!
-//! let schema = Schema::new(vec![
-//! Field::new("c1", DataType::Utf8, false),
-//! Field::new("c2", DataType::Float64, true),
-//! Field::new("c3", DataType::UInt32, false),
-//! Field::new("c3", DataType::Boolean, true),
-//! ]);
-//! let c1 = StringArray::from(vec![
-//! "Lorem ipsum dolor sit amet",
-//! "consectetur adipiscing elit",
-//! "sed do eiusmod tempor",
-//! ]);
-//! let c2 = PrimitiveArray::<Float64Type>::from(vec![
-//! Some(123.564532),
-//! None,
-//! Some(-556132.25),
-//! ]);
-//! let c3 = PrimitiveArray::<UInt32Type>::from(vec![3, 2, 1]);
-//! let c4 = BooleanArray::from(vec![Some(true), Some(false), None]);
-//!
-//! let batch = RecordBatch::try_new(
-//! Arc::new(schema),
-//! vec![Arc::new(c1), Arc::new(c2), Arc::new(c3), Arc::new(c4)],
-//! )
-//! .unwrap();
-//!
-//! let sw = StringWriter::new();
-//! let mut writer = csv::Writer::new(sw);
-//! writer.write(&batch).unwrap();
-//! }
-//! ```
-
-use std::io::{Error, ErrorKind, Result, Write};
-
-#[derive(Debug)]
-pub struct StringWriter {
- data: String,
-}
-
-impl StringWriter {
- pub fn new() -> Self {
- StringWriter {
- data: String::new(),
- }
- }
-}
-
-impl Default for StringWriter {
- fn default() -> Self {
- Self::new()
- }
-}
-
-impl ToString for StringWriter {
- fn to_string(&self) -> String {
- self.data.clone()
- }
-}
-
-impl Write for StringWriter {
- fn write(&mut self, buf: &[u8]) -> Result<usize> {
- let string = match String::from_utf8(buf.to_vec()) {
- Ok(x) => x,
- Err(e) => {
- return Err(Error::new(ErrorKind::InvalidData, e));
- }
- };
- self.data.push_str(&string);
- Ok(string.len())
- }
-
- fn flush(&mut self) -> Result<()> {
- Ok(())
- }
-}
diff --git a/arrow/src/util/test_util.rs b/arrow/src/util/test_util.rs
deleted file mode 100644
index 261b245..0000000
--- a/arrow/src/util/test_util.rs
+++ /dev/null
@@ -1,208 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-//! Utils to make testing easier
-
-use rand::{rngs::StdRng, Rng, SeedableRng};
-use std::{env, error::Error, fs, io::Write, path::PathBuf};
-
-/// Returns a vector of size `n`, filled with randomly generated bytes.
-pub fn random_bytes(n: usize) -> Vec<u8> {
- let mut result = vec![];
- let mut rng = seedable_rng();
- for _ in 0..n {
- result.push(rng.gen_range(0, 255));
- }
- result
-}
-
-/// Returns fixed seedable RNG
-pub fn seedable_rng() -> StdRng {
- StdRng::seed_from_u64(42)
-}
-
-/// Returns file handle for a temp file in 'target' directory with a provided content
-///
-/// TODO: Originates from `parquet` utils, can be merged in [ARROW-4064]
-pub fn get_temp_file(file_name: &str, content: &[u8]) -> fs::File {
- // build tmp path to a file in "target/debug/testdata"
- let mut path_buf = env::current_dir().unwrap();
- path_buf.push("target");
- path_buf.push("debug");
- path_buf.push("testdata");
- fs::create_dir_all(&path_buf).unwrap();
- path_buf.push(file_name);
-
- // write file content
- let mut tmp_file = fs::File::create(path_buf.as_path()).unwrap();
- tmp_file.write_all(content).unwrap();
- tmp_file.sync_all().unwrap();
-
- // return file handle for both read and write
- let file = fs::OpenOptions::new()
- .read(true)
- .write(true)
- .open(path_buf.as_path());
- assert!(file.is_ok());
- file.unwrap()
-}
-
-/// Returns the arrow test data directory, which is by default stored
-/// in a git submodule rooted at `arrow/testing/data`.
-///
-/// The default can be overridden by the optional environment
-/// variable `ARROW_TEST_DATA`
-///
-/// panics when the directory can not be found.
-///
-/// Example:
-/// ```
-/// let testdata = arrow::util::test_util::arrow_test_data();
-/// let csvdata = format!("{}/csv/aggregate_test_100.csv", testdata);
-/// assert!(std::path::PathBuf::from(csvdata).exists());
-/// ```
-pub fn arrow_test_data() -> String {
- match get_data_dir("ARROW_TEST_DATA", "../testing/data") {
- Ok(pb) => pb.display().to_string(),
- Err(err) => panic!("failed to get arrow data dir: {}", err),
- }
-}
-
-/// Returns the parquest test data directory, which is by default
-/// stored in a git submodule rooted at
-/// `arrow/parquest-testing/data`.
-///
-/// The default can be overridden by the optional environment variable
-/// `PARQUET_TEST_DATA`
-///
-/// panics when the directory can not be found.
-///
-/// Example:
-/// ```
-/// let testdata = arrow::util::test_util::parquet_test_data();
-/// let filename = format!("{}/binary.parquet", testdata);
-/// assert!(std::path::PathBuf::from(filename).exists());
-/// ```
-pub fn parquet_test_data() -> String {
- match get_data_dir("PARQUET_TEST_DATA", "../parquet-testing/data") {
- Ok(pb) => pb.display().to_string(),
- Err(err) => panic!("failed to get parquet data dir: {}", err),
- }
-}
-
-/// Returns a directory path for finding test data.
-///
-/// udf_env: name of an environment variable
-///
-/// submodule_dir: fallback path (relative to CARGO_MANIFEST_DIR)
-///
-/// Returns either:
-/// The path referred to in `udf_env` if that variable is set and refers to a directory
-/// The submodule_data directory relative to CARGO_MANIFEST_PATH
-fn get_data_dir(udf_env: &str, submodule_data: &str) -> Result<PathBuf, Box<dyn Error>> {
- // Try user defined env.
- if let Ok(dir) = env::var(udf_env) {
- let trimmed = dir.trim().to_string();
- if !trimmed.is_empty() {
- let pb = PathBuf::from(trimmed);
- if pb.is_dir() {
- return Ok(pb);
- } else {
- return Err(format!(
- "the data dir `{}` defined by env {} not found",
- pb.display().to_string(),
- udf_env
- )
- .into());
- }
- }
- }
-
- // The env is undefined or its value is trimmed to empty, let's try default dir.
-
- // env "CARGO_MANIFEST_DIR" is "the directory containing the manifest of your package",
- // set by `cargo run` or `cargo test`, see:
- // https://doc.rust-lang.org/cargo/reference/environment-variables.html
- let dir = env!("CARGO_MANIFEST_DIR");
-
- let pb = PathBuf::from(dir).join(submodule_data);
- if pb.is_dir() {
- Ok(pb)
- } else {
- Err(format!(
- "env `{}` is undefined or has empty value, and the pre-defined data dir `{}` not found\n\
- HINT: try running `git submodule update --init`",
- udf_env,
- pb.display().to_string(),
- ).into())
- }
-}
-
-#[cfg(test)]
-mod tests {
- use super::*;
- use std::env;
-
- #[test]
- fn test_data_dir() {
- let udf_env = "get_data_dir";
- let cwd = env::current_dir().unwrap();
-
- let existing_pb = cwd.join("..");
- let existing = existing_pb.display().to_string();
- let existing_str = existing.as_str();
-
- let non_existing = cwd.join("non-existing-dir").display().to_string();
- let non_existing_str = non_existing.as_str();
-
- env::set_var(udf_env, non_existing_str);
- let res = get_data_dir(udf_env, existing_str);
- assert!(res.is_err());
-
- env::set_var(udf_env, "");
- let res = get_data_dir(udf_env, existing_str);
- assert!(res.is_ok());
- assert_eq!(res.unwrap(), existing_pb);
-
- env::set_var(udf_env, " ");
- let res = get_data_dir(udf_env, existing_str);
- assert!(res.is_ok());
- assert_eq!(res.unwrap(), existing_pb);
-
- env::set_var(udf_env, existing_str);
- let res = get_data_dir(udf_env, existing_str);
- assert!(res.is_ok());
- assert_eq!(res.unwrap(), existing_pb);
-
- env::remove_var(udf_env);
- let res = get_data_dir(udf_env, non_existing_str);
- assert!(res.is_err());
-
- let res = get_data_dir(udf_env, existing_str);
- assert!(res.is_ok());
- assert_eq!(res.unwrap(), existing_pb);
- }
-
- #[test]
- fn test_happy() {
- let res = arrow_test_data();
- assert!(PathBuf::from(res).is_dir());
-
- let res = parquet_test_data();
- assert!(PathBuf::from(res).is_dir());
- }
-}
diff --git a/arrow/src/util/trusted_len.rs b/arrow/src/util/trusted_len.rs
deleted file mode 100644
index 84a6623..0000000
--- a/arrow/src/util/trusted_len.rs
+++ /dev/null
@@ -1,82 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-use super::bit_util;
-use crate::{
- buffer::{Buffer, MutableBuffer},
- datatypes::ArrowNativeType,
-};
-
-/// Creates two [`Buffer`]s from an iterator of `Option`.
-/// The first buffer corresponds to a bitmap buffer, the second one
-/// corresponds to a values buffer.
-/// # Safety
-/// The caller must ensure that `iterator` is `TrustedLen`.
-#[inline]
-pub(crate) unsafe fn trusted_len_unzip<I, P, T>(iterator: I) -> (Buffer, Buffer)
-where
- T: ArrowNativeType,
- P: std::borrow::Borrow<Option<T>>,
- I: Iterator<Item = P>,
-{
- let (_, upper) = iterator.size_hint();
- let upper = upper.expect("trusted_len_unzip requires an upper limit");
- let len = upper * std::mem::size_of::<T>();
-
- let mut null = MutableBuffer::from_len_zeroed(upper.saturating_add(7) / 8);
- let mut buffer = MutableBuffer::new(len);
-
- let dst_null = null.as_mut_ptr();
- let mut dst = buffer.as_mut_ptr() as *mut T;
- for (i, item) in iterator.enumerate() {
- let item = item.borrow();
- if let Some(item) = item {
- std::ptr::write(dst, *item);
- bit_util::set_bit_raw(dst_null, i);
- } else {
- std::ptr::write(dst, T::default());
- }
- dst = dst.add(1);
- }
- assert_eq!(
- dst.offset_from(buffer.as_ptr() as *mut T) as usize,
- upper,
- "Trusted iterator length was not accurately reported"
- );
- buffer.set_len(len);
- (null.into(), buffer.into())
-}
-
-#[cfg(test)]
-mod tests {
- use super::*;
-
- #[test]
- fn trusted_len_unzip_good() {
- let vec = vec![Some(1u32), None];
- let (null, buffer) = unsafe { trusted_len_unzip(vec.iter()) };
- assert_eq!(null.as_slice(), &[0b00000001]);
- assert_eq!(buffer.as_slice(), &[1u8, 0, 0, 0, 0, 0, 0, 0]);
- }
-
- #[test]
- #[should_panic(expected = "trusted_len_unzip requires an upper limit")]
- fn trusted_len_unzip_panic() {
- let iter = std::iter::repeat(Some(4i32));
- unsafe { trusted_len_unzip(iter) };
- }
-}
diff --git a/arrow/src/zz_memory_check.rs b/arrow/src/zz_memory_check.rs
deleted file mode 100644
index 70ec8eb..0000000
--- a/arrow/src/zz_memory_check.rs
+++ /dev/null
@@ -1,31 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-// This file is named like this so that it is the last one to be tested
-// It contains no content, it has a single test that verifies that there is no memory leak
-// on all unit-tests
-
-#[cfg(feature = "memory-check")]
-mod tests {
- use crate::memory::ALLOCATIONS;
-
- // verify that there is no data un-allocated
- #[test]
- fn test_memory_check() {
- unsafe { assert_eq!(ALLOCATIONS.load(std::sync::atomic::Ordering::SeqCst), 0) }
- }
-}
diff --git a/arrow/test/data/arrays.json b/arrow/test/data/arrays.json
deleted file mode 100644
index 5dbdd19..0000000
--- a/arrow/test/data/arrays.json
+++ /dev/null
@@ -1,3 +0,0 @@
-{"a":1, "b":[2.0, 1.3, -6.1], "c":[false, true], "d":"4"}
-{"a":-10, "b":[2.0, 1.3, -6.1], "c":[true, true], "d":"4"}
-{"a":2, "b":[2.0, null, -6.1], "c":[false, null], "d":"text"}
diff --git a/arrow/test/data/basic.json b/arrow/test/data/basic.json
deleted file mode 100644
index dafd2dd..0000000
--- a/arrow/test/data/basic.json
+++ /dev/null
@@ -1,12 +0,0 @@
-{"a":1, "b":2.0, "c":false, "d":"4"}
-{"a":-10, "b":-3.5, "c":true, "d":"4"}
-{"a":2, "b":0.6, "c":false, "d":"text"}
-{"a":1, "b":2.0, "c":false, "d":"4"}
-{"a":7, "b":-3.5, "c":true, "d":"4"}
-{"a":1, "b":0.6, "c":false, "d":"text"}
-{"a":1, "b":2.0, "c":false, "d":"4"}
-{"a":5, "b":-3.5, "c":true, "d":"4"}
-{"a":1, "b":0.6, "c":false, "d":"text"}
-{"a":1, "b":2.0, "c":false, "d":"4"}
-{"a":1, "b":-3.5, "c":true, "d":"4"}
-{"a":100000000000000, "b":0.6, "c":false, "d":"text"}
\ No newline at end of file
diff --git a/arrow/test/data/basic_nulls.json b/arrow/test/data/basic_nulls.json
deleted file mode 100644
index 1451df7..0000000
--- a/arrow/test/data/basic_nulls.json
+++ /dev/null
@@ -1,12 +0,0 @@
-{"a":1, "b":2.0, "c":false}
-{"a":null, "b":-3.5, "c":true, "d":"4"}
-{"c":false, "d":"text"}
-{"a":1, "b":2.0, "c":false, "d":"4"}
-{"a":7, "b":-3.5, "c":null, "d":null}
-{"a":1, "b":0.6, "c":false}
-{"a":1, "b":2.0, "d":"4"}
-{"a":5, "c":true}
-{"a":1, "b":0.6, "c":false, "d":"text"}
-{"a":1, "b":2.0, "c":false, "d":"4"}
-{"a":1, "b":-3.5, "c":true, "d":"4"}
-{}
\ No newline at end of file
diff --git a/arrow/test/data/integration.json b/arrow/test/data/integration.json
deleted file mode 100644
index 7e4a22c..0000000
--- a/arrow/test/data/integration.json
+++ /dev/null
@@ -1,808 +0,0 @@
-{
- "schema": {
- "fields": [
- {
- "name": "bools-with-metadata-map",
- "type": {
- "name": "bool"
- },
- "nullable": true,
- "metadata": {
- "k": "v"
- },
- "children": []
- },
- {
- "name": "bools-with-metadata-vec",
- "type": {
- "name": "bool"
- },
- "nullable": true,
- "metadata": [
- {
- "key": "k2",
- "value": "v2"
- }
- ],
- "children": []
- },
- {
- "name": "bools",
- "type": {
- "name": "bool"
- },
- "nullable": true,
- "children": []
- },
- {
- "name": "int8s",
- "type": {
- "name": "int",
- "isSigned": true,
- "bitWidth": 8
- },
- "nullable": true,
- "children": []
- },
- {
- "name": "int16s",
- "type": {
- "name": "int",
- "isSigned": true,
- "bitWidth": 16
- },
- "nullable": true,
- "children": []
- },
- {
- "name": "int32s",
- "type": {
- "name": "int",
- "isSigned": true,
- "bitWidth": 32
- },
- "nullable": true,
- "children": []
- },
- {
- "name": "int64s",
- "type": {
- "name": "int",
- "isSigned": true,
- "bitWidth": 64
- },
- "nullable": true,
- "children": []
- },
- {
- "name": "uint8s",
- "type": {
- "name": "int",
- "isSigned": false,
- "bitWidth": 8
- },
- "nullable": true,
- "children": []
- },
- {
- "name": "uint16s",
- "type": {
- "name": "int",
- "isSigned": false,
- "bitWidth": 16
- },
- "nullable": true,
- "children": []
- },
- {
- "name": "uint32s",
- "type": {
- "name": "int",
- "isSigned": false,
- "bitWidth": 32
- },
- "nullable": true,
- "children": []
- },
- {
- "name": "uint64s",
- "type": {
- "name": "int",
- "isSigned": false,
- "bitWidth": 64
- },
- "nullable": true,
- "children": []
- },
- {
- "name": "float32s",
- "type": {
- "name": "floatingpoint",
- "precision": "SINGLE"
- },
- "nullable": true,
- "children": []
- },
- {
- "name": "float64s",
- "type": {
- "name": "floatingpoint",
- "precision": "DOUBLE"
- },
- "nullable": true,
- "children": []
- },
- {
- "name": "date_days",
- "type": {
- "name": "date",
- "unit": "DAY"
- },
- "nullable": true,
- "children": []
- },
- {
- "name": "date_millis",
- "type": {
- "name": "date",
- "unit": "MILLISECOND"
- },
- "nullable": true,
- "children": []
- },
- {
- "name": "time_secs",
- "type": {
- "name": "time",
- "unit": "SECOND",
- "bitWidth": 32
- },
- "nullable": true,
- "children": []
- },
- {
- "name": "time_millis",
- "type": {
- "name": "time",
- "unit": "MILLISECOND",
- "bitWidth": 32
- },
- "nullable": true,
- "children": []
- },
- {
- "name": "time_micros",
- "type": {
- "name": "time",
- "unit": "MICROSECOND",
- "bitWidth": 64
- },
- "nullable": true,
- "children": []
- },
- {
- "name": "time_nanos",
- "type": {
- "name": "time",
- "unit": "NANOSECOND",
- "bitWidth": 64
- },
- "nullable": true,
- "children": []
- },
- {
- "name": "ts_secs",
- "type": {
- "name": "timestamp",
- "unit": "SECOND"
- },
- "nullable": true,
- "children": []
- },
- {
- "name": "ts_millis",
- "type": {
- "name": "timestamp",
- "unit": "MILLISECOND"
- },
- "nullable": true,
- "children": []
- },
- {
- "name": "ts_micros",
- "type": {
- "name": "timestamp",
- "unit": "MICROSECOND"
- },
- "nullable": true,
- "children": []
- },
- {
- "name": "ts_nanos",
- "type": {
- "name": "timestamp",
- "unit": "NANOSECOND"
- },
- "nullable": true,
- "children": []
- },
- {
- "name": "ts_secs_tz",
- "type": {
- "name": "timestamp",
- "unit": "SECOND",
- "timezone": "Europe/Budapest"
- },
- "nullable": true,
- "children": []
- },
- {
- "name": "ts_millis_tz",
- "type": {
- "name": "timestamp",
- "unit": "MILLISECOND",
- "timezone": "America/New_York"
- },
- "nullable": true,
- "children": []
- },
- {
- "name": "ts_micros_tz",
- "type": {
- "name": "timestamp",
- "unit": "MICROSECOND",
- "timezone": "UTC"
- },
- "nullable": true,
- "children": []
- },
- {
- "name": "ts_nanos_tz",
- "type": {
- "name": "timestamp",
- "unit": "NANOSECOND",
- "timezone": "Africa/Johannesburg"
- },
- "nullable": true,
- "children": []
- },
- {
- "name": "utf8s",
- "type": {
- "name": "utf8"
- },
- "nullable": true,
- "children": []
- },
- {
- "name": "lists",
- "nullable": true,
- "type": {
- "name": "list"
- },
- "children": [
- {
- "name": "item",
- "nullable": true,
- "type": {
- "name": "int",
- "bitWidth": 32,
- "isSigned": true
- },
- "children": []
- }
- ]
- },
- {
- "name": "structs",
- "type": {
- "name": "struct"
- },
- "nullable": true,
- "children": [
- {
- "name": "int32s",
- "type": {
- "name": "int",
- "isSigned": true,
- "bitWidth": 32
- },
- "nullable": true,
- "children": []
- },
- {
- "name": "utf8s",
- "type": {
- "name": "utf8"
- },
- "nullable": true,
- "children": []
- }
- ]
- }
- ]
- },
- "batches": [
- {
- "count": 3,
- "columns": [
- {
- "name": "bools-with-metadata-map",
- "count": 3,
- "VALIDITY": [
- 1,
- 0,
- 1
- ],
- "DATA": [
- true,
- true,
- false
- ]
- },
- {
- "name": "bools-with-metadata-vec",
- "count": 3,
- "VALIDITY": [
- 1,
- 0,
- 1
- ],
- "DATA": [
- true,
- true,
- false
- ]
- },
- {
- "name": "bools",
- "count": 3,
- "VALIDITY": [
- 1,
- 0,
- 1
- ],
- "DATA": [
- true,
- true,
- false
- ]
- },
- {
- "name": "int8s",
- "count": 3,
- "VALIDITY": [
- 1,
- 0,
- 1
- ],
- "DATA": [
- 1,
- 2,
- 3
- ]
- },
- {
- "name": "int16s",
- "count": 3,
- "VALIDITY": [
- 1,
- 0,
- 1
- ],
- "DATA": [
- 1,
- 2,
- 3
- ]
- },
- {
- "name": "int32s",
- "count": 3,
- "VALIDITY": [
- 1,
- 0,
- 1
- ],
- "DATA": [
- 1,
- 2,
- 3
- ]
- },
- {
- "name": "int64s",
- "count": 3,
- "VALIDITY": [
- 1,
- 0,
- 1
- ],
- "DATA": [
- 1,
- 2,
- 3
- ]
- },
- {
- "name": "uint8s",
- "count": 3,
- "VALIDITY": [
- 1,
- 0,
- 1
- ],
- "DATA": [
- 1,
- 2,
- 3
- ]
- },
- {
- "name": "uint16s",
- "count": 5,
- "VALIDITY": [
- 1,
- 0,
- 1
- ],
- "DATA": [
- 1,
- 2,
- 3
- ]
- },
- {
- "name": "uint32s",
- "count": 3,
- "VALIDITY": [
- 1,
- 0,
- 1
- ],
- "DATA": [
- 1,
- 2,
- 3
- ]
- },
- {
- "name": "uint64s",
- "count": 3,
- "VALIDITY": [
- 1,
- 0,
- 1
- ],
- "DATA": [
- 1,
- 2,
- 3
- ]
- },
- {
- "name": "float32s",
- "count": 3,
- "VALIDITY": [
- 1,
- 0,
- 1
- ],
- "DATA": [
- 1.0,
- 2.0,
- 3.0
- ]
- },
- {
- "name": "float64s",
- "count": 3,
- "VALIDITY": [
- 1,
- 0,
- 1
- ],
- "DATA": [
- 1.0,
- 2.0,
- 3.0
- ]
- },
- {
- "name": "date_days",
- "count": 3,
- "VALIDITY": [
- 1,
- 0,
- 0
- ],
- "DATA": [
- 1196848,
- 2319603,
- 2755982
- ]
- },
- {
- "name": "date_millis",
- "count": 3,
- "VALIDITY": [
- 1,
- 1,
- 1
- ],
- "DATA": [
- 167903550396207,
- 29923997007884,
- 30612271819236
- ]
- },
- {
- "name": "time_secs",
- "count": 3,
- "VALIDITY": [
- 1,
- 1,
- 1
- ],
- "DATA": [
- 27974,
- 78592,
- 43207
- ]
- },
- {
- "name": "time_millis",
- "count": 3,
- "VALIDITY": [
- 1,
- 1,
- 1
- ],
- "DATA": [
- 6613125,
- 74667230,
- 52260079
- ]
- },
- {
- "name": "time_micros",
- "count": 3,
- "VALIDITY": [
- 1,
- 0,
- 0
- ],
- "DATA": [
- 62522958593,
- 13470380050,
- 50797036705
- ]
- },
- {
- "name": "time_nanos",
- "count": 3,
- "VALIDITY": [
- 1,
- 0,
- 1
- ],
- "DATA": [
- 73380123595985,
- 52520995325145,
- 16584393546415
- ]
- },
- {
- "name": "ts_secs",
- "count": 3,
- "VALIDITY": [
- 0,
- 1,
- 0
- ],
- "DATA": [
- 209869064422,
- 193438817552,
- 51757838205
- ]
- },
- {
- "name": "ts_millis",
- "count": 3,
- "VALIDITY": [
- 0,
- 1,
- 1
- ],
- "DATA": [
- 228315043570185,
- 38606916383008,
- 58113709376587
- ]
- },
- {
- "name": "ts_micros",
- "count": 3,
- "VALIDITY": [
- 0,
- 0,
- 0
- ],
- "DATA": [
- 133457416537791415,
- 129522736067409280,
- 177110451066832967
- ]
- },
- {
- "name": "ts_nanos",
- "count": 3,
- "VALIDITY": [
- 0,
- 0,
- 1
- ],
- "DATA": [
- -804525722984600007,
- 8166038652634779458,
- -6473623571954960143
- ]
- },
- {
- "name": "ts_secs_tz",
- "count": 3,
- "VALIDITY": [
- 0,
- 1,
- 0
- ],
- "DATA": [
- 209869064422,
- 193438817552,
- 51757838205
- ]
- },
- {
- "name": "ts_millis_tz",
- "count": 3,
- "VALIDITY": [
- 0,
- 1,
- 1
- ],
- "DATA": [
- 228315043570185,
- 38606916383008,
- 58113709376587
- ]
- },
- {
- "name": "ts_micros_tz",
- "count": 3,
- "VALIDITY": [
- 0,
- 0,
- 0
- ],
- "DATA": [
- 133457416537791415,
- 129522736067409280,
- 177110451066832967
- ]
- },
- {
- "name": "ts_nanos_tz",
- "count": 3,
- "VALIDITY": [
- 0,
- 0,
- 1
- ],
- "DATA": [
- -804525722984600007,
- 8166038652634779458,
- -6473623571954960143
- ]
- },
- {
- "name": "utf8s",
- "count": 3,
- "VALIDITY": [
- 1,
- 0,
- 1
- ],
- "OFFSET": [
- 0,
- 2,
- 2,
- 5
- ],
- "DATA": [
- "aa",
- "",
- "bbb"
- ]
- },
- {
- "name": "lists",
- "count": 3,
- "VALIDITY": [
- 1,
- 1,
- 0
- ],
- "OFFSET": [
- 0,
- 3,
- 4,
- 4
- ],
- "children": [
- {
- "name": "item",
- "count": 4,
- "VALIDITY": [
- 0,
- 1,
- 0,
- 0
- ],
- "DATA": [
- 1,
- 2,
- 3,
- 4
- ]
- }
- ]
- },
- {
- "name": "structs",
- "count": 3,
- "VALIDITY": [
- 1,
- 1,
- 0
- ],
- "children": [
- {
- "name": "int32s",
- "count": 3,
- "VALIDITY": [
- 0,
- 1,
- 0
- ],
- "DATA": [
- -1,
- -2,
- -3
- ]
- },
- {
- "name": "utf8s",
- "count": 3,
- "VALIDITY": [
- 0,
- 0,
- 1
- ],
- "OFFSET": [
- 0,
- 0,
- 0,
- 7
- ],
- "DATA": [
- "",
- "",
- "aaaaaa"
- ]
- }
- ]
- }
- ]
- }
- ]
-}
\ No newline at end of file
diff --git a/arrow/test/data/list_string_dict_nested.json b/arrow/test/data/list_string_dict_nested.json
deleted file mode 100644
index d215b31..0000000
--- a/arrow/test/data/list_string_dict_nested.json
+++ /dev/null
@@ -1,3 +0,0 @@
-{"machine": "a", "events": ["Elect Leader", "Do Ballot"]}
-{"machine": "b", "events": ["Do Ballot", "Send Data", "Elect Leader"]}
-{"machine": "c", "events": ["Send Data"]}
diff --git a/arrow/test/data/list_string_dict_nested_nulls.json b/arrow/test/data/list_string_dict_nested_nulls.json
deleted file mode 100644
index 9300b14..0000000
--- a/arrow/test/data/list_string_dict_nested_nulls.json
+++ /dev/null
@@ -1,3 +0,0 @@
-{"machine": "a", "events": [null, "Elect Leader", "Do Ballot"]}
-{"machine": "b", "events": ["Do Ballot", null, "Send Data", "Elect Leader"]}
-{"machine": "c", "events": ["Send Data"]}
diff --git a/arrow/test/data/mixed_arrays.json b/arrow/test/data/mixed_arrays.json
deleted file mode 100644
index 1898728..0000000
--- a/arrow/test/data/mixed_arrays.json
+++ /dev/null
@@ -1,4 +0,0 @@
-{"a":1, "b":[2.0, 1.3, -6.1], "c":[false, true], "d":4.1}
-{"a":-10, "b":[2.0, 1.3, -6.1], "c":null, "d":null}
-{"a":2, "b":[2.0, null, -6.1], "c":[false, null], "d":"text"}
-{"a":3, "b":4, "c": true, "d":[1, false, "array", 2.4]}
diff --git a/arrow/test/data/mixed_arrays.json.gz b/arrow/test/data/mixed_arrays.json.gz
deleted file mode 100644
index 0f60400..0000000
--- a/arrow/test/data/mixed_arrays.json.gz
+++ /dev/null
Binary files differ
diff --git a/arrow/test/data/nested_structs.json b/arrow/test/data/nested_structs.json
deleted file mode 100644
index 32a3ac8..0000000
--- a/arrow/test/data/nested_structs.json
+++ /dev/null
@@ -1,4 +0,0 @@
-{"a": {"b": true, "c": {"d": "text"}}}
-{"a": {"b": false, "c": null}}
-{"a": {"b": true, "c": {"d": "text"}}}
-{"a": 1}
\ No newline at end of file
diff --git a/arrow/test/data/null_test.csv b/arrow/test/data/null_test.csv
deleted file mode 100644
index 7e0dde5..0000000
--- a/arrow/test/data/null_test.csv
+++ /dev/null
@@ -1,6 +0,0 @@
-c_int,c_float,c_string,c_bool
-1,1.1,"1.11",True
-2,2.2,"2.22",TRUE
-3,,"3.33",true
-4,4.4,,False
-5,6.6,"",FALSE
\ No newline at end of file
diff --git a/arrow/test/data/uk_cities.csv b/arrow/test/data/uk_cities.csv
deleted file mode 100644
index db9e6da..0000000
--- a/arrow/test/data/uk_cities.csv
+++ /dev/null
@@ -1,37 +0,0 @@
-"Elgin, Scotland, the UK",57.653484,-3.335724
-"Stoke-on-Trent, Staffordshire, the UK",53.002666,-2.179404
-"Solihull, Birmingham, UK",52.412811,-1.778197
-"Cardiff, Cardiff county, UK",51.481583,-3.179090
-"Eastbourne, East Sussex, UK",50.768036,0.290472
-"Oxford, Oxfordshire, UK",51.752022,-1.257677
-"London, UK",51.509865,-0.118092
-"Swindon, Swindon, UK",51.568535,-1.772232
-"Gravesend, Kent, UK",51.441883,0.370759
-"Northampton, Northamptonshire, UK",52.240479,-0.902656
-"Rugby, Warwickshire, UK",52.370876,-1.265032
-"Sutton Coldfield, West Midlands, UK",52.570385,-1.824042
-"Harlow, Essex, UK",51.772938,0.102310
-"Aberdeen, Aberdeen City, UK",57.149651,-2.099075
-"Swansea, Swansea, UK",51.621441,-3.943646
-"Chesterfield, Derbyshire, UK",53.235046,-1.421629
-"Londonderry, Derry, UK",55.006763,-7.318268
-"Salisbury, Wiltshire, UK",51.068787,-1.794472
-"Weymouth, Dorset, UK",50.614429,-2.457621
-"Wolverhampton, West Midlands, UK",52.591370,-2.110748
-"Preston, Lancashire, UK",53.765762,-2.692337
-"Bournemouth, UK",50.720806,-1.904755
-"Doncaster, South Yorkshire, UK",53.522820,-1.128462
-"Ayr, South Ayrshire, UK",55.458565,-4.629179
-"Hastings, East Sussex, UK",50.854259,0.573453
-"Bedford, UK",52.136436,-0.460739
-"Basildon, Essex, UK",51.572376,0.470009
-"Chippenham, Wiltshire, UK",51.458057,-2.116074
-"Belfast, UK",54.607868,-5.926437
-"Uckfield, East Sussex, UK",50.967941,0.085831
-"Worthing, West Sussex, UK",50.825024,-0.383835
-"Leeds, West Yorkshire, UK",53.801277,-1.548567
-"Kendal, Cumbria, UK",54.328506,-2.743870
-"Plymouth, UK",50.376289,-4.143841
-"Haverhill, Suffolk, UK",52.080875,0.444517
-"Frankton, Warwickshire, UK",52.328415,-1.377561
-"Inverness, the UK",57.477772,-4.224721
\ No newline at end of file
diff --git a/arrow/test/data/uk_cities_with_headers.csv b/arrow/test/data/uk_cities_with_headers.csv
deleted file mode 100644
index 92f5a17..0000000
--- a/arrow/test/data/uk_cities_with_headers.csv
+++ /dev/null
@@ -1,38 +0,0 @@
-city,lat,lng
-"Elgin, Scotland, the UK",57.653484,-3.335724
-"Stoke-on-Trent, Staffordshire, the UK",53.002666,-2.179404
-"Solihull, Birmingham, UK",52.412811,-1.778197
-"Cardiff, Cardiff county, UK",51.481583,-3.179090
-"Eastbourne, East Sussex, UK",50.768036,0.290472
-"Oxford, Oxfordshire, UK",51.752022,-1.257677
-"London, UK",51.509865,-0.118092
-"Swindon, Swindon, UK",51.568535,-1.772232
-"Gravesend, Kent, UK",51.441883,0.370759
-"Northampton, Northamptonshire, UK",52.240479,-0.902656
-"Rugby, Warwickshire, UK",52.370876,-1.265032
-"Sutton Coldfield, West Midlands, UK",52.570385,-1.824042
-"Harlow, Essex, UK",51.772938,0.102310
-"Aberdeen, Aberdeen City, UK",57.149651,-2.099075
-"Swansea, Swansea, UK",51.621441,-3.943646
-"Chesterfield, Derbyshire, UK",53.235046,-1.421629
-"Londonderry, Derry, UK",55.006763,-7.318268
-"Salisbury, Wiltshire, UK",51.068787,-1.794472
-"Weymouth, Dorset, UK",50.614429,-2.457621
-"Wolverhampton, West Midlands, UK",52.591370,-2.110748
-"Preston, Lancashire, UK",53.765762,-2.692337
-"Bournemouth, UK",50.720806,-1.904755
-"Doncaster, South Yorkshire, UK",53.522820,-1.128462
-"Ayr, South Ayrshire, UK",55.458565,-4.629179
-"Hastings, East Sussex, UK",50.854259,0.573453
-"Bedford, UK",52.136436,-0.460739
-"Basildon, Essex, UK",51.572376,0.470009
-"Chippenham, Wiltshire, UK",51.458057,-2.116074
-"Belfast, UK",54.607868,-5.926437
-"Uckfield, East Sussex, UK",50.967941,0.085831
-"Worthing, West Sussex, UK",50.825024,-0.383835
-"Leeds, West Yorkshire, UK",53.801277,-1.548567
-"Kendal, Cumbria, UK",54.328506,-2.743870
-"Plymouth, UK",50.376289,-4.143841
-"Haverhill, Suffolk, UK",52.080875,0.444517
-"Frankton, Warwickshire, UK",52.328415,-1.377561
-"Inverness, the UK",57.477772,-4.224721
\ No newline at end of file
diff --git a/arrow/test/data/various_types.csv b/arrow/test/data/various_types.csv
deleted file mode 100644
index 8f4466f..0000000
--- a/arrow/test/data/various_types.csv
+++ /dev/null
@@ -1,6 +0,0 @@
-c_int|c_float|c_string|c_bool|c_date|c_datetime
-1|1.1|"1.11"|true|1970-01-01|1970-01-01T00:00:00
-2|2.2|"2.22"|true|2020-11-08|2020-11-08T01:00:00
-3||"3.33"|true|1969-12-31|1969-11-08T02:00:00
-4|4.4||false||
-5|6.6|""|false|1990-01-01|1990-01-01T03:00:00
\ No newline at end of file
diff --git a/arrow/test/data/various_types_invalid.csv b/arrow/test/data/various_types_invalid.csv
deleted file mode 100644
index 6f059cb..0000000
--- a/arrow/test/data/various_types_invalid.csv
+++ /dev/null
@@ -1,6 +0,0 @@
-c_int|c_float|c_string|c_bool
-1|1.1|"1.11"|true
-2|2.2|"2.22"|true
-3||"3.33"|true
-4|4.x4||false
-5|6.6|""|false
\ No newline at end of file
diff --git a/dev/.gitignore b/dev/.gitignore
deleted file mode 100644
index c03a7c7..0000000
--- a/dev/.gitignore
+++ /dev/null
@@ -1,22 +0,0 @@
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-
-# Python virtual environments for dev tools
-.venv*/
-
-__pycache__
-*.egg-info
diff --git a/dev/README.md b/dev/README.md
deleted file mode 100644
index b4ea02b..0000000
--- a/dev/README.md
+++ /dev/null
@@ -1,57 +0,0 @@
-<!--
- ~ Licensed to the Apache Software Foundation (ASF) under one
- ~ or more contributor license agreements. See the NOTICE file
- ~ distributed with this work for additional information
- ~ regarding copyright ownership. The ASF licenses this file
- ~ to you under the Apache License, Version 2.0 (the
- ~ "License"); you may not use this file except in compliance
- ~ with the License. You may obtain a copy of the License at
- ~
- ~ http://www.apache.org/licenses/LICENSE-2.0
- ~
- ~ Unless required by applicable law or agreed to in writing,
- ~ software distributed under the License is distributed on an
- ~ "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
- ~ KIND, either express or implied. See the License for the
- ~ specific language governing permissions and limitations
- ~ under the License.
- -->
-
-# Arrow Developer Scripts
-
-This directory contains scripts useful to developers when packaging,
-testing, or committing to Arrow.
-
-## Verifying Release Candidates
-
-We have provided a script to assist with verifying release candidates:
-
-```shell
-bash dev/release/verify-release-candidate.sh 0.7.0 0
-```
-
-Currently this only works on Linux (patches to expand to macOS welcome!). Read
-the script for information about system dependencies.
-
-On Windows, we have a script that verifies C++ and Python (requires Visual
-Studio 2015):
-
-```
-dev/release/verify-release-candidate.bat apache-arrow-0.7.0.tar.gz
-```
-
-### Verifying the JavaScript release
-
-For JavaScript-specific releases, use a different verification script:
-
-```shell
-bash dev/release/js-verify-release-candidate.sh 0.7.0 0
-```
-
-# Integration testing
-
-Build the following base image used by multiple tests:
-
-```shell
-docker build -t arrow_integration_xenial_base -f docker_common/Dockerfile.xenial.base .
-```
diff --git a/dev/release/README.md b/dev/release/README.md
deleted file mode 100644
index 5b31083..0000000
--- a/dev/release/README.md
+++ /dev/null
@@ -1,266 +0,0 @@
-<!---
- Licensed to the Apache Software Foundation (ASF) under one
- or more contributor license agreements. See the NOTICE file
- distributed with this work for additional information
- regarding copyright ownership. The ASF licenses this file
- to you under the Apache License, Version 2.0 (the
- "License"); you may not use this file except in compliance
- with the License. You may obtain a copy of the License at
-
- http://www.apache.org/licenses/LICENSE-2.0
-
- Unless required by applicable law or agreed to in writing,
- software distributed under the License is distributed on an
- "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
- KIND, either express or implied. See the License for the
- specific language governing permissions and limitations
- under the License.
--->
-
-# Release Process
-
-## Branching
-
-We would maintain two branches: `active_release` and `master`.
-
-- All new PRs are created and merged against `master`
-- All versions are created from the `active_release` branch
-- Once merged to master, changes are "cherry-picked" (via a hopefully soon to be automated process), to the `active_release` branch based on the judgement of the original PR author and maintainers.
-
-- We do not merge breaking api changes, as defined in [Rust RFC 1105](https://github.com/rust-lang/rfcs/blob/master/text/1105-api-evolution.md) to the `active_release`
-
-Please see the [original proposal](https://docs.google.com/document/d/1tMQ67iu8XyGGZuj--h9WQYB9inCk6c2sL_4xMTwENGc/edit?ts=60961758) document the rational of this change.
-
-## Release Branching
-
-We aim to release every other week from the `active_release` branch.
-
-Every other Monday, a maintainer proposes a minor (e.g. `4.1.0` to `4.2.0`) or patch (e.g `4.1.0` to `4.1.1`) release, depending on changes to the `active_release` in the previous 2 weeks, following the process beloe.
-
-If this release is approved by at least three PMC members, a new version from that tarball is released to crates.io later in the week.
-
-Apache Arrow in general does synchronized major releases every three months. The Rust implementation aims to do its major releases in the same time frame.
-
-# Release Mechanics
-
-This directory contains the scripts used to manage an Apache Arrow Release.
-
-# Process Overview
-
-As part of the Apache governance model, official releases consist of
-signed source tarballs approved by the PMC.
-
-We then use the code in the approved source tarball to release to
-crates.io, the Rust ecosystem's package manager.
-
-## Branching
-
-# Release Preparation
-
-# Change Log
-
-We create a `CHANGELOG.md` so our users know what has been changed between releases.
-
-The CHANGELOG is created automatically using
-[change_log.sh](https://github.com/apache/arrow-rs/blob/master/change_log.sh)
-
-This script creates a changelog using github issues and the
-labels associated with them.
-
-# Mechanics of creating a release
-
-## Prepare the release branch and tags
-
-First, ensure that `active_release` contains the content of the desired release. For minor and patch releases, no additional steps are needed.
-
-To prepare for _a major release_, change `active release` to point at the latest `master` with commands such as:
-
-```
-git checkout active_release
-git fetch apache
-git reset --hard apache/master
-git push -f
-```
-
-### Update CHANGELOG.md + Version
-
-Now prepare a PR to update `CHANGELOG.md` and versions on `active_release` branch to reflect the planned release.
-
-See [#298](https://github.com/apache/arrow-rs/pull/298) for an example.
-
-Here are the commands used to prepare the 4.1.0 release:
-
-```bash
-git checkout active_release
-git pull
-git checkout -b make-release
-
-# manully edit ./dev/release/update_change_log.sh to reflect the release version
-# create the changelog
-CHANGELOG_GITHUB_TOKEN=<TOKEN> ./dev/release/update_change_log.sh
-# review change log / edit issues and labels if needed, rerun
-git commit -a -m 'Create changelog'
-
-# update versions
-sed -i '' -e 's/5.0.0-SNAPSHOT/4.1.0/g' `find . -name 'Cargo.toml'`
-git commit -a -m 'Update version'
-```
-
-Note that when reviewing the change log, rather than editing the
-`CHANGELOG.md`, it is preferred to update the issues and their labels
-(e.g. add `invalid` label to exclude them from release notes)
-
-## Prepare release candidate tarball
-
-(Note you need to be a committer to run these scripts as they upload to the apache svn distribution servers)
-
-### Create git tag for the release:
-
-While the official release artifact is a signed tarball, we also tag the commit it was created for convenience and code archaeology.
-
-Using a string such as `4.0.1` as the `<version>`, create and push the tag thusly:
-
-```shell
-git fetch apache
-git tag <version> apache/active_release
-# push tag to apache
-git push apache <version>
-```
-
-### Pick an Release Candidate (RC) number
-
-Pick numbers in sequential order, with `1` for `rc1`, `2` for `rc2`, etc.
-
-### Create, sign, and upload tarball
-
-Run the `create-tarball.sh` with the `<version>` tag and `<rc>` and you found in previous steps:
-
-```shell
-./dev/release/create-tarball.sh 4.1.0 2
-```
-
-This script
-
-1. creates and uploads a release candidate tarball to the [arrow
- dev](https://dist.apache.org/repos/dist/dev/arrow) location on the
- apache distribution svn server
-
-2. provide you an email template to
- send to dev@arrow.apache.org for release voting.
-
-### Vote on Release Candidate tarball
-
-Send the email output from the script to dev@arrow.apache.org. The email should look like
-
-```
-To: dev@arrow.apache.org
-Subject: [VOTE][RUST] Release Apache Arrow
-
-Hi,
-
-I would like to propose a release of Apache Arrow Rust
-Implementation, version 4.1.0.
-
-This release candidate is based on commit: a5dd428f57e62db20a945e8b1895de91405958c4 [1]
-
-The proposed release tarball and signatures are hosted at [2].
-The changelog is located at [3].
-
-Please download, verify checksums and signatures, run the unit tests,
-and vote on the release.
-
-The vote will be open for at least 72 hours.
-
-[ ] +1 Release this as Apache Arrow Rust
-[ ] +0
-[ ] -1 Do not release this as Apache Arrow Rust because...
-
-[1]: https://github.com/apache/arrow-rs/tree/a5dd428f57e62db20a945e8b1895de91405958c4
-[2]: https://dist.apache.org/repos/dist/dev/arrow/apache-arrow-rs-4.1.0
-[3]: https://github.com/apache/arrow-rs/blob/a5dd428f57e62db20a945e8b1895de91405958c4/CHANGELOG.md
-```
-
-For the release to become "official" it needs at least three PMC members to vote +1 on it.
-
-#### Verifying Release Candidates
-
-There is a script in this repository which can be used to help `dev/release/verify-release-candidate.sh` assist the verification process. Run it like:
-
-```
-./dev/release/verify-release-candidate.sh 4.1.0 2
-```
-
-#### If the release is not approved
-
-If the release is not approved, fix whatever the problem is and try again with the next RC number
-
-### If the release is approved,
-
-Move tarball to the release location in SVN, e.g. https://dist.apache.org/repos/dist/release/arrow/arrow-4.1.0/, using the `release-tarball.sh` script:
-
-```shell
-./dev/release/release-tarball.sh 4.1.0 2
-```
-
-### Publish on Crates.io
-
-Only approved releases of the tarball should be published to
-crates.io, in order to conform to Apache Software Foundation
-governance standards.
-
-An Arrow committer can publish this crate after an official project release has
-been made to crates.io using the following instructions.
-
-Follow [these
-instructions](https://doc.rust-lang.org/cargo/reference/publishing.html) to
-create an account and login to crates.io before asking to be added as an owner
-of the [arrow crate](https://crates.io/crates/arrow).
-
-Download and unpack the official release tarball
-
-Verify that the Cargo.toml in the tarball contains the correct version
-(e.g. `version = "0.11.0"`) and then publish the crate with the
-following commands
-
-```shell
-(cd arrow && cargo publish)
-(cd arrow_flight && cargo publish)
-(cd parquet && cargo publish)
-(cd parquet_derive && cargo publish)
-```
-
-# Backporting
-
-As of now, the plan for backporting to `active_release` is to do so semi-manually.
-
-Step 1: Pick the commit to cherry-pick
-
-Step 2: Create cherry-pick PR to active_release
-
-Step 3a: If CI passes, merge cherry-pick PR
-
-Step 3b: If CI doesn't pass or some other changes are needed, the PR should be reviewed / approved as normal prior to merge
-
-For example, to backport `b2de5446cc1e45a0559fb39039d0545df1ac0d26` to active_release use the folliwing
-
-```shell
-git clone git@github.com:apache/arrow-rs.git /tmp/arrow-rs
-
-ARROW_GITHUB_API_TOKEN=$ARROW_GITHUB_API_TOKEN CHECKOUT_ROOT=/tmp/arrow-rs CHERRY_PICK_SHA=b2de5446cc1e45a0559fb39039d0545df1ac0d26 python3 dev/release/cherry-pick-pr.py
-```
-
-## Tags
-
-There are two tags that help keep track of backporting:
-
-1. [`cherry-picked`](https://github.com/apache/arrow-rs/labels/cherry-picked) for PRs that have been cherry-picked/backported to `active_release`
-2. [`release-cherry-pick`](https://github.com/apache/arrow-rs/labels/release-cherry-pick) for the PRs that are the cherry pick
-
-You can find candidates to cherry pick using [this filter](https://github.com/apache/arrow-rs/pulls?q=is%3Apr+is%3Aclosed+-label%3Arelease-cherry-pick+-label%3Acherry-picked)
-
-## Rationale for creating PRs:
-
-1. PRs are a natural place to run the CI tests to make sure there are no logical conflicts
-2. PRs offer a place for the original author / committers to comment and say it should/should not be backported.
-3. PRs offer a way to make cleanups / fixups and approve (if needed) for non cherry pick PRs
-4. There is an additional control / review when the candidate release is created
diff --git a/dev/release/check-rat-report.py b/dev/release/check-rat-report.py
deleted file mode 100644
index e30d72b..0000000
--- a/dev/release/check-rat-report.py
+++ /dev/null
@@ -1,59 +0,0 @@
-#!/usr/bin/python
-##############################################################################
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT 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)
diff --git a/dev/release/cherry-pick-pr.py b/dev/release/cherry-pick-pr.py
deleted file mode 100755
index 2886a0d..0000000
--- a/dev/release/cherry-pick-pr.py
+++ /dev/null
@@ -1,154 +0,0 @@
-#!/usr/bin/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.
-##############################################################################
-
-# This script is designed to create a cherry pick PR to a target branch
-#
-# Usage: python3 cherry_pick_pr.py
-#
-# To test locally:
-#
-# git clone git@github.com:apache/arrow-rs.git /tmp/arrow-rs
-#
-# pip3 install PyGithub
-# ARROW_GITHUB_API_TOKEN=<..>
-# CHECKOUT_ROOT=<path>
-# CHERRY_PICK_SHA=<sha> python3 cherry-pick-pr.py
-#
-import os
-import sys
-import six
-import subprocess
-
-from pathlib import Path
-
-TARGET_BRANCH = 'active_release'
-TARGET_REPO = 'apache/arrow-rs'
-
-p = Path(__file__)
-
-# Use github workspace if specified
-repo_root = os.environ.get("CHECKOUT_ROOT")
-if repo_root is None:
- print("arrow-rs checkout must be supplied in CHECKOUT_ROOT environment")
- sys.exit(1)
-
-print("Using checkout in {}".format(repo_root))
-
-token = os.environ.get('ARROW_GITHUB_API_TOKEN', None)
-if token is None:
- print("GITHUB token must be supplied in ARROW_GITHUB_API_TOKEN environmet")
- sys.exit(1)
-
-new_sha = os.environ.get('CHERRY_PICK_SHA', None)
-if new_sha is None:
- print("SHA to cherry pick must be supplied in CHERRY_PICK_SHA environment")
- sys.exit(1)
-
-
-# from merge_pr.py from arrow repo
-def run_cmd(cmd):
- if isinstance(cmd, six.string_types):
- cmd = cmd.split(' ')
- try:
- output = subprocess.check_output(cmd)
- except subprocess.CalledProcessError as e:
- # this avoids hiding the stdout / stderr of failed processes
- print('Command failed: %s' % cmd)
- print('With output:')
- print('--------------')
- print(e.output)
- print('--------------')
- raise e
-
- if isinstance(output, six.binary_type):
- output = output.decode('utf-8')
-
- return output
-
-
-os.chdir(repo_root)
-new_sha_short = run_cmd("git rev-parse --short {}".format(new_sha)).strip()
-new_branch = 'cherry_pick_{}'.format(new_sha_short)
-
-
-def make_cherry_pick():
- if os.environ.get('GITHUB_SHA', None) is not None:
- print("Running on github runner, setting email/username")
- run_cmd(['git', 'config', 'user.email', 'dev@arrow.apache.com'])
- run_cmd(['git', 'config', 'user.name', 'Arrow-RS Automation'])
-
- #
- # Create a new branch from active_release
- # and cherry pick to there.
- #
-
- print("Creating cherry pick from {} to {}".format(
- new_sha_short, new_branch
- ))
-
- # The following tortured dance is required due to how the github
- # actions/checkout works (it doesn't pull other branches and pulls
- # only one commit back)
-
- # pull 10 commits back so we can get the proper cherry pick
- # (probably only need 2 but 10 must be better, right?)
- run_cmd(['git', 'fetch', '--depth', '10', 'origin', 'master'])
- run_cmd(['git', 'fetch', 'origin', 'active_release'])
- run_cmd(['git', 'checkout', '-b', new_branch])
- run_cmd(['git', 'reset', '--hard', 'origin/active_release'])
- run_cmd(['git', 'cherry-pick', new_sha])
- run_cmd(['git', 'push', '-u', 'origin', new_branch])
-
-
-
-def make_cherry_pick_pr():
- from github import Github
- g = Github(token)
- repo = g.get_repo(TARGET_REPO)
-
- release_cherry_pick_label = repo.get_label('release-cherry-pick')
- cherry_picked_label = repo.get_label('cherry-picked')
-
- # Default titles
- new_title = 'Cherry pick {} to active_release'.format(new_sha)
- new_commit_message = 'Automatic cherry-pick of {}\n'.format(new_sha)
-
- # try and get info from github api
- commit = repo.get_commit(new_sha)
- for orig_pull in commit.get_pulls():
- new_commit_message += '* Originally appeared in {}: {}\n'.format(
- orig_pull.html_url, orig_pull.title)
- new_title = 'Cherry pick {} to active_release'.format(orig_pull.title)
- orig_pull.add_to_labels(cherry_picked_label)
-
- pr = repo.create_pull(title=new_title,
- body=new_commit_message,
- base='refs/heads/active_release',
- head='refs/heads/{}'.format(new_branch),
- maintainer_can_modify=True,
- )
-
- pr.add_to_labels(release_cherry_pick_label)
-
- print('Created PR {}'.format(pr.html_url))
-
-
-make_cherry_pick()
-make_cherry_pick_pr()
diff --git a/dev/release/create-tarball.sh b/dev/release/create-tarball.sh
deleted file mode 100755
index 2311519..0000000
--- a/dev/release/create-tarball.sh
+++ /dev/null
@@ -1,123 +0,0 @@
-#!/bin/bash
-#
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-#
-
-
-# This script creates a signed tarball in
-# dev/dist/apache-arrow-rs-<version>-<sha>.tar.gz and uploads it to
-# the "dev" area of the dist.apache.arrow repository and prepares an
-# email for sending to the dev@arrow.apache.org list for a formal
-# vote.
-#
-# See release/README.md for full release instructions
-#
-# Requirements:
-#
-# 1. gpg setup for signing and have uploaded your public
-# signature to https://pgp.mit.edu/
-#
-# 2. Logged into the apache svn server with the appropriate
-# credentials
-#
-#
-# Based in part on 02-source.sh from apache/arrow
-#
-
-set -e
-
-SOURCE_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
-SOURCE_TOP_DIR="$(cd "${SOURCE_DIR}/../../" && pwd)"
-
-if [ "$#" -ne 2 ]; then
- echo "Usage: $0 <tag> <rc>"
- echo "ex. $0 4.1.0 2"
- exit
-fi
-
-tag=$1
-rc=$2
-
-release_hash=$(cd "${SOURCE_TOP_DIR}" && git rev-list --max-count=1 ${tag})
-
-release=apache-arrow-rs-${tag}
-distdir=${SOURCE_TOP_DIR}/dev/dist/${release}-rc${rc}
-tarname=${release}.tar.gz
-tarball=${distdir}/${tarname}
-url="https://dist.apache.org/repos/dist/dev/arrow/${release}-rc${rc}"
-
-echo "Attempting to create ${tarball} from tag ${tag}"
-
-
-if [ -z "$release_hash" ]; then
- echo "Cannot continue: unknown git tag: $tag"
-fi
-
-echo "Draft email for dev@arrow.apache.org mailing list"
-echo ""
-echo "---------------------------------------------------------"
-cat <<MAIL
-To: dev@arrow.apache.org
-Subject: [VOTE][RUST] Release Apache Arrow Rust ${tag} RC${rc}
-
-Hi,
-
-I would like to propose a release of Apache Arrow Rust Implementation, version ${tag}.
-
-This release candidate is based on commit: ${release_hash} [1]
-
-The proposed release tarball and signatures are hosted at [2].
-
-The changelog is located at [3].
-
-Please download, verify checksums and signatures, run the unit tests,
-and vote on the release.
-
-The vote will be open for at least 72 hours.
-
-[ ] +1 Release this as Apache Arrow Rust ${version}
-[ ] +0
-[ ] -1 Do not release this as Apache Arrow Rust ${version} because...
-
-[1]: https://github.com/apache/arrow-rs/tree/${release_hash}
-[2]: ${url}
-[3]: https://github.com/apache/arrow-rs/blob/${release_hash}/CHANGELOG.md
-MAIL
-echo "---------------------------------------------------------"
-
-
-# create <tarball> containing the files in git at $release_hash
-# the files in the tarball are prefixed with {tag} (e.g. 4.0.1)
-mkdir -p ${distdir}
-(cd "${SOURCE_TOP_DIR}" && git archive ${release_hash} --prefix ${release}/ | gzip > ${tarball})
-
-echo "Running rat license checker on ${tarball}"
-${SOURCE_DIR}/run-rat.sh ${tarball}
-
-echo "Signing tarball and creating checksums"
-gpg --armor --output ${tarball}.asc --detach-sig ${tarball}
-# create signing with relative path of tarball
-# so that they can be verified with a command such as
-# shasum --check apache-arrow-rs-4.1.0-rc2.tar.gz.sha512
-(cd ${distdir} && shasum -a 256 ${tarname}) > ${tarball}.sha256
-(cd ${distdir} && shasum -a 512 ${tarname}) > ${tarball}.sha512
-
-echo "Uploading to apache dist/dev to ${url}"
-svn co --depth=empty https://dist.apache.org/repos/dist/dev/arrow ${SOURCE_TOP_DIR}/dev/dist
-svn add ${distdir}
-svn ci -m "Apache Arrow Rust ${tag} ${rc}" ${distdir}
diff --git a/dev/release/rat_exclude_files.txt b/dev/release/rat_exclude_files.txt
deleted file mode 100644
index d2b5aad..0000000
--- a/dev/release/rat_exclude_files.txt
+++ /dev/null
@@ -1,27 +0,0 @@
-*.npmrc
-*.gitignore
-.gitmodules
-*.csv
-*.json
-*.snap
-.github/ISSUE_TEMPLATE/*.md
-.github/pull_request_template.md
-CHANGELOG.md
-pax_global_header
-MANIFEST.in
-requirements.txt
-*.html
-*.sgml
-*.css
-*.png
-*.ico
-*.svg
-*.devhelp2
-*.scss
-.gitattributes
-arrow/test/data/*.csv
-rust-toolchain
-arrow-flight/src/arrow.flight.protocol.rs
-ballista/rust/benchmarks/tpch/queries/q*.sql
-ballista/rust/scheduler/testdata/*
-ballista/ui/scheduler/yarn.lock
diff --git a/dev/release/release-tarball.sh b/dev/release/release-tarball.sh
deleted file mode 100755
index 9612921..0000000
--- a/dev/release/release-tarball.sh
+++ /dev/null
@@ -1,72 +0,0 @@
-#!/bin/bash
-#
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-#
-
-# This script copies a tarball from the "dev" area of the
-# dist.apache.arrow repository to the "release" area
-#
-# This script should only be run after the release has been approved
-# by the arrow PMC committee.
-#
-# See release/README.md for full release instructions
-#
-# Based in part on post-01-upload.sh from apache/arrow
-
-
-set -e
-set -u
-
-if [ "$#" -ne 2 ]; then
- echo "Usage: $0 <version> <rc-num>"
- echo "ex. $0 4.1.0 2"
- exit
-fi
-
-version=$1
-rc=$2
-
-tmp_dir=tmp-apache-arrow-dist
-
-echo "Recreate temporary directory: ${tmp_dir}"
-rm -rf ${tmp_dir}
-mkdir -p ${tmp_dir}
-
-echo "Clone dev dist repository"
-svn \
- co \
- https://dist.apache.org/repos/dist/dev/arrow/apache-arrow-rs-${version}-rc${rc} \
- ${tmp_dir}/dev
-
-echo "Clone release dist repository"
-svn co https://dist.apache.org/repos/dist/release/arrow ${tmp_dir}/release
-
-echo "Copy ${version}-rc${rc} to release working copy"
-release_version=arrow-rs-${version}
-mkdir -p ${tmp_dir}/release/${release_version}
-cp -r ${tmp_dir}/dev/* ${tmp_dir}/release/${release_version}/
-svn add ${tmp_dir}/release/${release_version}
-
-echo "Commit release"
-svn ci -m "Apache Arrow Rust ${version}" ${tmp_dir}/release
-
-echo "Clean up"
-rm -rf ${tmp_dir}
-
-echo "Success! The release is available here:"
-echo " https://dist.apache.org/repos/dist/release/arrow/${release_version}"
diff --git a/dev/release/run-rat.sh b/dev/release/run-rat.sh
deleted file mode 100755
index 94fa55f..0000000
--- a/dev/release/run-rat.sh
+++ /dev/null
@@ -1,43 +0,0 @@
-#!/bin/bash
-#
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-#
-
-RAT_VERSION=0.13
-
-# download apache rat
-if [ ! -f apache-rat-${RAT_VERSION}.jar ]; then
- curl -s https://repo1.maven.org/maven2/org/apache/rat/apache-rat/${RAT_VERSION}/apache-rat-${RAT_VERSION}.jar > apache-rat-${RAT_VERSION}.jar
-fi
-
-RAT="java -jar apache-rat-${RAT_VERSION}.jar -x "
-
-RELEASE_DIR=$(cd "$(dirname "$BASH_SOURCE")"; pwd)
-
-# generate the rat report
-$RAT $1 > rat.txt
-python $RELEASE_DIR/check-rat-report.py $RELEASE_DIR/rat_exclude_files.txt rat.txt > filtered_rat.txt
-cat filtered_rat.txt
-UNAPPROVED=`cat filtered_rat.txt | grep "NOT APPROVED" | wc -l`
-
-if [ "0" -eq "${UNAPPROVED}" ]; then
- echo "No unapproved licenses"
-else
- echo "${UNAPPROVED} unapproved licences. Check rat report: rat.txt"
- exit 1
-fi
diff --git a/dev/release/update_change_log.sh b/dev/release/update_change_log.sh
deleted file mode 100755
index a11f966..0000000
--- a/dev/release/update_change_log.sh
+++ /dev/null
@@ -1,39 +0,0 @@
-#!/bin/bash
-#
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-#
-
-# invokes the changelog generator with the config located in
-# arrow-rs/.github_changelog_generator
-#
-# Usage:
-# CHANGELOG_GITHUB_TOKEN=<TOKEN> ./update_change_log.sh
-
-set -e
-
-SOURCE_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
-SOURCE_TOP_DIR="$(cd "${SOURCE_DIR}/../../" && pwd)"
-
-pushd ${SOURCE_TOP_DIR}
-docker run -it --rm -e CHANGELOG_GITHUB_TOKEN=$CHANGELOG_GITHUB_TOKEN -v "$(pwd)":/usr/local/src/your-app githubchangeloggenerator/github-changelog-generator \
- --user apache \
- --project arrow-rs \
- --since-commit 2021-04-20 \
- --future-release 4.1.0
-
-sed -i "s/\\\n/\n\n/" CHANGELOG.md
diff --git a/dev/release/verify-release-candidate.sh b/dev/release/verify-release-candidate.sh
deleted file mode 100755
index 1ff34fd..0000000
--- a/dev/release/verify-release-candidate.sh
+++ /dev/null
@@ -1,154 +0,0 @@
-#!/bin/bash
-#
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-#
-
-case $# in
- 2) VERSION="$1"
- RC_NUMBER="$2"
- ;;
- *) echo "Usage: $0 X.Y.Z RC_NUMBER"
- exit 1
- ;;
-esac
-
-set -e
-set -x
-set -o pipefail
-
-SOURCE_DIR="$(cd "$(dirname "${BASH_SOURCE[0]:-$0}")" && pwd)"
-ARROW_DIR="$(dirname $(dirname ${SOURCE_DIR}))"
-ARROW_DIST_URL='https://dist.apache.org/repos/dist/dev/arrow'
-
-download_dist_file() {
- curl \
- --silent \
- --show-error \
- --fail \
- --location \
- --remote-name $ARROW_DIST_URL/$1
-}
-
-download_rc_file() {
- download_dist_file apache-arrow-rs-${VERSION}-rc${RC_NUMBER}/$1
-}
-
-import_gpg_keys() {
- download_dist_file KEYS
- gpg --import KEYS
-}
-
-fetch_archive() {
- local dist_name=$1
- download_rc_file ${dist_name}.tar.gz
- download_rc_file ${dist_name}.tar.gz.asc
- download_rc_file ${dist_name}.tar.gz.sha256
- download_rc_file ${dist_name}.tar.gz.sha512
- gpg --verify ${dist_name}.tar.gz.asc ${dist_name}.tar.gz
- shasum -a 256 -c ${dist_name}.tar.gz.sha256
- shasum -a 512 -c ${dist_name}.tar.gz.sha512
-}
-
-verify_dir_artifact_signatures() {
- # verify the signature and the checksums of each artifact
- find $1 -name '*.asc' | while read sigfile; do
- artifact=${sigfile/.asc/}
- gpg --verify $sigfile $artifact || exit 1
-
- # go into the directory because the checksum files contain only the
- # basename of the artifact
- pushd $(dirname $artifact)
- base_artifact=$(basename $artifact)
- if [ -f $base_artifact.sha256 ]; then
- shasum -a 256 -c $base_artifact.sha256 || exit 1
- fi
- shasum -a 512 -c $base_artifact.sha512 || exit 1
- popd
- done
-}
-
-setup_tempdir() {
- cleanup() {
- if [ "${TEST_SUCCESS}" = "yes" ]; then
- rm -fr "${ARROW_TMPDIR}"
- else
- echo "Failed to verify release candidate. See ${ARROW_TMPDIR} for details."
- fi
- }
-
- if [ -z "${ARROW_TMPDIR}" ]; then
- # clean up automatically if ARROW_TMPDIR is not defined
- ARROW_TMPDIR=$(mktemp -d -t "$1.XXXXX")
- trap cleanup EXIT
- else
- # don't clean up automatically
- mkdir -p "${ARROW_TMPDIR}"
- fi
-}
-
-test_source_distribution() {
- # install rust toolchain in a similar fashion like test-miniconda
- export RUSTUP_HOME=$PWD/test-rustup
- export CARGO_HOME=$PWD/test-rustup
-
- curl https://sh.rustup.rs -sSf | sh -s -- -y --no-modify-path
-
- export PATH=$RUSTUP_HOME/bin:$PATH
- source $RUSTUP_HOME/env
-
- # build and test rust
-
- # raises on any formatting errors
- rustup component add rustfmt --toolchain stable
- cargo fmt --all -- --check
-
- # Clone testing repositories if not cloned already
- git clone https://github.com/apache/arrow-testing.git arrow-testing-data
- git clone https://github.com/apache/parquet-testing.git parquet-testing-data
- export ARROW_TEST_DATA=$PWD/arrow-testing-data/data
- export PARQUET_TEST_DATA=$PWD/parquet-testing-data/data
-
- # use local modules because we don't publish modules to crates.io yet
- sed \
- -i.bak \
- -E \
- -e 's/^arrow = "([^"]*)"/arrow = { version = "\1", path = "..\/arrow" }/g' \
- -e 's/^parquet = "([^"]*)"/parquet = { version = "\1", path = "..\/parquet" }/g' \
- */Cargo.toml
-
- cargo build
- cargo test
-}
-
-TEST_SUCCESS=no
-
-setup_tempdir "arrow-${VERSION}"
-echo "Working in sandbox ${ARROW_TMPDIR}"
-cd ${ARROW_TMPDIR}
-
-dist_name="apache-arrow-rs-${VERSION}"
-import_gpg_keys
-fetch_archive ${dist_name}
-tar xf ${dist_name}.tar.gz
-pushd ${dist_name}
-test_source_distribution
-popd
-
-TEST_SUCCESS=yes
-echo 'Release candidate looks good!'
-exit 0
diff --git a/format/File.fbs b/format/File.fbs
deleted file mode 100644
index 906d494..0000000
--- a/format/File.fbs
+++ /dev/null
@@ -1,52 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-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;
diff --git a/format/Flight.proto b/format/Flight.proto
deleted file mode 100644
index b291d9d..0000000
--- a/format/Flight.proto
+++ /dev/null
@@ -1,335 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one
- * or more contributor license agreements. See the NOTICE file
- * distributed with this work for additional information
- * regarding copyright ownership. The ASF licenses this file
- * to you under the Apache License, Version 2.0 (the
- * "License"); you may not use this file except in compliance
- * with the License. You may obtain a copy of the License at
- * <p>
- * http://www.apache.org/licenses/LICENSE-2.0
- * <p>
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT 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.impl";
-option go_package = "github.com/apache/arrow/go/flight;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, 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 {
- // schema of the dataset as described in Schema.fbs::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 {
- // schema of the dataset as described in Schema.fbs::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 must be consumed.
- */
- repeated FlightEndpoint endpoint = 3;
-
- // Set these to -1 if unknown.
- int64 total_records = 4;
- int64 total_bytes = 5;
-}
-
-/*
- * 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. If the list is
- * empty, the expectation is that the ticket can only be redeemed on the
- * current service where the ticket was generated.
- */
- repeated Location location = 2;
-}
-
-/*
- * A location where a Flight service will accept retrieval of a particular
- * stream given a ticket.
- */
-message Location {
- string uri = 1;
-}
-
-/*
- * An opaque identifier that the service can use to retrieve a particular
- * portion of a stream.
- */
-message Ticket {
- bytes ticket = 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;
-}
diff --git a/format/Message.fbs b/format/Message.fbs
deleted file mode 100644
index f1c18d7..0000000
--- a/format/Message.fbs
+++ /dev/null
@@ -1,140 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-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<Int16> 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
- 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;
-}
-
-/// 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;
diff --git a/format/README.rst b/format/README.rst
deleted file mode 100644
index 0eaad49..0000000
--- a/format/README.rst
+++ /dev/null
@@ -1,25 +0,0 @@
-.. Licensed to the Apache Software Foundation (ASF) under one
-.. or more contributor license agreements. See the NOTICE file
-.. distributed with this work for additional information
-.. regarding copyright ownership. The ASF licenses this file
-.. to you under the Apache License, Version 2.0 (the
-.. "License"); you may not use this file except in compliance
-.. with the License. You may obtain a copy of the License at
-
-.. http://www.apache.org/licenses/LICENSE-2.0
-
-.. Unless required by applicable law or agreed to in writing,
-.. software distributed under the License is distributed on an
-.. "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-.. KIND, either express or implied. See the License for the
-.. specific language governing permissions and limitations
-.. under the License.
-
-Arrow Protocol Files
-====================
-
-This folder contains binary protocol definitions for the Arrow columnar format
-and other parts of the project, like the Flight RPC framework.
-
-For documentation about the Arrow format, see the `docs/source/format`
-directory.
diff --git a/format/Schema.fbs b/format/Schema.fbs
deleted file mode 100644
index 3b00dd4..0000000
--- a/format/Schema.fbs
+++ /dev/null
@@ -1,407 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-/// Logical types, vector layouts, and schemas
-
-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 {
-}
-
-table FixedSizeList {
- /// Number of list items per value
- listSize: int;
-}
-
-/// A Map is a logical nested type that is represented as
-///
-/// List<entries: Struct<key: K, value: V>>
-///
-/// 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 {
-}
-
-table FixedSizeBinary {
- /// Number of bytes per value
- byteWidth: int;
-}
-
-table Bool {
-}
-
-/// 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.
-table Decimal {
- /// Total number of decimal digits
- precision: int;
-
- /// Number of digits after the decimal point "."
- scale: int;
-
- /// Number of bits per value. The only accepted widths are 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 type representing 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 type. The physical storage type depends on the unit
-/// - SECOND and MILLISECOND: 32 bits
-/// - MICROSECOND and NANOSECOND: 64 bits
-table Time {
- unit: TimeUnit = MILLISECOND;
- bitWidth: int = 32;
-}
-
-/// Time elapsed from the Unix epoch, 00:00:00.000 on 1 January 1970, excluding
-/// leap seconds, as a 64-bit integer. Note that UNIX time does not include
-/// leap seconds.
-///
-/// The Timestamp metadata supports both "time zone naive" and "time zone
-/// aware" timestamps. Read about the timezone attribute for more detail
-table Timestamp {
- unit: TimeUnit;
-
- /// The time zone is a string indicating the name of a time zone, one of:
- ///
- /// * As used in the Olson time zone database (the "tz database" or
- /// "tzdata"), such as "America/New_York"
- /// * An absolute time zone 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:
- ///
- /// * If the time zone is null or equal to an empty string, the data is "time
- /// zone naive" and shall be displayed *as is* to the user, not localized
- /// to the locale of the user. This data can be though of as UTC but
- /// without having "UTC" as the time zone, it is not considered to be
- /// localized to any time zone
- ///
- /// * If the time zone is set to a valid value, values can be displayed as
- /// "localized" to that time zone, even though the underlying 64-bit
- /// integers are identical to the same data stored in UTC. Converting
- /// between time zones is a metadata-only operation and does not change the
- /// underlying values
- timezone: string;
-}
-
-enum IntervalUnit: short { YEAR_MONTH, DAY_TIME}
-// 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).
-// YEAR_MONTH - Indicates the number of elapsed whole months, stored as
-// 4-byte integers.
-// DAY_TIME - Indicates the number of elapsed days and milliseconds,
-// stored as 2 contiguous 32-bit integers (8-bytes in total). Support
-// of this IntervalUnit is not required for full arrow compatibility.
-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,
-}
-
-/// ----------------------------------------------------------------------
-/// 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;
diff --git a/format/SparseTensor.fbs b/format/SparseTensor.fbs
deleted file mode 100644
index a6fd2f9..0000000
--- a/format/SparseTensor.fbs
+++ /dev/null
@@ -1,228 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-/// 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;
diff --git a/format/Tensor.fbs b/format/Tensor.fbs
deleted file mode 100644
index 409297c..0000000
--- a/format/Tensor.fbs
+++ /dev/null
@@ -1,54 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-/// 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;
diff --git a/header b/header
deleted file mode 100644
index 70665d1..0000000
--- a/header
+++ /dev/null
@@ -1,16 +0,0 @@
-Licensed to the Apache Software Foundation (ASF) under one
-or more contributor license agreements. See the NOTICE file
-distributed with this work for additional information
-regarding copyright ownership. The ASF licenses this file
-to you under the Apache License, Version 2.0 (the
-"License"); you may not use this file except in compliance
-with the License. You may obtain a copy of the License at
-
-http://www.apache.org/licenses/LICENSE-2.0
-
-Unless required by applicable law or agreed to in writing, software
-distributed under the License is distributed on an "AS IS" BASIS,
-WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-See the License for the specific language governing permissions and
-limitations under the License.
-
diff --git a/integration-testing/Cargo.toml b/integration-testing/Cargo.toml
deleted file mode 100644
index af97787..0000000
--- a/integration-testing/Cargo.toml
+++ /dev/null
@@ -1,45 +0,0 @@
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-
-[package]
-name = "arrow-integration-testing"
-description = "Binaries used in the Arrow integration tests"
-version = "5.0.0-SNAPSHOT"
-homepage = "https://github.com/apache/arrow-rs"
-repository = "https://github.com/apache/arrow-rs"
-authors = ["Apache Arrow <dev@arrow.apache.org>"]
-license = "Apache-2.0"
-edition = "2018"
-publish = false
-
-[features]
-logging = ["tracing-subscriber"]
-
-[dependencies]
-arrow = { path = "../arrow" }
-arrow-flight = { path = "../arrow-flight" }
-async-trait = "0.1.41"
-clap = "2.33"
-futures = "0.3"
-hex = "0.4"
-prost = "0.7"
-serde = { version = "1.0", features = ["rc"] }
-serde_derive = "1.0"
-serde_json = { version = "1.0", features = ["preserve_order"] }
-tokio = { version = "1.0", features = ["macros", "rt", "rt-multi-thread"] }
-tonic = "0.4"
-tracing-subscriber = { version = "0.2.15", optional = true }
diff --git a/integration-testing/README.md b/integration-testing/README.md
deleted file mode 100644
index ff36542..0000000
--- a/integration-testing/README.md
+++ /dev/null
@@ -1,30 +0,0 @@
-<!---
- Licensed to the Apache Software Foundation (ASF) under one
- or more contributor license agreements. See the NOTICE file
- distributed with this work for additional information
- regarding copyright ownership. The ASF licenses this file
- to you under the Apache License, Version 2.0 (the
- "License"); you may not use this file except in compliance
- with the License. You may obtain a copy of the License at
-
- http://www.apache.org/licenses/LICENSE-2.0
-
- Unless required by applicable law or agreed to in writing,
- software distributed under the License is distributed on an
- "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
- KIND, either express or implied. See the License for the
- specific language governing permissions and limitations
- under the License.
--->
-
-# Apache Arrow Rust Integration Testing
-
-See [Integration.rst](../../docs/source/format/Integration.rst) for an overview of integration testing.
-
-This crate contains the following binaries, which are invoked by Archery during integration testing with other Arrow implementations.
-
-| Binary | Purpose |
-| --------------------------- | ----------------------------------------- |
-| arrow-file-to-stream | Converts an Arrow file to an Arrow stream |
-| arrow-stream-to-file | Converts an Arrow stream to an Arrow file |
-| arrow-json-integration-test | Converts between Arrow and JSON formats |
diff --git a/integration-testing/src/bin/arrow-file-to-stream.rs b/integration-testing/src/bin/arrow-file-to-stream.rs
deleted file mode 100644
index d6bb042..0000000
--- a/integration-testing/src/bin/arrow-file-to-stream.rs
+++ /dev/null
@@ -1,43 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-use std::env;
-use std::fs::File;
-use std::io::{self, BufReader};
-
-use arrow::error::Result;
-use arrow::ipc::reader::FileReader;
-use arrow::ipc::writer::StreamWriter;
-
-fn main() -> Result<()> {
- let args: Vec<String> = env::args().collect();
- let filename = &args[1];
- let f = File::open(filename)?;
- let reader = BufReader::new(f);
- let mut reader = FileReader::try_new(reader)?;
- let schema = reader.schema();
-
- let mut writer = StreamWriter::try_new(io::stdout(), &schema)?;
-
- reader.try_for_each(|batch| {
- let batch = batch?;
- writer.write(&batch)
- })?;
- writer.finish()?;
-
- Ok(())
-}
diff --git a/integration-testing/src/bin/arrow-json-integration-test.rs b/integration-testing/src/bin/arrow-json-integration-test.rs
deleted file mode 100644
index 2578020..0000000
--- a/integration-testing/src/bin/arrow-json-integration-test.rs
+++ /dev/null
@@ -1,180 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-use std::fs::File;
-
-use clap::{App, Arg};
-
-use arrow::error::{ArrowError, Result};
-use arrow::ipc::reader::FileReader;
-use arrow::ipc::writer::FileWriter;
-use arrow::util::integration_util::*;
-use arrow_integration_testing::read_json_file;
-
-fn main() -> Result<()> {
- let matches = App::new("rust arrow-json-integration-test")
- .arg(Arg::with_name("integration")
- .long("integration"))
- .arg(Arg::with_name("arrow")
- .long("arrow")
- .help("path to ARROW file")
- .takes_value(true))
- .arg(Arg::with_name("json")
- .long("json")
- .help("path to JSON file")
- .takes_value(true))
- .arg(Arg::with_name("mode")
- .long("mode")
- .help("mode of integration testing tool (ARROW_TO_JSON, JSON_TO_ARROW, VALIDATE)")
- .takes_value(true)
- .default_value("VALIDATE"))
- .arg(Arg::with_name("verbose")
- .long("verbose")
- .help("enable/disable verbose mode"))
- .get_matches();
-
- let arrow_file = matches
- .value_of("arrow")
- .expect("must provide path to arrow file");
- let json_file = matches
- .value_of("json")
- .expect("must provide path to json file");
- let mode = matches.value_of("mode").unwrap();
- let verbose = true; //matches.value_of("verbose").is_some();
-
- match mode {
- "JSON_TO_ARROW" => json_to_arrow(json_file, arrow_file, verbose),
- "ARROW_TO_JSON" => arrow_to_json(arrow_file, json_file, verbose),
- "VALIDATE" => validate(arrow_file, json_file, verbose),
- _ => panic!("mode {} not supported", mode),
- }
-}
-
-fn json_to_arrow(json_name: &str, arrow_name: &str, verbose: bool) -> Result<()> {
- if verbose {
- eprintln!("Converting {} to {}", json_name, arrow_name);
- }
-
- let json_file = read_json_file(json_name)?;
-
- let arrow_file = File::create(arrow_name)?;
- let mut writer = FileWriter::try_new(arrow_file, &json_file.schema)?;
-
- for b in json_file.batches {
- writer.write(&b)?;
- }
-
- writer.finish()?;
-
- Ok(())
-}
-
-fn arrow_to_json(arrow_name: &str, json_name: &str, verbose: bool) -> Result<()> {
- if verbose {
- eprintln!("Converting {} to {}", arrow_name, json_name);
- }
-
- let arrow_file = File::open(arrow_name)?;
- let reader = FileReader::try_new(arrow_file)?;
-
- let mut fields: Vec<ArrowJsonField> = vec![];
- for f in reader.schema().fields() {
- fields.push(ArrowJsonField::from(f));
- }
- let schema = ArrowJsonSchema { fields };
-
- let batches = reader
- .map(|batch| Ok(ArrowJsonBatch::from_batch(&batch?)))
- .collect::<Result<Vec<_>>>()?;
-
- let arrow_json = ArrowJson {
- schema,
- batches,
- dictionaries: None,
- };
-
- let json_file = File::create(json_name)?;
- serde_json::to_writer(&json_file, &arrow_json).unwrap();
-
- Ok(())
-}
-
-fn validate(arrow_name: &str, json_name: &str, verbose: bool) -> Result<()> {
- if verbose {
- eprintln!("Validating {} and {}", arrow_name, json_name);
- }
-
- // open JSON file
- let json_file = read_json_file(json_name)?;
-
- // open Arrow file
- let arrow_file = File::open(arrow_name)?;
- let mut arrow_reader = FileReader::try_new(arrow_file)?;
- let arrow_schema = arrow_reader.schema().as_ref().to_owned();
-
- // compare schemas
- if json_file.schema != arrow_schema {
- return Err(ArrowError::ComputeError(format!(
- "Schemas do not match. JSON: {:?}. Arrow: {:?}",
- json_file.schema, arrow_schema
- )));
- }
-
- let json_batches = &json_file.batches;
-
- // compare number of batches
- assert!(
- json_batches.len() == arrow_reader.num_batches(),
- "JSON batches and Arrow batches are unequal"
- );
-
- if verbose {
- eprintln!(
- "Schemas match. JSON file has {} batches.",
- json_batches.len()
- );
- }
-
- for json_batch in json_batches {
- if let Some(Ok(arrow_batch)) = arrow_reader.next() {
- // compare batches
- let num_columns = arrow_batch.num_columns();
- assert!(num_columns == json_batch.num_columns());
- assert!(arrow_batch.num_rows() == json_batch.num_rows());
-
- for i in 0..num_columns {
- assert_eq!(
- arrow_batch.column(i).data(),
- json_batch.column(i).data(),
- "Arrow and JSON batch columns not the same"
- );
- }
- } else {
- return Err(ArrowError::ComputeError(
- "no more arrow batches left".to_owned(),
- ));
- }
- }
-
- if arrow_reader.next().is_some() {
- return Err(ArrowError::ComputeError(
- "no more json batches left".to_owned(),
- ));
- }
-
- Ok(())
-}
diff --git a/integration-testing/src/bin/arrow-stream-to-file.rs b/integration-testing/src/bin/arrow-stream-to-file.rs
deleted file mode 100644
index f81d42e..0000000
--- a/integration-testing/src/bin/arrow-stream-to-file.rs
+++ /dev/null
@@ -1,34 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-use std::io;
-
-use arrow::error::Result;
-use arrow::ipc::reader::StreamReader;
-use arrow::ipc::writer::FileWriter;
-
-fn main() -> Result<()> {
- let mut arrow_stream_reader = StreamReader::try_new(io::stdin())?;
- let schema = arrow_stream_reader.schema();
-
- let mut writer = FileWriter::try_new(io::stdout(), &schema)?;
-
- arrow_stream_reader.try_for_each(|batch| writer.write(&batch?))?;
- writer.finish()?;
-
- Ok(())
-}
diff --git a/integration-testing/src/bin/flight-test-integration-client.rs b/integration-testing/src/bin/flight-test-integration-client.rs
deleted file mode 100644
index 1901553..0000000
--- a/integration-testing/src/bin/flight-test-integration-client.rs
+++ /dev/null
@@ -1,62 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-use arrow_integration_testing::flight_client_scenarios;
-
-use clap::{App, Arg};
-
-type Error = Box<dyn std::error::Error + Send + Sync + 'static>;
-type Result<T = (), E = Error> = std::result::Result<T, E>;
-
-#[tokio::main]
-async fn main() -> Result {
- #[cfg(feature = "logging")]
- tracing_subscriber::fmt::init();
-
- let matches = App::new("rust flight-test-integration-client")
- .arg(Arg::with_name("host").long("host").takes_value(true))
- .arg(Arg::with_name("port").long("port").takes_value(true))
- .arg(Arg::with_name("path").long("path").takes_value(true))
- .arg(
- Arg::with_name("scenario")
- .long("scenario")
- .takes_value(true),
- )
- .get_matches();
-
- let host = matches.value_of("host").expect("Host is required");
- let port = matches.value_of("port").expect("Port is required");
-
- match matches.value_of("scenario") {
- Some("middleware") => {
- flight_client_scenarios::middleware::run_scenario(host, port).await?
- }
- Some("auth:basic_proto") => {
- flight_client_scenarios::auth_basic_proto::run_scenario(host, port).await?
- }
- Some(scenario_name) => unimplemented!("Scenario not found: {}", scenario_name),
- None => {
- let path = matches
- .value_of("path")
- .expect("Path is required if scenario is not specified");
- flight_client_scenarios::integration_test::run_scenario(host, port, path)
- .await?;
- }
- }
-
- Ok(())
-}
diff --git a/integration-testing/src/bin/flight-test-integration-server.rs b/integration-testing/src/bin/flight-test-integration-server.rs
deleted file mode 100644
index b1b2807..0000000
--- a/integration-testing/src/bin/flight-test-integration-server.rs
+++ /dev/null
@@ -1,55 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-use clap::{App, Arg};
-
-use arrow_integration_testing::flight_server_scenarios;
-
-type Error = Box<dyn std::error::Error + Send + Sync + 'static>;
-type Result<T = (), E = Error> = std::result::Result<T, E>;
-
-#[tokio::main]
-async fn main() -> Result {
- #[cfg(feature = "logging")]
- tracing_subscriber::fmt::init();
-
- let matches = App::new("rust flight-test-integration-server")
- .about("Integration testing server for Flight.")
- .arg(Arg::with_name("port").long("port").takes_value(true))
- .arg(
- Arg::with_name("scenario")
- .long("scenario")
- .takes_value(true),
- )
- .get_matches();
-
- let port = matches.value_of("port").unwrap_or("0");
-
- match matches.value_of("scenario") {
- Some("middleware") => {
- flight_server_scenarios::middleware::scenario_setup(port).await?
- }
- Some("auth:basic_proto") => {
- flight_server_scenarios::auth_basic_proto::scenario_setup(port).await?
- }
- Some(scenario_name) => unimplemented!("Scenario not found: {}", scenario_name),
- None => {
- flight_server_scenarios::integration_test::scenario_setup(port).await?;
- }
- }
- Ok(())
-}
diff --git a/integration-testing/src/flight_client_scenarios.rs b/integration-testing/src/flight_client_scenarios.rs
deleted file mode 100644
index 66cced5..0000000
--- a/integration-testing/src/flight_client_scenarios.rs
+++ /dev/null
@@ -1,20 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-pub mod auth_basic_proto;
-pub mod integration_test;
-pub mod middleware;
diff --git a/integration-testing/src/flight_client_scenarios/auth_basic_proto.rs b/integration-testing/src/flight_client_scenarios/auth_basic_proto.rs
deleted file mode 100644
index 5e8cd46..0000000
--- a/integration-testing/src/flight_client_scenarios/auth_basic_proto.rs
+++ /dev/null
@@ -1,109 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-use crate::{AUTH_PASSWORD, AUTH_USERNAME};
-
-use arrow_flight::{
- flight_service_client::FlightServiceClient, BasicAuth, HandshakeRequest,
-};
-use futures::{stream, StreamExt};
-use prost::Message;
-use tonic::{metadata::MetadataValue, Request, Status};
-
-type Error = Box<dyn std::error::Error + Send + Sync + 'static>;
-type Result<T = (), E = Error> = std::result::Result<T, E>;
-
-type Client = FlightServiceClient<tonic::transport::Channel>;
-
-pub async fn run_scenario(host: &str, port: &str) -> Result {
- let url = format!("http://{}:{}", host, port);
- let mut client = FlightServiceClient::connect(url).await?;
-
- let action = arrow_flight::Action::default();
-
- let resp = client.do_action(Request::new(action.clone())).await;
- // This client is unauthenticated and should fail.
- match resp {
- Err(e) => {
- if e.code() != tonic::Code::Unauthenticated {
- return Err(Box::new(Status::internal(format!(
- "Expected UNAUTHENTICATED but got {:?}",
- e
- ))));
- }
- }
- Ok(other) => {
- return Err(Box::new(Status::internal(format!(
- "Expected UNAUTHENTICATED but got {:?}",
- other
- ))));
- }
- }
-
- let token = authenticate(&mut client, AUTH_USERNAME, AUTH_PASSWORD)
- .await
- .expect("must respond successfully from handshake");
-
- let mut request = Request::new(action);
- let metadata = request.metadata_mut();
- metadata.insert_bin(
- "auth-token-bin",
- MetadataValue::from_bytes(token.as_bytes()),
- );
-
- let resp = client.do_action(request).await?;
- let mut resp = resp.into_inner();
-
- let r = resp
- .next()
- .await
- .expect("No response received")
- .expect("Invalid response received");
-
- let body = String::from_utf8(r.body).unwrap();
- assert_eq!(body, AUTH_USERNAME);
-
- Ok(())
-}
-
-async fn authenticate(
- client: &mut Client,
- username: &str,
- password: &str,
-) -> Result<String> {
- let auth = BasicAuth {
- username: username.into(),
- password: password.into(),
- };
- let mut payload = vec![];
- auth.encode(&mut payload)?;
-
- let req = stream::once(async {
- HandshakeRequest {
- payload,
- ..HandshakeRequest::default()
- }
- });
-
- let rx = client.handshake(Request::new(req)).await?;
- let mut rx = rx.into_inner();
-
- let r = rx.next().await.expect("must respond from handshake")?;
- assert!(rx.next().await.is_none(), "must not respond a second time");
-
- Ok(String::from_utf8(r.payload).unwrap())
-}
diff --git a/integration-testing/src/flight_client_scenarios/integration_test.rs b/integration-testing/src/flight_client_scenarios/integration_test.rs
deleted file mode 100644
index ff61b5c..0000000
--- a/integration-testing/src/flight_client_scenarios/integration_test.rs
+++ /dev/null
@@ -1,271 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-use crate::{read_json_file, ArrowFile};
-
-use arrow::{
- array::ArrayRef,
- datatypes::SchemaRef,
- ipc::{self, reader, writer},
- record_batch::RecordBatch,
-};
-use arrow_flight::{
- flight_descriptor::DescriptorType, flight_service_client::FlightServiceClient,
- utils::flight_data_to_arrow_batch, FlightData, FlightDescriptor, Location, Ticket,
-};
-use futures::{channel::mpsc, sink::SinkExt, stream, StreamExt};
-use tonic::{Request, Streaming};
-
-use std::sync::Arc;
-
-type Error = Box<dyn std::error::Error + Send + Sync + 'static>;
-type Result<T = (), E = Error> = std::result::Result<T, E>;
-
-type Client = FlightServiceClient<tonic::transport::Channel>;
-
-pub async fn run_scenario(host: &str, port: &str, path: &str) -> Result {
- let url = format!("http://{}:{}", host, port);
-
- let client = FlightServiceClient::connect(url).await?;
-
- let ArrowFile {
- schema, batches, ..
- } = read_json_file(path)?;
-
- let schema = Arc::new(schema);
-
- let mut descriptor = FlightDescriptor::default();
- descriptor.set_type(DescriptorType::Path);
- descriptor.path = vec![path.to_string()];
-
- upload_data(
- client.clone(),
- schema.clone(),
- descriptor.clone(),
- batches.clone(),
- )
- .await?;
- verify_data(client, descriptor, schema, &batches).await?;
-
- Ok(())
-}
-
-async fn upload_data(
- mut client: Client,
- schema: SchemaRef,
- descriptor: FlightDescriptor,
- original_data: Vec<RecordBatch>,
-) -> Result {
- let (mut upload_tx, upload_rx) = mpsc::channel(10);
-
- let options = arrow::ipc::writer::IpcWriteOptions::default();
- let mut schema_flight_data =
- arrow_flight::utils::flight_data_from_arrow_schema(&schema, &options);
- schema_flight_data.flight_descriptor = Some(descriptor.clone());
- upload_tx.send(schema_flight_data).await?;
-
- let mut original_data_iter = original_data.iter().enumerate();
-
- if let Some((counter, first_batch)) = original_data_iter.next() {
- let metadata = counter.to_string().into_bytes();
- // Preload the first batch into the channel before starting the request
- send_batch(&mut upload_tx, &metadata, first_batch, &options).await?;
-
- let outer = client.do_put(Request::new(upload_rx)).await?;
- let mut inner = outer.into_inner();
-
- let r = inner
- .next()
- .await
- .expect("No response received")
- .expect("Invalid response received");
- assert_eq!(metadata, r.app_metadata);
-
- // Stream the rest of the batches
- for (counter, batch) in original_data_iter {
- let metadata = counter.to_string().into_bytes();
- send_batch(&mut upload_tx, &metadata, batch, &options).await?;
-
- let r = inner
- .next()
- .await
- .expect("No response received")
- .expect("Invalid response received");
- assert_eq!(metadata, r.app_metadata);
- }
- drop(upload_tx);
- assert!(
- inner.next().await.is_none(),
- "Should not receive more results"
- );
- } else {
- drop(upload_tx);
- client.do_put(Request::new(upload_rx)).await?;
- }
-
- Ok(())
-}
-
-async fn send_batch(
- upload_tx: &mut mpsc::Sender<FlightData>,
- metadata: &[u8],
- batch: &RecordBatch,
- options: &writer::IpcWriteOptions,
-) -> Result {
- let (dictionary_flight_data, mut batch_flight_data) =
- arrow_flight::utils::flight_data_from_arrow_batch(batch, &options);
-
- upload_tx
- .send_all(&mut stream::iter(dictionary_flight_data).map(Ok))
- .await?;
-
- // Only the record batch's FlightData gets app_metadata
- batch_flight_data.app_metadata = metadata.to_vec();
- upload_tx.send(batch_flight_data).await?;
- Ok(())
-}
-
-async fn verify_data(
- mut client: Client,
- descriptor: FlightDescriptor,
- expected_schema: SchemaRef,
- expected_data: &[RecordBatch],
-) -> Result {
- let resp = client.get_flight_info(Request::new(descriptor)).await?;
- let info = resp.into_inner();
-
- assert!(
- !info.endpoint.is_empty(),
- "No endpoints returned from Flight server",
- );
- for endpoint in info.endpoint {
- let ticket = endpoint
- .ticket
- .expect("No ticket returned from Flight server");
-
- assert!(
- !endpoint.location.is_empty(),
- "No locations returned from Flight server",
- );
- for location in endpoint.location {
- consume_flight_location(
- location,
- ticket.clone(),
- &expected_data,
- expected_schema.clone(),
- )
- .await?;
- }
- }
-
- Ok(())
-}
-
-async fn consume_flight_location(
- location: Location,
- ticket: Ticket,
- expected_data: &[RecordBatch],
- schema: SchemaRef,
-) -> Result {
- let mut location = location;
- // The other Flight implementations use the `grpc+tcp` scheme, but the Rust http libs
- // don't recognize this as valid.
- location.uri = location.uri.replace("grpc+tcp://", "grpc://");
-
- let mut client = FlightServiceClient::connect(location.uri).await?;
- let resp = client.do_get(ticket).await?;
- let mut resp = resp.into_inner();
-
- // We already have the schema from the FlightInfo, but the server sends it again as the
- // first FlightData. Ignore this one.
- let _schema_again = resp.next().await.unwrap();
-
- let mut dictionaries_by_field = vec![None; schema.fields().len()];
-
- for (counter, expected_batch) in expected_data.iter().enumerate() {
- let data = receive_batch_flight_data(
- &mut resp,
- schema.clone(),
- &mut dictionaries_by_field,
- )
- .await
- .unwrap_or_else(|| {
- panic!(
- "Got fewer batches than expected, received so far: {} expected: {}",
- counter,
- expected_data.len(),
- )
- });
-
- let metadata = counter.to_string().into_bytes();
- assert_eq!(metadata, data.app_metadata);
-
- let actual_batch =
- flight_data_to_arrow_batch(&data, schema.clone(), &dictionaries_by_field)
- .expect("Unable to convert flight data to Arrow batch");
-
- assert_eq!(expected_batch.schema(), actual_batch.schema());
- assert_eq!(expected_batch.num_columns(), actual_batch.num_columns());
- assert_eq!(expected_batch.num_rows(), actual_batch.num_rows());
- let schema = expected_batch.schema();
- for i in 0..expected_batch.num_columns() {
- let field = schema.field(i);
- let field_name = field.name();
-
- let expected_data = expected_batch.column(i).data();
- let actual_data = actual_batch.column(i).data();
-
- assert_eq!(expected_data, actual_data, "Data for field {}", field_name);
- }
- }
-
- assert!(
- resp.next().await.is_none(),
- "Got more batches than the expected: {}",
- expected_data.len(),
- );
-
- Ok(())
-}
-
-async fn receive_batch_flight_data(
- resp: &mut Streaming<FlightData>,
- schema: SchemaRef,
- dictionaries_by_field: &mut [Option<ArrayRef>],
-) -> Option<FlightData> {
- let mut data = resp.next().await?.ok()?;
- let mut message = arrow::ipc::root_as_message(&data.data_header[..])
- .expect("Error parsing first message");
-
- while message.header_type() == ipc::MessageHeader::DictionaryBatch {
- reader::read_dictionary(
- &data.data_body,
- message
- .header_as_dictionary_batch()
- .expect("Error parsing dictionary"),
- &schema,
- dictionaries_by_field,
- )
- .expect("Error reading dictionary");
-
- data = resp.next().await?.ok()?;
- message = arrow::ipc::root_as_message(&data.data_header[..])
- .expect("Error parsing message");
- }
-
- Some(data)
-}
diff --git a/integration-testing/src/flight_client_scenarios/middleware.rs b/integration-testing/src/flight_client_scenarios/middleware.rs
deleted file mode 100644
index cbca879..0000000
--- a/integration-testing/src/flight_client_scenarios/middleware.rs
+++ /dev/null
@@ -1,83 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-use arrow_flight::{
- flight_descriptor::DescriptorType, flight_service_client::FlightServiceClient,
- FlightDescriptor,
-};
-use tonic::{Request, Status};
-
-type Error = Box<dyn std::error::Error + Send + Sync + 'static>;
-type Result<T = (), E = Error> = std::result::Result<T, E>;
-
-pub async fn run_scenario(host: &str, port: &str) -> Result {
- let url = format!("http://{}:{}", host, port);
- let conn = tonic::transport::Endpoint::new(url)?.connect().await?;
- let mut client = FlightServiceClient::with_interceptor(conn, middleware_interceptor);
-
- let mut descriptor = FlightDescriptor::default();
- descriptor.set_type(DescriptorType::Cmd);
- descriptor.cmd = b"".to_vec();
-
- // This call is expected to fail.
- match client
- .get_flight_info(Request::new(descriptor.clone()))
- .await
- {
- Ok(_) => return Err(Box::new(Status::internal("Expected call to fail"))),
- Err(e) => {
- let headers = e.metadata();
- let middleware_header = headers.get("x-middleware");
- let value = middleware_header.map(|v| v.to_str().unwrap()).unwrap_or("");
-
- if value != "expected value" {
- let msg = format!(
- "On failing call: Expected to receive header 'x-middleware: expected value', \
- but instead got: '{}'",
- value
- );
- return Err(Box::new(Status::internal(msg)));
- }
- }
- }
-
- // This call should succeed
- descriptor.cmd = b"success".to_vec();
- let resp = client.get_flight_info(Request::new(descriptor)).await?;
-
- let headers = resp.metadata();
- let middleware_header = headers.get("x-middleware");
- let value = middleware_header.map(|v| v.to_str().unwrap()).unwrap_or("");
-
- if value != "expected value" {
- let msg = format!(
- "On success call: Expected to receive header 'x-middleware: expected value', \
- but instead got: '{}'",
- value
- );
- return Err(Box::new(Status::internal(msg)));
- }
-
- Ok(())
-}
-
-#[allow(clippy::unnecessary_wraps)]
-fn middleware_interceptor(mut req: Request<()>) -> Result<Request<()>, Status> {
- let metadata = req.metadata_mut();
- metadata.insert("x-middleware", "expected value".parse().unwrap());
- Ok(req)
-}
diff --git a/integration-testing/src/flight_server_scenarios.rs b/integration-testing/src/flight_server_scenarios.rs
deleted file mode 100644
index 9163b69..0000000
--- a/integration-testing/src/flight_server_scenarios.rs
+++ /dev/null
@@ -1,48 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-use std::net::SocketAddr;
-
-use arrow_flight::{FlightEndpoint, Location, Ticket};
-use tokio::net::TcpListener;
-
-pub mod auth_basic_proto;
-pub mod integration_test;
-pub mod middleware;
-
-type Error = Box<dyn std::error::Error + Send + Sync + 'static>;
-type Result<T = (), E = Error> = std::result::Result<T, E>;
-
-pub async fn listen_on(port: &str) -> Result<SocketAddr> {
- let addr: SocketAddr = format!("0.0.0.0:{}", port).parse()?;
-
- let listener = TcpListener::bind(addr).await?;
- let addr = listener.local_addr()?;
-
- Ok(addr)
-}
-
-pub fn endpoint(ticket: &str, location_uri: impl Into<String>) -> FlightEndpoint {
- FlightEndpoint {
- ticket: Some(Ticket {
- ticket: ticket.as_bytes().to_vec(),
- }),
- location: vec![Location {
- uri: location_uri.into(),
- }],
- }
-}
diff --git a/integration-testing/src/flight_server_scenarios/auth_basic_proto.rs b/integration-testing/src/flight_server_scenarios/auth_basic_proto.rs
deleted file mode 100644
index ea7ad3c..0000000
--- a/integration-testing/src/flight_server_scenarios/auth_basic_proto.rs
+++ /dev/null
@@ -1,225 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-use std::pin::Pin;
-use std::sync::Arc;
-
-use arrow_flight::{
- flight_service_server::FlightService, flight_service_server::FlightServiceServer,
- Action, ActionType, BasicAuth, Criteria, Empty, FlightData, FlightDescriptor,
- FlightInfo, HandshakeRequest, HandshakeResponse, PutResult, SchemaResult, Ticket,
-};
-use futures::{channel::mpsc, sink::SinkExt, Stream, StreamExt};
-use tokio::sync::Mutex;
-use tonic::{
- metadata::MetadataMap, transport::Server, Request, Response, Status, Streaming,
-};
-type TonicStream<T> = Pin<Box<dyn Stream<Item = T> + Send + Sync + 'static>>;
-
-type Error = Box<dyn std::error::Error + Send + Sync + 'static>;
-type Result<T = (), E = Error> = std::result::Result<T, E>;
-
-use prost::Message;
-
-use crate::{AUTH_PASSWORD, AUTH_USERNAME};
-
-pub async fn scenario_setup(port: &str) -> Result {
- let service = AuthBasicProtoScenarioImpl {
- username: AUTH_USERNAME.into(),
- password: AUTH_PASSWORD.into(),
- peer_identity: Arc::new(Mutex::new(None)),
- };
- let addr = super::listen_on(port).await?;
- let svc = FlightServiceServer::new(service);
-
- let server = Server::builder().add_service(svc).serve(addr);
-
- // NOTE: Log output used in tests to signal server is ready
- println!("Server listening on localhost:{}", addr.port());
- server.await?;
- Ok(())
-}
-
-#[derive(Clone)]
-pub struct AuthBasicProtoScenarioImpl {
- username: Arc<str>,
- password: Arc<str>,
- peer_identity: Arc<Mutex<Option<String>>>,
-}
-
-impl AuthBasicProtoScenarioImpl {
- async fn check_auth(
- &self,
- metadata: &MetadataMap,
- ) -> Result<GrpcServerCallContext, Status> {
- let token = metadata
- .get_bin("auth-token-bin")
- .and_then(|v| v.to_bytes().ok())
- .and_then(|b| String::from_utf8(b.to_vec()).ok());
- self.is_valid(token).await
- }
-
- async fn is_valid(
- &self,
- token: Option<String>,
- ) -> Result<GrpcServerCallContext, Status> {
- match token {
- Some(t) if t == *self.username => Ok(GrpcServerCallContext {
- peer_identity: self.username.to_string(),
- }),
- _ => Err(Status::unauthenticated("Invalid token")),
- }
- }
-}
-
-struct GrpcServerCallContext {
- peer_identity: String,
-}
-
-impl GrpcServerCallContext {
- pub fn peer_identity(&self) -> &str {
- &self.peer_identity
- }
-}
-
-#[tonic::async_trait]
-impl FlightService for AuthBasicProtoScenarioImpl {
- type HandshakeStream = TonicStream<Result<HandshakeResponse, Status>>;
- type ListFlightsStream = TonicStream<Result<FlightInfo, Status>>;
- type DoGetStream = TonicStream<Result<FlightData, Status>>;
- type DoPutStream = TonicStream<Result<PutResult, Status>>;
- type DoActionStream = TonicStream<Result<arrow_flight::Result, Status>>;
- type ListActionsStream = TonicStream<Result<ActionType, Status>>;
- type DoExchangeStream = TonicStream<Result<FlightData, Status>>;
-
- async fn get_schema(
- &self,
- request: Request<FlightDescriptor>,
- ) -> Result<Response<SchemaResult>, Status> {
- self.check_auth(request.metadata()).await?;
- Err(Status::unimplemented("Not yet implemented"))
- }
-
- async fn do_get(
- &self,
- request: Request<Ticket>,
- ) -> Result<Response<Self::DoGetStream>, Status> {
- self.check_auth(request.metadata()).await?;
- Err(Status::unimplemented("Not yet implemented"))
- }
-
- async fn handshake(
- &self,
- request: Request<Streaming<HandshakeRequest>>,
- ) -> Result<Response<Self::HandshakeStream>, Status> {
- let (tx, rx) = mpsc::channel(10);
-
- tokio::spawn({
- let username = self.username.clone();
- let password = self.password.clone();
-
- async move {
- let requests = request.into_inner();
-
- requests
- .for_each(move |req| {
- let mut tx = tx.clone();
- let req = req.expect("Error reading handshake request");
- let HandshakeRequest { payload, .. } = req;
-
- let auth = BasicAuth::decode(&*payload)
- .expect("Error parsing handshake request");
-
- let resp = if *auth.username == *username
- && *auth.password == *password
- {
- Ok(HandshakeResponse {
- payload: username.as_bytes().to_vec(),
- ..HandshakeResponse::default()
- })
- } else {
- Err(Status::unauthenticated(format!(
- "Don't know user {}",
- auth.username
- )))
- };
-
- async move {
- tx.send(resp)
- .await
- .expect("Error sending handshake response");
- }
- })
- .await;
- }
- });
-
- Ok(Response::new(Box::pin(rx)))
- }
-
- async fn list_flights(
- &self,
- request: Request<Criteria>,
- ) -> Result<Response<Self::ListFlightsStream>, Status> {
- self.check_auth(request.metadata()).await?;
- Err(Status::unimplemented("Not yet implemented"))
- }
-
- async fn get_flight_info(
- &self,
- request: Request<FlightDescriptor>,
- ) -> Result<Response<FlightInfo>, Status> {
- self.check_auth(request.metadata()).await?;
- Err(Status::unimplemented("Not yet implemented"))
- }
-
- async fn do_put(
- &self,
- request: Request<Streaming<FlightData>>,
- ) -> Result<Response<Self::DoPutStream>, Status> {
- self.check_auth(request.metadata()).await?;
- Err(Status::unimplemented("Not yet implemented"))
- }
-
- async fn do_action(
- &self,
- request: Request<Action>,
- ) -> Result<Response<Self::DoActionStream>, Status> {
- let flight_context = self.check_auth(request.metadata()).await?;
- // Respond with the authenticated username.
- let buf = flight_context.peer_identity().as_bytes().to_vec();
- let result = arrow_flight::Result { body: buf };
- let output = futures::stream::once(async { Ok(result) });
- Ok(Response::new(Box::pin(output) as Self::DoActionStream))
- }
-
- async fn list_actions(
- &self,
- request: Request<Empty>,
- ) -> Result<Response<Self::ListActionsStream>, Status> {
- self.check_auth(request.metadata()).await?;
- Err(Status::unimplemented("Not yet implemented"))
- }
-
- async fn do_exchange(
- &self,
- request: Request<Streaming<FlightData>>,
- ) -> Result<Response<Self::DoExchangeStream>, Status> {
- self.check_auth(request.metadata()).await?;
- Err(Status::unimplemented("Not yet implemented"))
- }
-}
diff --git a/integration-testing/src/flight_server_scenarios/integration_test.rs b/integration-testing/src/flight_server_scenarios/integration_test.rs
deleted file mode 100644
index ee42a47..0000000
--- a/integration-testing/src/flight_server_scenarios/integration_test.rs
+++ /dev/null
@@ -1,385 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-use std::collections::HashMap;
-use std::convert::TryFrom;
-use std::pin::Pin;
-use std::sync::Arc;
-
-use arrow::{
- array::ArrayRef,
- datatypes::Schema,
- datatypes::SchemaRef,
- ipc::{self, reader},
- record_batch::RecordBatch,
-};
-use arrow_flight::{
- flight_descriptor::DescriptorType, flight_service_server::FlightService,
- flight_service_server::FlightServiceServer, Action, ActionType, Criteria, Empty,
- FlightData, FlightDescriptor, FlightEndpoint, FlightInfo, HandshakeRequest,
- HandshakeResponse, PutResult, SchemaResult, Ticket,
-};
-use futures::{channel::mpsc, sink::SinkExt, Stream, StreamExt};
-use tokio::sync::Mutex;
-use tonic::{transport::Server, Request, Response, Status, Streaming};
-
-type TonicStream<T> = Pin<Box<dyn Stream<Item = T> + Send + Sync + 'static>>;
-
-type Error = Box<dyn std::error::Error + Send + Sync + 'static>;
-type Result<T = (), E = Error> = std::result::Result<T, E>;
-
-pub async fn scenario_setup(port: &str) -> Result {
- let addr = super::listen_on(port).await?;
-
- let service = FlightServiceImpl {
- server_location: format!("grpc+tcp://{}", addr),
- ..Default::default()
- };
- let svc = FlightServiceServer::new(service);
-
- let server = Server::builder().add_service(svc).serve(addr);
-
- // NOTE: Log output used in tests to signal server is ready
- println!("Server listening on localhost:{}", addr.port());
- server.await?;
- Ok(())
-}
-
-#[derive(Debug, Clone)]
-struct IntegrationDataset {
- schema: Schema,
- chunks: Vec<RecordBatch>,
-}
-
-#[derive(Clone, Default)]
-pub struct FlightServiceImpl {
- server_location: String,
- uploaded_chunks: Arc<Mutex<HashMap<String, IntegrationDataset>>>,
-}
-
-impl FlightServiceImpl {
- fn endpoint_from_path(&self, path: &str) -> FlightEndpoint {
- super::endpoint(path, &self.server_location)
- }
-}
-
-#[tonic::async_trait]
-impl FlightService for FlightServiceImpl {
- type HandshakeStream = TonicStream<Result<HandshakeResponse, Status>>;
- type ListFlightsStream = TonicStream<Result<FlightInfo, Status>>;
- type DoGetStream = TonicStream<Result<FlightData, Status>>;
- type DoPutStream = TonicStream<Result<PutResult, Status>>;
- type DoActionStream = TonicStream<Result<arrow_flight::Result, Status>>;
- type ListActionsStream = TonicStream<Result<ActionType, Status>>;
- type DoExchangeStream = TonicStream<Result<FlightData, Status>>;
-
- async fn get_schema(
- &self,
- _request: Request<FlightDescriptor>,
- ) -> Result<Response<SchemaResult>, Status> {
- Err(Status::unimplemented("Not yet implemented"))
- }
-
- async fn do_get(
- &self,
- request: Request<Ticket>,
- ) -> Result<Response<Self::DoGetStream>, Status> {
- let ticket = request.into_inner();
-
- let key = String::from_utf8(ticket.ticket.to_vec())
- .map_err(|e| Status::invalid_argument(format!("Invalid ticket: {:?}", e)))?;
-
- let uploaded_chunks = self.uploaded_chunks.lock().await;
-
- let flight = uploaded_chunks.get(&key).ok_or_else(|| {
- Status::not_found(format!("Could not find flight. {}", key))
- })?;
-
- let options = arrow::ipc::writer::IpcWriteOptions::default();
-
- let schema = std::iter::once({
- Ok(arrow_flight::utils::flight_data_from_arrow_schema(
- &flight.schema,
- &options,
- ))
- });
-
- let batches = flight
- .chunks
- .iter()
- .enumerate()
- .flat_map(|(counter, batch)| {
- let (dictionary_flight_data, mut batch_flight_data) =
- arrow_flight::utils::flight_data_from_arrow_batch(batch, &options);
-
- // Only the record batch's FlightData gets app_metadata
- let metadata = counter.to_string().into_bytes();
- batch_flight_data.app_metadata = metadata;
-
- dictionary_flight_data
- .into_iter()
- .chain(std::iter::once(batch_flight_data))
- .map(Ok)
- });
-
- let output = futures::stream::iter(schema.chain(batches).collect::<Vec<_>>());
-
- Ok(Response::new(Box::pin(output) as Self::DoGetStream))
- }
-
- async fn handshake(
- &self,
- _request: Request<Streaming<HandshakeRequest>>,
- ) -> Result<Response<Self::HandshakeStream>, Status> {
- Err(Status::unimplemented("Not yet implemented"))
- }
-
- async fn list_flights(
- &self,
- _request: Request<Criteria>,
- ) -> Result<Response<Self::ListFlightsStream>, Status> {
- Err(Status::unimplemented("Not yet implemented"))
- }
-
- async fn get_flight_info(
- &self,
- request: Request<FlightDescriptor>,
- ) -> Result<Response<FlightInfo>, Status> {
- let descriptor = request.into_inner();
-
- match descriptor.r#type {
- t if t == DescriptorType::Path as i32 => {
- let path = &descriptor.path;
- if path.is_empty() {
- return Err(Status::invalid_argument("Invalid path"));
- }
-
- let uploaded_chunks = self.uploaded_chunks.lock().await;
- let flight = uploaded_chunks.get(&path[0]).ok_or_else(|| {
- Status::not_found(format!("Could not find flight. {}", path[0]))
- })?;
-
- let endpoint = self.endpoint_from_path(&path[0]);
-
- let total_records: usize =
- flight.chunks.iter().map(|chunk| chunk.num_rows()).sum();
-
- let options = arrow::ipc::writer::IpcWriteOptions::default();
- let schema = arrow_flight::utils::ipc_message_from_arrow_schema(
- &flight.schema,
- &options,
- )
- .expect(
- "Could not generate schema bytes from schema stored by a DoPut; \
- this should be impossible",
- );
-
- let info = FlightInfo {
- schema,
- flight_descriptor: Some(descriptor.clone()),
- endpoint: vec![endpoint],
- total_records: total_records as i64,
- total_bytes: -1,
- };
-
- Ok(Response::new(info))
- }
- other => Err(Status::unimplemented(format!("Request type: {}", other))),
- }
- }
-
- async fn do_put(
- &self,
- request: Request<Streaming<FlightData>>,
- ) -> Result<Response<Self::DoPutStream>, Status> {
- let mut input_stream = request.into_inner();
- let flight_data = input_stream
- .message()
- .await?
- .ok_or_else(|| Status::invalid_argument("Must send some FlightData"))?;
-
- let descriptor = flight_data
- .flight_descriptor
- .clone()
- .ok_or_else(|| Status::invalid_argument("Must have a descriptor"))?;
-
- if descriptor.r#type != DescriptorType::Path as i32 || descriptor.path.is_empty()
- {
- return Err(Status::invalid_argument("Must specify a path"));
- }
-
- let key = descriptor.path[0].clone();
-
- let schema = Schema::try_from(&flight_data)
- .map_err(|e| Status::invalid_argument(format!("Invalid schema: {:?}", e)))?;
- let schema_ref = Arc::new(schema.clone());
-
- let (response_tx, response_rx) = mpsc::channel(10);
-
- let uploaded_chunks = self.uploaded_chunks.clone();
-
- tokio::spawn(async {
- let mut error_tx = response_tx.clone();
- if let Err(e) = save_uploaded_chunks(
- uploaded_chunks,
- schema_ref,
- input_stream,
- response_tx,
- schema,
- key,
- )
- .await
- {
- error_tx.send(Err(e)).await.expect("Error sending error")
- }
- });
-
- Ok(Response::new(Box::pin(response_rx) as Self::DoPutStream))
- }
-
- async fn do_action(
- &self,
- _request: Request<Action>,
- ) -> Result<Response<Self::DoActionStream>, Status> {
- Err(Status::unimplemented("Not yet implemented"))
- }
-
- async fn list_actions(
- &self,
- _request: Request<Empty>,
- ) -> Result<Response<Self::ListActionsStream>, Status> {
- Err(Status::unimplemented("Not yet implemented"))
- }
-
- async fn do_exchange(
- &self,
- _request: Request<Streaming<FlightData>>,
- ) -> Result<Response<Self::DoExchangeStream>, Status> {
- Err(Status::unimplemented("Not yet implemented"))
- }
-}
-
-async fn send_app_metadata(
- tx: &mut mpsc::Sender<Result<PutResult, Status>>,
- app_metadata: &[u8],
-) -> Result<(), Status> {
- tx.send(Ok(PutResult {
- app_metadata: app_metadata.to_vec(),
- }))
- .await
- .map_err(|e| Status::internal(format!("Could not send PutResult: {:?}", e)))
-}
-
-async fn record_batch_from_message(
- message: ipc::Message<'_>,
- data_body: &[u8],
- schema_ref: SchemaRef,
- dictionaries_by_field: &[Option<ArrayRef>],
-) -> Result<RecordBatch, Status> {
- let ipc_batch = message.header_as_record_batch().ok_or_else(|| {
- Status::internal("Could not parse message header as record batch")
- })?;
-
- let arrow_batch_result = reader::read_record_batch(
- data_body,
- ipc_batch,
- schema_ref,
- &dictionaries_by_field,
- );
-
- arrow_batch_result.map_err(|e| {
- Status::internal(format!("Could not convert to RecordBatch: {:?}", e))
- })
-}
-
-async fn dictionary_from_message(
- message: ipc::Message<'_>,
- data_body: &[u8],
- schema_ref: SchemaRef,
- dictionaries_by_field: &mut [Option<ArrayRef>],
-) -> Result<(), Status> {
- let ipc_batch = message.header_as_dictionary_batch().ok_or_else(|| {
- Status::internal("Could not parse message header as dictionary batch")
- })?;
-
- let dictionary_batch_result =
- reader::read_dictionary(data_body, ipc_batch, &schema_ref, dictionaries_by_field);
- dictionary_batch_result.map_err(|e| {
- Status::internal(format!("Could not convert to Dictionary: {:?}", e))
- })
-}
-
-async fn save_uploaded_chunks(
- uploaded_chunks: Arc<Mutex<HashMap<String, IntegrationDataset>>>,
- schema_ref: Arc<Schema>,
- mut input_stream: Streaming<FlightData>,
- mut response_tx: mpsc::Sender<Result<PutResult, Status>>,
- schema: Schema,
- key: String,
-) -> Result<(), Status> {
- let mut chunks = vec![];
- let mut uploaded_chunks = uploaded_chunks.lock().await;
-
- let mut dictionaries_by_field = vec![None; schema_ref.fields().len()];
-
- while let Some(Ok(data)) = input_stream.next().await {
- let message = arrow::ipc::root_as_message(&data.data_header[..])
- .map_err(|e| Status::internal(format!("Could not parse message: {:?}", e)))?;
-
- match message.header_type() {
- ipc::MessageHeader::Schema => {
- return Err(Status::internal(
- "Not expecting a schema when messages are read",
- ))
- }
- ipc::MessageHeader::RecordBatch => {
- send_app_metadata(&mut response_tx, &data.app_metadata).await?;
-
- let batch = record_batch_from_message(
- message,
- &data.data_body,
- schema_ref.clone(),
- &dictionaries_by_field,
- )
- .await?;
-
- chunks.push(batch);
- }
- ipc::MessageHeader::DictionaryBatch => {
- dictionary_from_message(
- message,
- &data.data_body,
- schema_ref.clone(),
- &mut dictionaries_by_field,
- )
- .await?;
- }
- t => {
- return Err(Status::internal(format!(
- "Reading types other than record batches not yet supported, \
- unable to read {:?}",
- t
- )));
- }
- }
- }
-
- let dataset = IntegrationDataset { schema, chunks };
- uploaded_chunks.insert(key, dataset);
-
- Ok(())
-}
diff --git a/integration-testing/src/flight_server_scenarios/middleware.rs b/integration-testing/src/flight_server_scenarios/middleware.rs
deleted file mode 100644
index 1416acc..0000000
--- a/integration-testing/src/flight_server_scenarios/middleware.rs
+++ /dev/null
@@ -1,150 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-use std::pin::Pin;
-
-use arrow_flight::{
- flight_descriptor::DescriptorType, flight_service_server::FlightService,
- flight_service_server::FlightServiceServer, Action, ActionType, Criteria, Empty,
- FlightData, FlightDescriptor, FlightInfo, HandshakeRequest, HandshakeResponse,
- PutResult, SchemaResult, Ticket,
-};
-use futures::Stream;
-use tonic::{transport::Server, Request, Response, Status, Streaming};
-
-type TonicStream<T> = Pin<Box<dyn Stream<Item = T> + Send + Sync + 'static>>;
-
-type Error = Box<dyn std::error::Error + Send + Sync + 'static>;
-type Result<T = (), E = Error> = std::result::Result<T, E>;
-
-pub async fn scenario_setup(port: &str) -> Result {
- let service = MiddlewareScenarioImpl {};
- let svc = FlightServiceServer::new(service);
- let addr = super::listen_on(port).await?;
-
- let server = Server::builder().add_service(svc).serve(addr);
-
- // NOTE: Log output used in tests to signal server is ready
- println!("Server listening on localhost:{}", addr.port());
- server.await?;
- Ok(())
-}
-
-#[derive(Clone, Default)]
-pub struct MiddlewareScenarioImpl {}
-
-#[tonic::async_trait]
-impl FlightService for MiddlewareScenarioImpl {
- type HandshakeStream = TonicStream<Result<HandshakeResponse, Status>>;
- type ListFlightsStream = TonicStream<Result<FlightInfo, Status>>;
- type DoGetStream = TonicStream<Result<FlightData, Status>>;
- type DoPutStream = TonicStream<Result<PutResult, Status>>;
- type DoActionStream = TonicStream<Result<arrow_flight::Result, Status>>;
- type ListActionsStream = TonicStream<Result<ActionType, Status>>;
- type DoExchangeStream = TonicStream<Result<FlightData, Status>>;
-
- async fn get_schema(
- &self,
- _request: Request<FlightDescriptor>,
- ) -> Result<Response<SchemaResult>, Status> {
- Err(Status::unimplemented("Not yet implemented"))
- }
-
- async fn do_get(
- &self,
- _request: Request<Ticket>,
- ) -> Result<Response<Self::DoGetStream>, Status> {
- Err(Status::unimplemented("Not yet implemented"))
- }
-
- async fn handshake(
- &self,
- _request: Request<Streaming<HandshakeRequest>>,
- ) -> Result<Response<Self::HandshakeStream>, Status> {
- Err(Status::unimplemented("Not yet implemented"))
- }
-
- async fn list_flights(
- &self,
- _request: Request<Criteria>,
- ) -> Result<Response<Self::ListFlightsStream>, Status> {
- Err(Status::unimplemented("Not yet implemented"))
- }
-
- async fn get_flight_info(
- &self,
- request: Request<FlightDescriptor>,
- ) -> Result<Response<FlightInfo>, Status> {
- let middleware_header = request.metadata().get("x-middleware").cloned();
-
- let descriptor = request.into_inner();
-
- if descriptor.r#type == DescriptorType::Cmd as i32 && descriptor.cmd == b"success"
- {
- // Return a fake location - the test doesn't read it
- let endpoint = super::endpoint("foo", "grpc+tcp://localhost:10010");
-
- let info = FlightInfo {
- flight_descriptor: Some(descriptor),
- endpoint: vec![endpoint],
- ..Default::default()
- };
-
- let mut response = Response::new(info);
- if let Some(value) = middleware_header {
- response.metadata_mut().insert("x-middleware", value);
- }
-
- return Ok(response);
- }
-
- let mut status = Status::unknown("Unknown");
- if let Some(value) = middleware_header {
- status.metadata_mut().insert("x-middleware", value);
- }
-
- Err(status)
- }
-
- async fn do_put(
- &self,
- _request: Request<Streaming<FlightData>>,
- ) -> Result<Response<Self::DoPutStream>, Status> {
- Err(Status::unimplemented("Not yet implemented"))
- }
-
- async fn do_action(
- &self,
- _request: Request<Action>,
- ) -> Result<Response<Self::DoActionStream>, Status> {
- Err(Status::unimplemented("Not yet implemented"))
- }
-
- async fn list_actions(
- &self,
- _request: Request<Empty>,
- ) -> Result<Response<Self::ListActionsStream>, Status> {
- Err(Status::unimplemented("Not yet implemented"))
- }
-
- async fn do_exchange(
- &self,
- _request: Request<Streaming<FlightData>>,
- ) -> Result<Response<Self::DoExchangeStream>, Status> {
- Err(Status::unimplemented("Not yet implemented"))
- }
-}
diff --git a/integration-testing/src/lib.rs b/integration-testing/src/lib.rs
deleted file mode 100644
index 22eed03..0000000
--- a/integration-testing/src/lib.rs
+++ /dev/null
@@ -1,601 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-//! Common code used in the integration test binaries
-
-use hex::decode;
-use serde_json::Value;
-
-use arrow::util::integration_util::ArrowJsonBatch;
-
-use arrow::array::*;
-use arrow::datatypes::{DataType, Field, IntervalUnit, Schema};
-use arrow::error::{ArrowError, Result};
-use arrow::record_batch::RecordBatch;
-use arrow::{
- buffer::Buffer,
- buffer::MutableBuffer,
- datatypes::ToByteSlice,
- util::{bit_util, integration_util::*},
-};
-
-use std::collections::HashMap;
-use std::fs::File;
-use std::io::BufReader;
-use std::sync::Arc;
-
-/// The expected username for the basic auth integration test.
-pub const AUTH_USERNAME: &str = "arrow";
-/// The expected password for the basic auth integration test.
-pub const AUTH_PASSWORD: &str = "flight";
-
-pub mod flight_client_scenarios;
-pub mod flight_server_scenarios;
-
-pub struct ArrowFile {
- pub schema: Schema,
- // we can evolve this into a concrete Arrow type
- // this is temporarily not being read from
- pub _dictionaries: HashMap<i64, ArrowJsonDictionaryBatch>,
- pub batches: Vec<RecordBatch>,
-}
-
-pub fn read_json_file(json_name: &str) -> Result<ArrowFile> {
- let json_file = File::open(json_name)?;
- let reader = BufReader::new(json_file);
- let arrow_json: Value = serde_json::from_reader(reader).unwrap();
- let schema = Schema::from(&arrow_json["schema"])?;
- // read dictionaries
- let mut dictionaries = HashMap::new();
- if let Some(dicts) = arrow_json.get("dictionaries") {
- for d in dicts
- .as_array()
- .expect("Unable to get dictionaries as array")
- {
- let json_dict: ArrowJsonDictionaryBatch = serde_json::from_value(d.clone())
- .expect("Unable to get dictionary from JSON");
- // TODO: convert to a concrete Arrow type
- dictionaries.insert(json_dict.id, json_dict);
- }
- }
-
- let mut batches = vec![];
- for b in arrow_json["batches"].as_array().unwrap() {
- let json_batch: ArrowJsonBatch = serde_json::from_value(b.clone()).unwrap();
- let batch = record_batch_from_json(&schema, json_batch, Some(&dictionaries))?;
- batches.push(batch);
- }
- Ok(ArrowFile {
- schema,
- _dictionaries: dictionaries,
- batches,
- })
-}
-
-fn record_batch_from_json(
- schema: &Schema,
- json_batch: ArrowJsonBatch,
- json_dictionaries: Option<&HashMap<i64, ArrowJsonDictionaryBatch>>,
-) -> Result<RecordBatch> {
- let mut columns = vec![];
-
- for (field, json_col) in schema.fields().iter().zip(json_batch.columns) {
- let col = array_from_json(field, json_col, json_dictionaries)?;
- columns.push(col);
- }
-
- RecordBatch::try_new(Arc::new(schema.clone()), columns)
-}
-
-/// Construct an Arrow array from a partially typed JSON column
-fn array_from_json(
- field: &Field,
- json_col: ArrowJsonColumn,
- dictionaries: Option<&HashMap<i64, ArrowJsonDictionaryBatch>>,
-) -> Result<ArrayRef> {
- match field.data_type() {
- DataType::Null => Ok(Arc::new(NullArray::new(json_col.count))),
- DataType::Boolean => {
- let mut b = BooleanBuilder::new(json_col.count);
- for (is_valid, value) in json_col
- .validity
- .as_ref()
- .unwrap()
- .iter()
- .zip(json_col.data.unwrap())
- {
- match is_valid {
- 1 => b.append_value(value.as_bool().unwrap()),
- _ => b.append_null(),
- }?;
- }
- Ok(Arc::new(b.finish()))
- }
- DataType::Int8 => {
- let mut b = Int8Builder::new(json_col.count);
- for (is_valid, value) in json_col
- .validity
- .as_ref()
- .unwrap()
- .iter()
- .zip(json_col.data.unwrap())
- {
- match is_valid {
- 1 => b.append_value(value.as_i64().ok_or_else(|| {
- ArrowError::JsonError(format!(
- "Unable to get {:?} as int64",
- value
- ))
- })? as i8),
- _ => b.append_null(),
- }?;
- }
- Ok(Arc::new(b.finish()))
- }
- DataType::Int16 => {
- let mut b = Int16Builder::new(json_col.count);
- for (is_valid, value) in json_col
- .validity
- .as_ref()
- .unwrap()
- .iter()
- .zip(json_col.data.unwrap())
- {
- match is_valid {
- 1 => b.append_value(value.as_i64().unwrap() as i16),
- _ => b.append_null(),
- }?;
- }
- Ok(Arc::new(b.finish()))
- }
- DataType::Int32
- | DataType::Date32
- | DataType::Time32(_)
- | DataType::Interval(IntervalUnit::YearMonth) => {
- let mut b = Int32Builder::new(json_col.count);
- for (is_valid, value) in json_col
- .validity
- .as_ref()
- .unwrap()
- .iter()
- .zip(json_col.data.unwrap())
- {
- match is_valid {
- 1 => b.append_value(value.as_i64().unwrap() as i32),
- _ => b.append_null(),
- }?;
- }
- let array = Arc::new(b.finish()) as ArrayRef;
- arrow::compute::cast(&array, field.data_type())
- }
- DataType::Int64
- | DataType::Date64
- | DataType::Time64(_)
- | DataType::Timestamp(_, _)
- | DataType::Duration(_)
- | DataType::Interval(IntervalUnit::DayTime) => {
- let mut b = Int64Builder::new(json_col.count);
- for (is_valid, value) in json_col
- .validity
- .as_ref()
- .unwrap()
- .iter()
- .zip(json_col.data.unwrap())
- {
- match is_valid {
- 1 => b.append_value(match value {
- Value::Number(n) => n.as_i64().unwrap(),
- Value::String(s) => {
- s.parse().expect("Unable to parse string as i64")
- }
- _ => panic!("Unable to parse {:?} as number", value),
- }),
- _ => b.append_null(),
- }?;
- }
- let array = Arc::new(b.finish()) as ArrayRef;
- arrow::compute::cast(&array, field.data_type())
- }
- DataType::UInt8 => {
- let mut b = UInt8Builder::new(json_col.count);
- for (is_valid, value) in json_col
- .validity
- .as_ref()
- .unwrap()
- .iter()
- .zip(json_col.data.unwrap())
- {
- match is_valid {
- 1 => b.append_value(value.as_u64().unwrap() as u8),
- _ => b.append_null(),
- }?;
- }
- Ok(Arc::new(b.finish()))
- }
- DataType::UInt16 => {
- let mut b = UInt16Builder::new(json_col.count);
- for (is_valid, value) in json_col
- .validity
- .as_ref()
- .unwrap()
- .iter()
- .zip(json_col.data.unwrap())
- {
- match is_valid {
- 1 => b.append_value(value.as_u64().unwrap() as u16),
- _ => b.append_null(),
- }?;
- }
- Ok(Arc::new(b.finish()))
- }
- DataType::UInt32 => {
- let mut b = UInt32Builder::new(json_col.count);
- for (is_valid, value) in json_col
- .validity
- .as_ref()
- .unwrap()
- .iter()
- .zip(json_col.data.unwrap())
- {
- match is_valid {
- 1 => b.append_value(value.as_u64().unwrap() as u32),
- _ => b.append_null(),
- }?;
- }
- Ok(Arc::new(b.finish()))
- }
- DataType::UInt64 => {
- let mut b = UInt64Builder::new(json_col.count);
- for (is_valid, value) in json_col
- .validity
- .as_ref()
- .unwrap()
- .iter()
- .zip(json_col.data.unwrap())
- {
- match is_valid {
- 1 => b.append_value(
- value
- .as_str()
- .unwrap()
- .parse()
- .expect("Unable to parse string as u64"),
- ),
- _ => b.append_null(),
- }?;
- }
- Ok(Arc::new(b.finish()))
- }
- DataType::Float32 => {
- let mut b = Float32Builder::new(json_col.count);
- for (is_valid, value) in json_col
- .validity
- .as_ref()
- .unwrap()
- .iter()
- .zip(json_col.data.unwrap())
- {
- match is_valid {
- 1 => b.append_value(value.as_f64().unwrap() as f32),
- _ => b.append_null(),
- }?;
- }
- Ok(Arc::new(b.finish()))
- }
- DataType::Float64 => {
- let mut b = Float64Builder::new(json_col.count);
- for (is_valid, value) in json_col
- .validity
- .as_ref()
- .unwrap()
- .iter()
- .zip(json_col.data.unwrap())
- {
- match is_valid {
- 1 => b.append_value(value.as_f64().unwrap()),
- _ => b.append_null(),
- }?;
- }
- Ok(Arc::new(b.finish()))
- }
- DataType::Binary => {
- let mut b = BinaryBuilder::new(json_col.count);
- for (is_valid, value) in json_col
- .validity
- .as_ref()
- .unwrap()
- .iter()
- .zip(json_col.data.unwrap())
- {
- match is_valid {
- 1 => {
- let v = decode(value.as_str().unwrap()).unwrap();
- b.append_value(&v)
- }
- _ => b.append_null(),
- }?;
- }
- Ok(Arc::new(b.finish()))
- }
- DataType::LargeBinary => {
- let mut b = LargeBinaryBuilder::new(json_col.count);
- for (is_valid, value) in json_col
- .validity
- .as_ref()
- .unwrap()
- .iter()
- .zip(json_col.data.unwrap())
- {
- match is_valid {
- 1 => {
- let v = decode(value.as_str().unwrap()).unwrap();
- b.append_value(&v)
- }
- _ => b.append_null(),
- }?;
- }
- Ok(Arc::new(b.finish()))
- }
- DataType::Utf8 => {
- let mut b = StringBuilder::new(json_col.count);
- for (is_valid, value) in json_col
- .validity
- .as_ref()
- .unwrap()
- .iter()
- .zip(json_col.data.unwrap())
- {
- match is_valid {
- 1 => b.append_value(value.as_str().unwrap()),
- _ => b.append_null(),
- }?;
- }
- Ok(Arc::new(b.finish()))
- }
- DataType::LargeUtf8 => {
- let mut b = LargeStringBuilder::new(json_col.count);
- for (is_valid, value) in json_col
- .validity
- .as_ref()
- .unwrap()
- .iter()
- .zip(json_col.data.unwrap())
- {
- match is_valid {
- 1 => b.append_value(value.as_str().unwrap()),
- _ => b.append_null(),
- }?;
- }
- Ok(Arc::new(b.finish()))
- }
- DataType::FixedSizeBinary(len) => {
- let mut b = FixedSizeBinaryBuilder::new(json_col.count, *len);
- for (is_valid, value) in json_col
- .validity
- .as_ref()
- .unwrap()
- .iter()
- .zip(json_col.data.unwrap())
- {
- match is_valid {
- 1 => {
- let v = hex::decode(value.as_str().unwrap()).unwrap();
- b.append_value(&v)
- }
- _ => b.append_null(),
- }?;
- }
- Ok(Arc::new(b.finish()))
- }
- DataType::List(child_field) => {
- let null_buf = create_null_buf(&json_col);
- let children = json_col.children.clone().unwrap();
- let child_array = array_from_json(
- &child_field,
- children.get(0).unwrap().clone(),
- dictionaries,
- )?;
- let offsets: Vec<i32> = json_col
- .offset
- .unwrap()
- .iter()
- .map(|v| v.as_i64().unwrap() as i32)
- .collect();
- let list_data = ArrayData::builder(field.data_type().clone())
- .len(json_col.count)
- .offset(0)
- .add_buffer(Buffer::from(&offsets.to_byte_slice()))
- .add_child_data(child_array.data().clone())
- .null_bit_buffer(null_buf)
- .build();
- Ok(Arc::new(ListArray::from(list_data)))
- }
- DataType::LargeList(child_field) => {
- let null_buf = create_null_buf(&json_col);
- let children = json_col.children.clone().unwrap();
- let child_array = array_from_json(
- &child_field,
- children.get(0).unwrap().clone(),
- dictionaries,
- )?;
- let offsets: Vec<i64> = json_col
- .offset
- .unwrap()
- .iter()
- .map(|v| match v {
- Value::Number(n) => n.as_i64().unwrap(),
- Value::String(s) => s.parse::<i64>().unwrap(),
- _ => panic!("64-bit offset must be either string or number"),
- })
- .collect();
- let list_data = ArrayData::builder(field.data_type().clone())
- .len(json_col.count)
- .offset(0)
- .add_buffer(Buffer::from(&offsets.to_byte_slice()))
- .add_child_data(child_array.data().clone())
- .null_bit_buffer(null_buf)
- .build();
- Ok(Arc::new(LargeListArray::from(list_data)))
- }
- DataType::FixedSizeList(child_field, _) => {
- let children = json_col.children.clone().unwrap();
- let child_array = array_from_json(
- &child_field,
- children.get(0).unwrap().clone(),
- dictionaries,
- )?;
- let null_buf = create_null_buf(&json_col);
- let list_data = ArrayData::builder(field.data_type().clone())
- .len(json_col.count)
- .add_child_data(child_array.data().clone())
- .null_bit_buffer(null_buf)
- .build();
- Ok(Arc::new(FixedSizeListArray::from(list_data)))
- }
- DataType::Struct(fields) => {
- // construct struct with null data
- let null_buf = create_null_buf(&json_col);
- let mut array_data = ArrayData::builder(field.data_type().clone())
- .len(json_col.count)
- .null_bit_buffer(null_buf);
-
- for (field, col) in fields.iter().zip(json_col.children.unwrap()) {
- let array = array_from_json(field, col, dictionaries)?;
- array_data = array_data.add_child_data(array.data().clone());
- }
-
- let array = StructArray::from(array_data.build());
- Ok(Arc::new(array))
- }
- DataType::Dictionary(key_type, value_type) => {
- let dict_id = field.dict_id().ok_or_else(|| {
- ArrowError::JsonError(format!(
- "Unable to find dict_id for field {:?}",
- field
- ))
- })?;
- // find dictionary
- let dictionary = dictionaries
- .ok_or_else(|| {
- ArrowError::JsonError(format!(
- "Unable to find any dictionaries for field {:?}",
- field
- ))
- })?
- .get(&dict_id);
- match dictionary {
- Some(dictionary) => dictionary_array_from_json(
- field, json_col, key_type, value_type, dictionary,
- ),
- None => Err(ArrowError::JsonError(format!(
- "Unable to find dictionary for field {:?}",
- field
- ))),
- }
- }
- t => Err(ArrowError::JsonError(format!(
- "data type {:?} not supported",
- t
- ))),
- }
-}
-
-fn dictionary_array_from_json(
- field: &Field,
- json_col: ArrowJsonColumn,
- dict_key: &DataType,
- dict_value: &DataType,
- dictionary: &ArrowJsonDictionaryBatch,
-) -> Result<ArrayRef> {
- match dict_key {
- DataType::Int8
- | DataType::Int16
- | DataType::Int32
- | DataType::Int64
- | DataType::UInt8
- | DataType::UInt16
- | DataType::UInt32
- | DataType::UInt64 => {
- let null_buf = create_null_buf(&json_col);
-
- // build the key data into a buffer, then construct values separately
- let key_field = Field::new_dict(
- "key",
- dict_key.clone(),
- field.is_nullable(),
- field
- .dict_id()
- .expect("Dictionary fields must have a dict_id value"),
- field
- .dict_is_ordered()
- .expect("Dictionary fields must have a dict_is_ordered value"),
- );
- let keys = array_from_json(&key_field, json_col, None)?;
- // note: not enough info on nullability of dictionary
- let value_field = Field::new("value", dict_value.clone(), true);
- println!("dictionary value type: {:?}", dict_value);
- let values =
- array_from_json(&value_field, dictionary.data.columns[0].clone(), None)?;
-
- // convert key and value to dictionary data
- let dict_data = ArrayData::builder(field.data_type().clone())
- .len(keys.len())
- .add_buffer(keys.data().buffers()[0].clone())
- .null_bit_buffer(null_buf)
- .add_child_data(values.data().clone())
- .build();
-
- let array = match dict_key {
- DataType::Int8 => {
- Arc::new(Int8DictionaryArray::from(dict_data)) as ArrayRef
- }
- DataType::Int16 => Arc::new(Int16DictionaryArray::from(dict_data)),
- DataType::Int32 => Arc::new(Int32DictionaryArray::from(dict_data)),
- DataType::Int64 => Arc::new(Int64DictionaryArray::from(dict_data)),
- DataType::UInt8 => Arc::new(UInt8DictionaryArray::from(dict_data)),
- DataType::UInt16 => Arc::new(UInt16DictionaryArray::from(dict_data)),
- DataType::UInt32 => Arc::new(UInt32DictionaryArray::from(dict_data)),
- DataType::UInt64 => Arc::new(UInt64DictionaryArray::from(dict_data)),
- _ => unreachable!(),
- };
- Ok(array)
- }
- _ => Err(ArrowError::JsonError(format!(
- "Dictionary key type {:?} not supported",
- dict_key
- ))),
- }
-}
-
-/// A helper to create a null buffer from a Vec<bool>
-fn create_null_buf(json_col: &ArrowJsonColumn) -> Buffer {
- let num_bytes = bit_util::ceil(json_col.count, 8);
- let mut null_buf = MutableBuffer::new(num_bytes).with_bitset(num_bytes, false);
- json_col
- .validity
- .clone()
- .unwrap()
- .iter()
- .enumerate()
- .for_each(|(i, v)| {
- let null_slice = null_buf.as_slice_mut();
- if *v != 0 {
- bit_util::set_bit(null_slice, i);
- }
- });
- null_buf.into()
-}
diff --git a/parquet-testing b/parquet-testing
deleted file mode 160000
index 8e7badc..0000000
--- a/parquet-testing
+++ /dev/null
@@ -1 +0,0 @@
-Subproject commit 8e7badc6a3817a02e06d17b5d8ab6b6dc356e890
diff --git a/parquet/Cargo.toml b/parquet/Cargo.toml
deleted file mode 100644
index d66ee23..0000000
--- a/parquet/Cargo.toml
+++ /dev/null
@@ -1,83 +0,0 @@
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-
-[package]
-name = "parquet"
-version = "5.0.0-SNAPSHOT"
-license = "Apache-2.0"
-description = "Apache Parquet implementation in Rust"
-homepage = "https://github.com/apache/arrow-rs"
-repository = "https://github.com/apache/arrow-rs"
-authors = ["Apache Arrow <dev@arrow.apache.org>"]
-keywords = [ "arrow", "parquet", "hadoop" ]
-readme = "README.md"
-build = "build.rs"
-edition = "2018"
-
-[dependencies]
-# update note: pin `parquet-format` to specific version until it does not break at minor
-# version, see ARROW-11187.
-parquet-format = "~2.6.1"
-byteorder = "1"
-thrift = "0.13"
-snap = { version = "1.0", optional = true }
-brotli = { version = "3.3", optional = true }
-flate2 = { version = "1.0", optional = true }
-lz4 = { version = "1.23", optional = true }
-zstd = { version = "0.8", optional = true }
-chrono = "0.4"
-num-bigint = "0.4"
-arrow = { path = "../arrow", version = "5.0.0-SNAPSHOT", optional = true }
-base64 = { version = "0.13", optional = true }
-clap = { version = "2.33.3", optional = true }
-serde_json = { version = "1.0", features = ["preserve_order"], optional = true }
-rand = "0.8"
-
-[dev-dependencies]
-criterion = "0.3"
-rand = "0.8"
-snap = "1.0"
-brotli = "3.3"
-flate2 = "1.0"
-lz4 = "1.23"
-zstd = "0.8"
-arrow = { path = "../arrow", version = "5.0.0-SNAPSHOT" }
-serde_json = { version = "1.0", features = ["preserve_order"] }
-
-[features]
-default = ["arrow", "snap", "brotli", "flate2", "lz4", "zstd", "base64"]
-cli = ["serde_json", "base64", "clap"]
-
-[[ bin ]]
-name = "parquet-read"
-required-features = ["cli"]
-
-[[ bin ]]
-name = "parquet-schema"
-required-features = ["cli"]
-
-[[ bin ]]
-name = "parquet-rowcount"
-required-features = ["cli"]
-
-[[bench]]
-name = "arrow_writer"
-harness = false
-
-[[bench]]
-name = "arrow_array_reader"
-harness = false
\ No newline at end of file
diff --git a/parquet/README.md b/parquet/README.md
deleted file mode 100644
index 338dba7..0000000
--- a/parquet/README.md
+++ /dev/null
@@ -1,143 +0,0 @@
-<!---
- Licensed to the Apache Software Foundation (ASF) under one
- or more contributor license agreements. See the NOTICE file
- distributed with this work for additional information
- regarding copyright ownership. The ASF licenses this file
- to you under the Apache License, Version 2.0 (the
- "License"); you may not use this file except in compliance
- with the License. You may obtain a copy of the License at
-
- http://www.apache.org/licenses/LICENSE-2.0
-
- Unless required by applicable law or agreed to in writing,
- software distributed under the License is distributed on an
- "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
- KIND, either express or implied. See the License for the
- specific language governing permissions and limitations
- under the License.
--->
-
-# An Apache Parquet implementation in Rust
-
-[](https://crates.io/crates/parquet)
-
-## Usage
-
-Add this to your Cargo.toml:
-
-```toml
-[dependencies]
-parquet = "^4"
-```
-
-and this to your crate root:
-
-```rust
-extern crate parquet;
-```
-
-Example usage of reading data:
-
-```rust
-use std::fs::File;
-use std::path::Path;
-use parquet::file::reader::{FileReader, SerializedFileReader};
-
-let file = File::open(&Path::new("/path/to/file")).unwrap();
-let reader = SerializedFileReader::new(file).unwrap();
-let mut iter = reader.get_row_iter(None).unwrap();
-while let Some(record) = iter.next() {
- println!("{}", record);
-}
-```
-
-See [crate documentation](https://docs.rs/crate/parquet/5.0.0-SNAPSHOT) on available API.
-
-## Upgrading from versions prior to 4.0
-
-If you are upgrading from version 3.0 or previous of this crate, you
-likely need to change your code to use [`ConvertedType`] rather than
-[`LogicalType`] to preserve existing behaviour in your code.
-
-Version 2.4.0 of the Parquet format introduced a `LogicalType` to replace the existing `ConvertedType`.
-This crate used `parquet::basic::LogicalType` to map to the `ConvertedType`, but this has been renamed to `parquet::basic::ConvertedType` from version 4.0 of this crate.
-
-The `ConvertedType` is deprecated in the format, but is still written
-to preserve backward compatibility.
-It is preferred that `LogicalType` is used, as it supports nanosecond
-precision timestamps without using the deprecated `Int96` Parquet type.
-
-## Supported Parquet Version
-
-- Parquet-format 2.6.0
-
-To update Parquet format to a newer version, check if [parquet-format](https://github.com/sunchao/parquet-format-rs)
-version is available. Then simply update version of `parquet-format` crate in Cargo.toml.
-
-## Features
-
-- [x] All encodings supported
-- [x] All compression codecs supported
-- [x] Read support
- - [x] Primitive column value readers
- - [x] Row record reader
- - [x] Arrow record reader
-- [ ] Statistics support
-- [x] Write support
- - [x] Primitive column value writers
- - [ ] Row record writer
- - [x] Arrow record writer
-- [ ] Predicate pushdown
-- [x] Parquet format 2.6.0 support
-
-## Requirements
-
-Parquet requires LLVM. Our windows CI image includes LLVM but to build the libraries locally windows
-users will have to install LLVM. Follow [this](https://github.com/appveyor/ci/issues/2651) link for info.
-
-## Build
-
-Run `cargo build` or `cargo build --release` to build in release mode.
-Some features take advantage of SSE4.2 instructions, which can be
-enabled by adding `RUSTFLAGS="-C target-feature=+sse4.2"` before the
-`cargo build` command.
-
-## Test
-
-Run `cargo test` for unit tests. To also run tests related to the binaries, use `cargo test --features cli`.
-
-## Binaries
-
-The following binaries are provided (use `cargo install --features cli` to install them):
-
-- **parquet-schema** for printing Parquet file schema and metadata.
- `Usage: parquet-schema <file-path>`, where `file-path` is the path to a Parquet file. Use `-v/--verbose` flag
- to print full metadata or schema only (when not specified only schema will be printed).
-
-- **parquet-read** for reading records from a Parquet file.
- `Usage: parquet-read <file-path> [num-records]`, where `file-path` is the path to a Parquet file,
- and `num-records` is the number of records to read from a file (when not specified all records will
- be printed). Use `-j/--json` to print records in JSON lines format.
-
-- **parquet-rowcount** for reporting the number of records in one or more Parquet files.
- `Usage: parquet-rowcount <file-paths>...`, where `<file-paths>...` is a space separated list of one or more
- files to read.
-
-If you see `Library not loaded` error, please make sure `LD_LIBRARY_PATH` is set properly:
-
-```
-export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:$(rustc --print sysroot)/lib
-```
-
-## Benchmarks
-
-Run `cargo bench` for benchmarks.
-
-## Docs
-
-To build documentation, run `cargo doc --no-deps`.
-To compile and view in the browser, run `cargo doc --no-deps --open`.
-
-## License
-
-Licensed under the Apache License, Version 2.0: http://www.apache.org/licenses/LICENSE-2.0.
diff --git a/parquet/benches/arrow_array_reader.rs b/parquet/benches/arrow_array_reader.rs
deleted file mode 100644
index 6e87512..0000000
--- a/parquet/benches/arrow_array_reader.rs
+++ /dev/null
@@ -1,766 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-use criterion::{criterion_group, criterion_main, Criterion};
-use parquet::util::{DataPageBuilder, DataPageBuilderImpl, InMemoryPageIterator};
-use parquet::{
- arrow::array_reader::ArrayReader,
- basic::Encoding,
- column::page::PageIterator,
- data_type::{ByteArrayType, Int32Type},
- schema::types::{ColumnDescPtr, SchemaDescPtr},
-};
-use std::{collections::VecDeque, sync::Arc};
-
-fn build_test_schema() -> SchemaDescPtr {
- use parquet::schema::{parser::parse_message_type, types::SchemaDescriptor};
- let message_type = "
- message test_schema {
- REQUIRED INT32 mandatory_int32_leaf;
- REPEATED Group test_mid_int32 {
- OPTIONAL INT32 optional_int32_leaf;
- }
- REQUIRED BYTE_ARRAY mandatory_string_leaf (UTF8);
- REPEATED Group test_mid_string {
- OPTIONAL BYTE_ARRAY optional_string_leaf (UTF8);
- }
- }
- ";
- parse_message_type(message_type)
- .map(|t| Arc::new(SchemaDescriptor::new(Arc::new(t))))
- .unwrap()
-}
-
-// test data params
-const NUM_ROW_GROUPS: usize = 1;
-const PAGES_PER_GROUP: usize = 2;
-const VALUES_PER_PAGE: usize = 10_000;
-const BATCH_SIZE: usize = 8192;
-
-use rand::{rngs::StdRng, Rng, SeedableRng};
-
-pub fn seedable_rng() -> StdRng {
- StdRng::seed_from_u64(42)
-}
-
-fn build_plain_encoded_int32_page_iterator(
- schema: SchemaDescPtr,
- column_desc: ColumnDescPtr,
- null_density: f32,
-) -> impl PageIterator + Clone {
- let max_def_level = column_desc.max_def_level();
- let max_rep_level = column_desc.max_rep_level();
- let rep_levels = vec![0; VALUES_PER_PAGE];
- let mut rng = seedable_rng();
- let mut pages: Vec<Vec<parquet::column::page::Page>> = Vec::new();
- let mut int32_value = 0;
- for _i in 0..NUM_ROW_GROUPS {
- let mut column_chunk_pages = Vec::new();
- for _j in 0..PAGES_PER_GROUP {
- // generate page
- let mut values = Vec::with_capacity(VALUES_PER_PAGE);
- let mut def_levels = Vec::with_capacity(VALUES_PER_PAGE);
- for _k in 0..VALUES_PER_PAGE {
- let def_level = if rng.gen::<f32>() < null_density {
- max_def_level - 1
- } else {
- max_def_level
- };
- if def_level == max_def_level {
- int32_value += 1;
- values.push(int32_value);
- }
- def_levels.push(def_level);
- }
- let mut page_builder =
- DataPageBuilderImpl::new(column_desc.clone(), values.len() as u32, true);
- page_builder.add_rep_levels(max_rep_level, &rep_levels);
- page_builder.add_def_levels(max_def_level, &def_levels);
- page_builder.add_values::<Int32Type>(Encoding::PLAIN, &values);
- column_chunk_pages.push(page_builder.consume());
- }
- pages.push(column_chunk_pages);
- }
-
- InMemoryPageIterator::new(schema, column_desc, pages)
-}
-
-fn build_dictionary_encoded_int32_page_iterator(
- schema: SchemaDescPtr,
- column_desc: ColumnDescPtr,
- null_density: f32,
-) -> impl PageIterator + Clone {
- use parquet::encoding::{DictEncoder, Encoder};
- let max_def_level = column_desc.max_def_level();
- let max_rep_level = column_desc.max_rep_level();
- let rep_levels = vec![0; VALUES_PER_PAGE];
- // generate 1% unique values
- const NUM_UNIQUE_VALUES: usize = VALUES_PER_PAGE / 100;
- let unique_values = (0..NUM_UNIQUE_VALUES)
- .map(|x| (x + 1) as i32)
- .collect::<Vec<_>>();
- let mut rng = seedable_rng();
- let mut pages: Vec<Vec<parquet::column::page::Page>> = Vec::new();
- for _i in 0..NUM_ROW_GROUPS {
- let mut column_chunk_pages = VecDeque::new();
- let mem_tracker = Arc::new(parquet::memory::MemTracker::new());
- let mut dict_encoder =
- DictEncoder::<Int32Type>::new(column_desc.clone(), mem_tracker);
- // add data pages
- for _j in 0..PAGES_PER_GROUP {
- // generate page
- let mut values = Vec::with_capacity(VALUES_PER_PAGE);
- let mut def_levels = Vec::with_capacity(VALUES_PER_PAGE);
- for _k in 0..VALUES_PER_PAGE {
- let def_level = if rng.gen::<f32>() < null_density {
- max_def_level - 1
- } else {
- max_def_level
- };
- if def_level == max_def_level {
- // select random value from list of unique values
- let int32_value = unique_values[rng.gen_range(0..NUM_UNIQUE_VALUES)];
- values.push(int32_value);
- }
- def_levels.push(def_level);
- }
- let mut page_builder =
- DataPageBuilderImpl::new(column_desc.clone(), values.len() as u32, true);
- page_builder.add_rep_levels(max_rep_level, &rep_levels);
- page_builder.add_def_levels(max_def_level, &def_levels);
- let _ = dict_encoder.put(&values);
- let indices = dict_encoder
- .write_indices()
- .expect("write_indices() should be OK");
- page_builder.add_indices(indices);
- column_chunk_pages.push_back(page_builder.consume());
- }
- // add dictionary page
- let dict = dict_encoder
- .write_dict()
- .expect("write_dict() should be OK");
- let dict_page = parquet::column::page::Page::DictionaryPage {
- buf: dict,
- num_values: dict_encoder.num_entries() as u32,
- encoding: Encoding::RLE_DICTIONARY,
- is_sorted: false,
- };
- column_chunk_pages.push_front(dict_page);
- pages.push(column_chunk_pages.into());
- }
-
- InMemoryPageIterator::new(schema, column_desc, pages)
-}
-
-fn build_plain_encoded_string_page_iterator(
- schema: SchemaDescPtr,
- column_desc: ColumnDescPtr,
- null_density: f32,
-) -> impl PageIterator + Clone {
- let max_def_level = column_desc.max_def_level();
- let max_rep_level = column_desc.max_rep_level();
- let rep_levels = vec![0; VALUES_PER_PAGE];
- let mut rng = seedable_rng();
- let mut pages: Vec<Vec<parquet::column::page::Page>> = Vec::new();
- for i in 0..NUM_ROW_GROUPS {
- let mut column_chunk_pages = Vec::new();
- for j in 0..PAGES_PER_GROUP {
- // generate page
- let mut values = Vec::with_capacity(VALUES_PER_PAGE);
- let mut def_levels = Vec::with_capacity(VALUES_PER_PAGE);
- for k in 0..VALUES_PER_PAGE {
- let def_level = if rng.gen::<f32>() < null_density {
- max_def_level - 1
- } else {
- max_def_level
- };
- if def_level == max_def_level {
- let string_value =
- format!("Test value {}, row group: {}, page: {}", k, i, j);
- values
- .push(parquet::data_type::ByteArray::from(string_value.as_str()));
- }
- def_levels.push(def_level);
- }
- let mut page_builder =
- DataPageBuilderImpl::new(column_desc.clone(), values.len() as u32, true);
- page_builder.add_rep_levels(max_rep_level, &rep_levels);
- page_builder.add_def_levels(max_def_level, &def_levels);
- page_builder.add_values::<ByteArrayType>(Encoding::PLAIN, &values);
- column_chunk_pages.push(page_builder.consume());
- }
- pages.push(column_chunk_pages);
- }
-
- InMemoryPageIterator::new(schema, column_desc, pages)
-}
-
-fn build_dictionary_encoded_string_page_iterator(
- schema: SchemaDescPtr,
- column_desc: ColumnDescPtr,
- null_density: f32,
-) -> impl PageIterator + Clone {
- use parquet::encoding::{DictEncoder, Encoder};
- let max_def_level = column_desc.max_def_level();
- let max_rep_level = column_desc.max_rep_level();
- let rep_levels = vec![0; VALUES_PER_PAGE];
- // generate 1% unique values
- const NUM_UNIQUE_VALUES: usize = VALUES_PER_PAGE / 100;
- let unique_values = (0..NUM_UNIQUE_VALUES)
- .map(|x| format!("Dictionary value {}", x))
- .collect::<Vec<_>>();
- let mut rng = seedable_rng();
- let mut pages: Vec<Vec<parquet::column::page::Page>> = Vec::new();
- for _i in 0..NUM_ROW_GROUPS {
- let mut column_chunk_pages = VecDeque::new();
- let mem_tracker = Arc::new(parquet::memory::MemTracker::new());
- let mut dict_encoder =
- DictEncoder::<ByteArrayType>::new(column_desc.clone(), mem_tracker);
- // add data pages
- for _j in 0..PAGES_PER_GROUP {
- // generate page
- let mut values = Vec::with_capacity(VALUES_PER_PAGE);
- let mut def_levels = Vec::with_capacity(VALUES_PER_PAGE);
- for _k in 0..VALUES_PER_PAGE {
- let def_level = if rng.gen::<f32>() < null_density {
- max_def_level - 1
- } else {
- max_def_level
- };
- if def_level == max_def_level {
- // select random value from list of unique values
- let string_value =
- unique_values[rng.gen_range(0..NUM_UNIQUE_VALUES)].as_str();
- values.push(parquet::data_type::ByteArray::from(string_value));
- }
- def_levels.push(def_level);
- }
- let mut page_builder =
- DataPageBuilderImpl::new(column_desc.clone(), values.len() as u32, true);
- page_builder.add_rep_levels(max_rep_level, &rep_levels);
- page_builder.add_def_levels(max_def_level, &def_levels);
- let _ = dict_encoder.put(&values);
- let indices = dict_encoder
- .write_indices()
- .expect("write_indices() should be OK");
- page_builder.add_indices(indices);
- column_chunk_pages.push_back(page_builder.consume());
- }
- // add dictionary page
- let dict = dict_encoder
- .write_dict()
- .expect("write_dict() should be OK");
- let dict_page = parquet::column::page::Page::DictionaryPage {
- buf: dict,
- num_values: dict_encoder.num_entries() as u32,
- encoding: Encoding::RLE_DICTIONARY,
- is_sorted: false,
- };
- column_chunk_pages.push_front(dict_page);
- pages.push(column_chunk_pages.into());
- }
-
- InMemoryPageIterator::new(schema, column_desc, pages)
-}
-
-fn bench_array_reader(mut array_reader: impl ArrayReader) -> usize {
- // test procedure: read data in batches of 8192 until no more data
- let mut total_count = 0;
- loop {
- let array = array_reader.next_batch(BATCH_SIZE);
- let array_len = array.unwrap().len();
- total_count += array_len;
- if array_len < BATCH_SIZE {
- break;
- }
- }
- total_count
-}
-
-fn create_int32_arrow_array_reader(
- page_iterator: impl PageIterator + 'static,
- column_desc: ColumnDescPtr,
-) -> impl ArrayReader {
- use parquet::arrow::arrow_array_reader::{ArrowArrayReader, PrimitiveArrayConverter};
- let converter = PrimitiveArrayConverter::<arrow::datatypes::Int32Type>::new();
- ArrowArrayReader::try_new(page_iterator, column_desc, converter, None).unwrap()
-}
-
-fn create_int32_primitive_array_reader(
- page_iterator: impl PageIterator + 'static,
- column_desc: ColumnDescPtr,
-) -> impl ArrayReader {
- use parquet::arrow::array_reader::PrimitiveArrayReader;
- PrimitiveArrayReader::<Int32Type>::new(Box::new(page_iterator), column_desc, None)
- .unwrap()
-}
-
-fn create_string_arrow_array_reader(
- page_iterator: impl PageIterator + 'static,
- column_desc: ColumnDescPtr,
-) -> impl ArrayReader {
- use parquet::arrow::arrow_array_reader::{ArrowArrayReader, StringArrayConverter};
- let converter = StringArrayConverter::new();
- ArrowArrayReader::try_new(page_iterator, column_desc, converter, None).unwrap()
-}
-
-fn create_string_complex_array_reader(
- page_iterator: impl PageIterator + 'static,
- column_desc: ColumnDescPtr,
-) -> impl ArrayReader {
- use parquet::arrow::array_reader::ComplexObjectArrayReader;
- use parquet::arrow::converter::{Utf8ArrayConverter, Utf8Converter};
- let converter = Utf8Converter::new(Utf8ArrayConverter {});
- ComplexObjectArrayReader::<parquet::data_type::ByteArrayType, Utf8Converter>::new(
- Box::new(page_iterator),
- column_desc,
- converter,
- None,
- )
- .unwrap()
-}
-
-fn add_benches(c: &mut Criterion) {
- const EXPECTED_VALUE_COUNT: usize =
- NUM_ROW_GROUPS * PAGES_PER_GROUP * VALUES_PER_PAGE;
- let mut group = c.benchmark_group("arrow_array_reader");
-
- let mut count: usize = 0;
-
- let schema = build_test_schema();
- let mandatory_int32_column_desc = schema.column(0);
- let optional_int32_column_desc = schema.column(1);
- let mandatory_string_column_desc = schema.column(2);
- // println!("mandatory_string_column_desc: {:?}", mandatory_string_column_desc);
- let optional_string_column_desc = schema.column(3);
- // println!("optional_string_column_desc: {:?}", optional_string_column_desc);
-
- // primitive / int32 benchmarks
- // =============================
-
- // int32, plain encoded, no NULLs
- let plain_int32_no_null_data = build_plain_encoded_int32_page_iterator(
- schema.clone(),
- mandatory_int32_column_desc.clone(),
- 0.0,
- );
- group.bench_function(
- "read Int32Array, plain encoded, mandatory, no NULLs - old",
- |b| {
- b.iter(|| {
- let array_reader = create_int32_primitive_array_reader(
- plain_int32_no_null_data.clone(),
- mandatory_int32_column_desc.clone(),
- );
- count = bench_array_reader(array_reader);
- })
- },
- );
- assert_eq!(count, EXPECTED_VALUE_COUNT);
-
- group.bench_function(
- "read Int32Array, plain encoded, mandatory, no NULLs - new",
- |b| {
- b.iter(|| {
- let array_reader = create_int32_arrow_array_reader(
- plain_int32_no_null_data.clone(),
- mandatory_int32_column_desc.clone(),
- );
- count = bench_array_reader(array_reader);
- })
- },
- );
- assert_eq!(count, EXPECTED_VALUE_COUNT);
-
- let plain_int32_no_null_data = build_plain_encoded_int32_page_iterator(
- schema.clone(),
- optional_int32_column_desc.clone(),
- 0.0,
- );
- group.bench_function(
- "read Int32Array, plain encoded, optional, no NULLs - old",
- |b| {
- b.iter(|| {
- let array_reader = create_int32_primitive_array_reader(
- plain_int32_no_null_data.clone(),
- optional_int32_column_desc.clone(),
- );
- count = bench_array_reader(array_reader);
- })
- },
- );
- assert_eq!(count, EXPECTED_VALUE_COUNT);
-
- group.bench_function(
- "read Int32Array, plain encoded, optional, no NULLs - new",
- |b| {
- b.iter(|| {
- let array_reader = create_int32_arrow_array_reader(
- plain_int32_no_null_data.clone(),
- optional_int32_column_desc.clone(),
- );
- count = bench_array_reader(array_reader);
- })
- },
- );
- assert_eq!(count, EXPECTED_VALUE_COUNT);
-
- // int32, plain encoded, half NULLs
- let plain_int32_half_null_data = build_plain_encoded_int32_page_iterator(
- schema.clone(),
- optional_int32_column_desc.clone(),
- 0.5,
- );
- group.bench_function(
- "read Int32Array, plain encoded, optional, half NULLs - old",
- |b| {
- b.iter(|| {
- let array_reader = create_int32_primitive_array_reader(
- plain_int32_half_null_data.clone(),
- optional_int32_column_desc.clone(),
- );
- count = bench_array_reader(array_reader);
- })
- },
- );
- assert_eq!(count, EXPECTED_VALUE_COUNT);
-
- group.bench_function(
- "read Int32Array, plain encoded, optional, half NULLs - new",
- |b| {
- b.iter(|| {
- let array_reader = create_int32_arrow_array_reader(
- plain_int32_half_null_data.clone(),
- optional_int32_column_desc.clone(),
- );
- count = bench_array_reader(array_reader);
- })
- },
- );
- assert_eq!(count, EXPECTED_VALUE_COUNT);
-
- // int32, dictionary encoded, no NULLs
- let dictionary_int32_no_null_data = build_dictionary_encoded_int32_page_iterator(
- schema.clone(),
- mandatory_int32_column_desc.clone(),
- 0.0,
- );
- group.bench_function(
- "read Int32Array, dictionary encoded, mandatory, no NULLs - old",
- |b| {
- b.iter(|| {
- let array_reader = create_int32_primitive_array_reader(
- dictionary_int32_no_null_data.clone(),
- mandatory_int32_column_desc.clone(),
- );
- count = bench_array_reader(array_reader);
- })
- },
- );
- assert_eq!(count, EXPECTED_VALUE_COUNT);
-
- group.bench_function(
- "read Int32Array, dictionary encoded, mandatory, no NULLs - new",
- |b| {
- b.iter(|| {
- let array_reader = create_int32_arrow_array_reader(
- dictionary_int32_no_null_data.clone(),
- mandatory_int32_column_desc.clone(),
- );
- count = bench_array_reader(array_reader);
- })
- },
- );
- assert_eq!(count, EXPECTED_VALUE_COUNT);
-
- let dictionary_int32_no_null_data = build_dictionary_encoded_int32_page_iterator(
- schema.clone(),
- optional_int32_column_desc.clone(),
- 0.0,
- );
- group.bench_function(
- "read Int32Array, dictionary encoded, optional, no NULLs - old",
- |b| {
- b.iter(|| {
- let array_reader = create_int32_primitive_array_reader(
- dictionary_int32_no_null_data.clone(),
- optional_int32_column_desc.clone(),
- );
- count = bench_array_reader(array_reader);
- })
- },
- );
- assert_eq!(count, EXPECTED_VALUE_COUNT);
-
- group.bench_function(
- "read Int32Array, dictionary encoded, optional, no NULLs - new",
- |b| {
- b.iter(|| {
- let array_reader = create_int32_arrow_array_reader(
- dictionary_int32_no_null_data.clone(),
- optional_int32_column_desc.clone(),
- );
- count = bench_array_reader(array_reader);
- })
- },
- );
- assert_eq!(count, EXPECTED_VALUE_COUNT);
-
- // int32, dictionary encoded, half NULLs
- let dictionary_int32_half_null_data = build_dictionary_encoded_int32_page_iterator(
- schema.clone(),
- optional_int32_column_desc.clone(),
- 0.5,
- );
- group.bench_function(
- "read Int32Array, dictionary encoded, optional, half NULLs - old",
- |b| {
- b.iter(|| {
- let array_reader = create_int32_primitive_array_reader(
- dictionary_int32_half_null_data.clone(),
- optional_int32_column_desc.clone(),
- );
- count = bench_array_reader(array_reader);
- })
- },
- );
- assert_eq!(count, EXPECTED_VALUE_COUNT);
-
- group.bench_function(
- "read Int32Array, dictionary encoded, optional, half NULLs - new",
- |b| {
- b.iter(|| {
- let array_reader = create_int32_arrow_array_reader(
- dictionary_int32_half_null_data.clone(),
- optional_int32_column_desc.clone(),
- );
- count = bench_array_reader(array_reader);
- })
- },
- );
- assert_eq!(count, EXPECTED_VALUE_COUNT);
-
- // string benchmarks
- //==============================
-
- // string, plain encoded, no NULLs
- let plain_string_no_null_data = build_plain_encoded_string_page_iterator(
- schema.clone(),
- mandatory_string_column_desc.clone(),
- 0.0,
- );
- group.bench_function(
- "read StringArray, plain encoded, mandatory, no NULLs - old",
- |b| {
- b.iter(|| {
- let array_reader = create_string_complex_array_reader(
- plain_string_no_null_data.clone(),
- mandatory_string_column_desc.clone(),
- );
- count = bench_array_reader(array_reader);
- })
- },
- );
- assert_eq!(count, EXPECTED_VALUE_COUNT);
-
- group.bench_function(
- "read StringArray, plain encoded, mandatory, no NULLs - new",
- |b| {
- b.iter(|| {
- let array_reader = create_string_arrow_array_reader(
- plain_string_no_null_data.clone(),
- mandatory_string_column_desc.clone(),
- );
- count = bench_array_reader(array_reader);
- })
- },
- );
- assert_eq!(count, EXPECTED_VALUE_COUNT);
-
- let plain_string_no_null_data = build_plain_encoded_string_page_iterator(
- schema.clone(),
- optional_string_column_desc.clone(),
- 0.0,
- );
- group.bench_function(
- "read StringArray, plain encoded, optional, no NULLs - old",
- |b| {
- b.iter(|| {
- let array_reader = create_string_complex_array_reader(
- plain_string_no_null_data.clone(),
- optional_string_column_desc.clone(),
- );
- count = bench_array_reader(array_reader);
- })
- },
- );
- assert_eq!(count, EXPECTED_VALUE_COUNT);
-
- group.bench_function(
- "read StringArray, plain encoded, optional, no NULLs - new",
- |b| {
- b.iter(|| {
- let array_reader = create_string_arrow_array_reader(
- plain_string_no_null_data.clone(),
- optional_string_column_desc.clone(),
- );
- count = bench_array_reader(array_reader);
- })
- },
- );
- assert_eq!(count, EXPECTED_VALUE_COUNT);
-
- // string, plain encoded, half NULLs
- let plain_string_half_null_data = build_plain_encoded_string_page_iterator(
- schema.clone(),
- optional_string_column_desc.clone(),
- 0.5,
- );
- group.bench_function(
- "read StringArray, plain encoded, optional, half NULLs - old",
- |b| {
- b.iter(|| {
- let array_reader = create_string_complex_array_reader(
- plain_string_half_null_data.clone(),
- optional_string_column_desc.clone(),
- );
- count = bench_array_reader(array_reader);
- })
- },
- );
- assert_eq!(count, EXPECTED_VALUE_COUNT);
-
- group.bench_function(
- "read StringArray, plain encoded, optional, half NULLs - new",
- |b| {
- b.iter(|| {
- let array_reader = create_string_arrow_array_reader(
- plain_string_half_null_data.clone(),
- optional_string_column_desc.clone(),
- );
- count = bench_array_reader(array_reader);
- })
- },
- );
- assert_eq!(count, EXPECTED_VALUE_COUNT);
-
- // string, dictionary encoded, no NULLs
- let dictionary_string_no_null_data = build_dictionary_encoded_string_page_iterator(
- schema.clone(),
- mandatory_string_column_desc.clone(),
- 0.0,
- );
- group.bench_function(
- "read StringArray, dictionary encoded, mandatory, no NULLs - old",
- |b| {
- b.iter(|| {
- let array_reader = create_string_complex_array_reader(
- dictionary_string_no_null_data.clone(),
- mandatory_string_column_desc.clone(),
- );
- count = bench_array_reader(array_reader);
- })
- },
- );
- assert_eq!(count, EXPECTED_VALUE_COUNT);
-
- group.bench_function(
- "read StringArray, dictionary encoded, mandatory, no NULLs - new",
- |b| {
- b.iter(|| {
- let array_reader = create_string_arrow_array_reader(
- dictionary_string_no_null_data.clone(),
- mandatory_string_column_desc.clone(),
- );
- count = bench_array_reader(array_reader);
- })
- },
- );
- assert_eq!(count, EXPECTED_VALUE_COUNT);
-
- let dictionary_string_no_null_data = build_dictionary_encoded_string_page_iterator(
- schema.clone(),
- optional_string_column_desc.clone(),
- 0.0,
- );
- group.bench_function(
- "read StringArray, dictionary encoded, optional, no NULLs - old",
- |b| {
- b.iter(|| {
- let array_reader = create_string_complex_array_reader(
- dictionary_string_no_null_data.clone(),
- optional_string_column_desc.clone(),
- );
- count = bench_array_reader(array_reader);
- })
- },
- );
- assert_eq!(count, EXPECTED_VALUE_COUNT);
-
- group.bench_function(
- "read StringArray, dictionary encoded, optional, no NULLs - new",
- |b| {
- b.iter(|| {
- let array_reader = create_string_arrow_array_reader(
- dictionary_string_no_null_data.clone(),
- optional_string_column_desc.clone(),
- );
- count = bench_array_reader(array_reader);
- })
- },
- );
- assert_eq!(count, EXPECTED_VALUE_COUNT);
-
- // string, dictionary encoded, half NULLs
- let dictionary_string_half_null_data = build_dictionary_encoded_string_page_iterator(
- schema,
- optional_string_column_desc.clone(),
- 0.5,
- );
- group.bench_function(
- "read StringArray, dictionary encoded, optional, half NULLs - old",
- |b| {
- b.iter(|| {
- let array_reader = create_string_complex_array_reader(
- dictionary_string_half_null_data.clone(),
- optional_string_column_desc.clone(),
- );
- count = bench_array_reader(array_reader);
- })
- },
- );
- assert_eq!(count, EXPECTED_VALUE_COUNT);
-
- group.bench_function(
- "read StringArray, dictionary encoded, optional, half NULLs - new",
- |b| {
- b.iter(|| {
- let array_reader = create_string_arrow_array_reader(
- dictionary_string_half_null_data.clone(),
- optional_string_column_desc.clone(),
- );
- count = bench_array_reader(array_reader);
- })
- },
- );
- assert_eq!(count, EXPECTED_VALUE_COUNT);
-
- group.finish();
-}
-
-criterion_group!(benches, add_benches);
-criterion_main!(benches);
diff --git a/parquet/benches/arrow_writer.rs b/parquet/benches/arrow_writer.rs
deleted file mode 100644
index 069ed39..0000000
--- a/parquet/benches/arrow_writer.rs
+++ /dev/null
@@ -1,202 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-#[macro_use]
-extern crate criterion;
-use criterion::{Criterion, Throughput};
-
-extern crate arrow;
-extern crate parquet;
-
-use std::sync::Arc;
-
-use arrow::datatypes::*;
-use arrow::{record_batch::RecordBatch, util::data_gen::*};
-use parquet::{
- arrow::ArrowWriter, errors::Result, file::writer::InMemoryWriteableCursor,
-};
-
-fn create_primitive_bench_batch(
- size: usize,
- null_density: f32,
- true_density: f32,
-) -> Result<RecordBatch> {
- let fields = vec![
- Field::new("_1", DataType::Int8, true),
- Field::new("_2", DataType::Int16, true),
- Field::new("_3", DataType::Int32, true),
- Field::new("_4", DataType::Int64, true),
- Field::new("_5", DataType::UInt8, true),
- Field::new("_6", DataType::UInt16, true),
- Field::new("_7", DataType::UInt32, true),
- Field::new("_8", DataType::UInt64, true),
- Field::new("_9", DataType::Float32, true),
- Field::new("_10", DataType::Float64, true),
- Field::new("_11", DataType::Date32, true),
- Field::new("_12", DataType::Date64, true),
- Field::new("_13", DataType::Time32(TimeUnit::Second), true),
- Field::new("_14", DataType::Time32(TimeUnit::Millisecond), true),
- Field::new("_15", DataType::Time64(TimeUnit::Microsecond), true),
- Field::new("_16", DataType::Time64(TimeUnit::Nanosecond), true),
- Field::new("_17", DataType::Utf8, true),
- Field::new("_18", DataType::LargeUtf8, true),
- Field::new("_19", DataType::Boolean, true),
- ];
- let schema = Schema::new(fields);
- Ok(create_random_batch(
- Arc::new(schema),
- size,
- null_density,
- true_density,
- )?)
-}
-
-fn _create_nested_bench_batch(
- size: usize,
- null_density: f32,
- true_density: f32,
-) -> Result<RecordBatch> {
- let fields = vec![
- Field::new(
- "_1",
- DataType::Struct(vec![
- Field::new("_1", DataType::Int8, true),
- Field::new(
- "_2",
- DataType::Struct(vec![
- Field::new("_1", DataType::Int8, true),
- Field::new(
- "_1",
- DataType::Struct(vec![
- Field::new("_1", DataType::Int8, true),
- Field::new("_2", DataType::Utf8, true),
- ]),
- true,
- ),
- Field::new("_2", DataType::UInt8, true),
- ]),
- true,
- ),
- ]),
- true,
- ),
- Field::new(
- "_2",
- DataType::LargeList(Box::new(Field::new(
- "item",
- DataType::List(Box::new(Field::new(
- "item",
- DataType::Struct(vec![
- Field::new(
- "_1",
- DataType::Struct(vec![
- Field::new("_1", DataType::Int8, true),
- Field::new("_2", DataType::Int16, true),
- Field::new("_3", DataType::Int32, true),
- ]),
- true,
- ),
- Field::new(
- "_2",
- DataType::List(Box::new(Field::new(
- "",
- DataType::FixedSizeBinary(2),
- true,
- ))),
- true,
- ),
- ]),
- true,
- ))),
- true,
- ))),
- true,
- ),
- ];
- let schema = Schema::new(fields);
- Ok(create_random_batch(
- Arc::new(schema),
- size,
- null_density,
- true_density,
- )?)
-}
-
-#[inline]
-fn write_batch(batch: &RecordBatch) -> Result<()> {
- // Write batch to an in-memory writer
- let cursor = InMemoryWriteableCursor::default();
- let mut writer = ArrowWriter::try_new(cursor, batch.schema(), None)?;
-
- writer.write(&batch)?;
- writer.close()?;
- Ok(())
-}
-
-fn bench_primitive_writer(c: &mut Criterion) {
- let batch = create_primitive_bench_batch(1024, 0.25, 0.75).unwrap();
- let mut group = c.benchmark_group("write_batch primitive");
- group.throughput(Throughput::Bytes(
- batch
- .columns()
- .iter()
- .map(|f| f.get_array_memory_size() as u64)
- .sum(),
- ));
- group.bench_function("1024 values", |b| b.iter(|| write_batch(&batch).unwrap()));
-
- let batch = create_primitive_bench_batch(4096, 0.25, 0.75).unwrap();
- group.throughput(Throughput::Bytes(
- batch
- .columns()
- .iter()
- .map(|f| f.get_array_memory_size() as u64)
- .sum(),
- ));
- group.bench_function("4096 values", |b| b.iter(|| write_batch(&batch).unwrap()));
-
- group.finish();
-}
-
-// This bench triggers a write error, it is ignored for now
-fn _bench_nested_writer(c: &mut Criterion) {
- let batch = _create_nested_bench_batch(1024, 0.25, 0.75).unwrap();
- let mut group = c.benchmark_group("write_batch nested");
- group.throughput(Throughput::Bytes(
- batch
- .columns()
- .iter()
- .map(|f| f.get_array_memory_size() as u64)
- .sum(),
- ));
- group.bench_function("1024 values", |b| b.iter(|| write_batch(&batch).unwrap()));
-
- let batch = create_primitive_bench_batch(4096, 0.25, 0.75).unwrap();
- group.throughput(Throughput::Bytes(
- batch
- .columns()
- .iter()
- .map(|f| f.get_array_memory_size() as u64)
- .sum(),
- ));
- group.bench_function("4096 values", |b| b.iter(|| write_batch(&batch).unwrap()));
-
- group.finish();
-}
-
-criterion_group!(benches, bench_primitive_writer);
-criterion_main!(benches);
diff --git a/parquet/build.rs b/parquet/build.rs
deleted file mode 100644
index b42b2a4..0000000
--- a/parquet/build.rs
+++ /dev/null
@@ -1,43 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-use std::process::Command;
-
-fn main() {
- // Set Parquet version, build hash and "created by" string.
- let version = env!("CARGO_PKG_VERSION");
- let mut created_by = format!("parquet-rs version {}", version);
- if let Ok(git_hash) = run(Command::new("git").arg("rev-parse").arg("HEAD")) {
- created_by.push_str(format!(" (build {})", git_hash).as_str());
- println!("cargo:rustc-env=PARQUET_BUILD={}", git_hash);
- }
- println!("cargo:rustc-env=PARQUET_VERSION={}", version);
- println!("cargo:rustc-env=PARQUET_CREATED_BY={}", created_by);
-}
-
-/// Runs command and returns either content of stdout for successful execution,
-/// or an error message otherwise.
-fn run(command: &mut Command) -> Result<String, String> {
- println!("Running: `{:?}`", command);
- match command.output() {
- Ok(ref output) if output.status.success() => {
- Ok(String::from_utf8_lossy(&output.stdout).trim().to_string())
- }
- Ok(ref output) => Err(format!("Failed: `{:?}` ({})", command, output.status)),
- Err(error) => Err(format!("Failed: `{:?}` ({})", command, error)),
- }
-}
diff --git a/parquet/src/arrow/array_reader.rs b/parquet/src/arrow/array_reader.rs
deleted file mode 100644
index bd57cf3..0000000
--- a/parquet/src/arrow/array_reader.rs
+++ /dev/null
@@ -1,2556 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-use std::cmp::{max, min};
-use std::collections::{HashMap, HashSet};
-use std::marker::PhantomData;
-use std::mem::size_of;
-use std::result::Result::Ok;
-use std::sync::Arc;
-use std::vec::Vec;
-
-use arrow::array::{
- new_empty_array, Array, ArrayData, ArrayDataBuilder, ArrayRef, BinaryArray,
- BinaryBuilder, BooleanArray, BooleanBufferBuilder, BooleanBuilder, DecimalBuilder,
- FixedSizeBinaryArray, FixedSizeBinaryBuilder, GenericListArray, Int16BufferBuilder,
- Int32Array, Int64Array, OffsetSizeTrait, PrimitiveArray, PrimitiveBuilder,
- StringArray, StringBuilder, StructArray,
-};
-use arrow::buffer::{Buffer, MutableBuffer};
-use arrow::datatypes::{
- ArrowPrimitiveType, BooleanType as ArrowBooleanType, DataType as ArrowType,
- Date32Type as ArrowDate32Type, Date64Type as ArrowDate64Type,
- DurationMicrosecondType as ArrowDurationMicrosecondType,
- DurationMillisecondType as ArrowDurationMillisecondType,
- DurationNanosecondType as ArrowDurationNanosecondType,
- DurationSecondType as ArrowDurationSecondType, Field,
- Float32Type as ArrowFloat32Type, Float64Type as ArrowFloat64Type,
- Int16Type as ArrowInt16Type, Int32Type as ArrowInt32Type,
- Int64Type as ArrowInt64Type, Int8Type as ArrowInt8Type, IntervalUnit, Schema,
- Time32MillisecondType as ArrowTime32MillisecondType,
- Time32SecondType as ArrowTime32SecondType,
- Time64MicrosecondType as ArrowTime64MicrosecondType,
- Time64NanosecondType as ArrowTime64NanosecondType, TimeUnit as ArrowTimeUnit,
- TimestampMicrosecondType as ArrowTimestampMicrosecondType,
- TimestampMillisecondType as ArrowTimestampMillisecondType,
- TimestampNanosecondType as ArrowTimestampNanosecondType,
- TimestampSecondType as ArrowTimestampSecondType, ToByteSlice,
- UInt16Type as ArrowUInt16Type, UInt32Type as ArrowUInt32Type,
- UInt64Type as ArrowUInt64Type, UInt8Type as ArrowUInt8Type,
-};
-use arrow::util::bit_util;
-
-use crate::arrow::converter::{
- BinaryArrayConverter, BinaryConverter, Converter, DecimalArrayConverter,
- DecimalConverter, FixedLenBinaryConverter, FixedSizeArrayConverter,
- Int96ArrayConverter, Int96Converter, IntervalDayTimeArrayConverter,
- IntervalDayTimeConverter, IntervalYearMonthArrayConverter,
- IntervalYearMonthConverter, LargeBinaryArrayConverter, LargeBinaryConverter,
- LargeUtf8ArrayConverter, LargeUtf8Converter,
-};
-use crate::arrow::record_reader::RecordReader;
-use crate::arrow::schema::parquet_to_arrow_field;
-use crate::basic::{ConvertedType, Repetition, Type as PhysicalType};
-use crate::column::page::PageIterator;
-use crate::column::reader::ColumnReaderImpl;
-use crate::data_type::{
- BoolType, ByteArrayType, DataType, DoubleType, FixedLenByteArrayType, FloatType,
- Int32Type, Int64Type, Int96Type,
-};
-use crate::errors::{ParquetError, ParquetError::ArrowError, Result};
-use crate::file::reader::{FilePageIterator, FileReader};
-use crate::schema::types::{
- ColumnDescPtr, ColumnDescriptor, ColumnPath, SchemaDescPtr, Type, TypePtr,
-};
-use crate::schema::visitor::TypeVisitor;
-use std::any::Any;
-
-/// Array reader reads parquet data into arrow array.
-pub trait ArrayReader {
- fn as_any(&self) -> &dyn Any;
-
- /// Returns the arrow type of this array reader.
- fn get_data_type(&self) -> &ArrowType;
-
- /// Reads at most `batch_size` records into an arrow array and return it.
- fn next_batch(&mut self, batch_size: usize) -> Result<ArrayRef>;
-
- /// Returns the definition levels of data from last call of `next_batch`.
- /// The result is used by parent array reader to calculate its own definition
- /// levels and repetition levels, so that its parent can calculate null bitmap.
- fn get_def_levels(&self) -> Option<&[i16]>;
-
- /// Return the repetition levels of data from last call of `next_batch`.
- /// The result is used by parent array reader to calculate its own definition
- /// levels and repetition levels, so that its parent can calculate null bitmap.
- fn get_rep_levels(&self) -> Option<&[i16]>;
-}
-
-/// A NullArrayReader reads Parquet columns stored as null int32s with an Arrow
-/// NullArray type.
-pub struct NullArrayReader<T: DataType> {
- data_type: ArrowType,
- pages: Box<dyn PageIterator>,
- def_levels_buffer: Option<Buffer>,
- rep_levels_buffer: Option<Buffer>,
- column_desc: ColumnDescPtr,
- record_reader: RecordReader<T>,
- _type_marker: PhantomData<T>,
-}
-
-impl<T: DataType> NullArrayReader<T> {
- /// Construct null array reader.
- pub fn new(
- mut pages: Box<dyn PageIterator>,
- column_desc: ColumnDescPtr,
- ) -> Result<Self> {
- let mut record_reader = RecordReader::<T>::new(column_desc.clone());
- if let Some(page_reader) = pages.next() {
- record_reader.set_page_reader(page_reader?)?;
- }
-
- Ok(Self {
- data_type: ArrowType::Null,
- pages,
- def_levels_buffer: None,
- rep_levels_buffer: None,
- column_desc,
- record_reader,
- _type_marker: PhantomData,
- })
- }
-}
-
-/// Implementation of primitive array reader.
-impl<T: DataType> ArrayReader for NullArrayReader<T> {
- fn as_any(&self) -> &dyn Any {
- self
- }
-
- /// Returns data type of primitive array.
- fn get_data_type(&self) -> &ArrowType {
- &self.data_type
- }
-
- /// Reads at most `batch_size` records into array.
- fn next_batch(&mut self, batch_size: usize) -> Result<ArrayRef> {
- let mut records_read = 0usize;
- while records_read < batch_size {
- let records_to_read = batch_size - records_read;
-
- // NB can be 0 if at end of page
- let records_read_once = self.record_reader.read_records(records_to_read)?;
- records_read += records_read_once;
-
- // Record reader exhausted
- if records_read_once < records_to_read {
- if let Some(page_reader) = self.pages.next() {
- // Read from new page reader
- self.record_reader.set_page_reader(page_reader?)?;
- } else {
- // Page reader also exhausted
- break;
- }
- }
- }
-
- // convert to arrays
- let array = arrow::array::NullArray::new(records_read);
-
- // save definition and repetition buffers
- self.def_levels_buffer = self.record_reader.consume_def_levels()?;
- self.rep_levels_buffer = self.record_reader.consume_rep_levels()?;
- self.record_reader.reset();
- Ok(Arc::new(array))
- }
-
- fn get_def_levels(&self) -> Option<&[i16]> {
- self.def_levels_buffer
- .as_ref()
- .map(|buf| unsafe { buf.typed_data() })
- }
-
- fn get_rep_levels(&self) -> Option<&[i16]> {
- self.rep_levels_buffer
- .as_ref()
- .map(|buf| unsafe { buf.typed_data() })
- }
-}
-
-/// Primitive array readers are leaves of array reader tree. They accept page iterator
-/// and read them into primitive arrays.
-pub struct PrimitiveArrayReader<T: DataType> {
- data_type: ArrowType,
- pages: Box<dyn PageIterator>,
- def_levels_buffer: Option<Buffer>,
- rep_levels_buffer: Option<Buffer>,
- column_desc: ColumnDescPtr,
- record_reader: RecordReader<T>,
- _type_marker: PhantomData<T>,
-}
-
-impl<T: DataType> PrimitiveArrayReader<T> {
- /// Construct primitive array reader.
- pub fn new(
- mut pages: Box<dyn PageIterator>,
- column_desc: ColumnDescPtr,
- arrow_type: Option<ArrowType>,
- ) -> Result<Self> {
- // Check if Arrow type is specified, else create it from Parquet type
- let data_type = match arrow_type {
- Some(t) => t,
- None => parquet_to_arrow_field(column_desc.as_ref())?
- .data_type()
- .clone(),
- };
-
- let mut record_reader = RecordReader::<T>::new(column_desc.clone());
- if let Some(page_reader) = pages.next() {
- record_reader.set_page_reader(page_reader?)?;
- }
-
- Ok(Self {
- data_type,
- pages,
- def_levels_buffer: None,
- rep_levels_buffer: None,
- column_desc,
- record_reader,
- _type_marker: PhantomData,
- })
- }
-}
-
-/// Implementation of primitive array reader.
-impl<T: DataType> ArrayReader for PrimitiveArrayReader<T> {
- fn as_any(&self) -> &dyn Any {
- self
- }
-
- /// Returns data type of primitive array.
- fn get_data_type(&self) -> &ArrowType {
- &self.data_type
- }
-
- /// Reads at most `batch_size` records into array.
- fn next_batch(&mut self, batch_size: usize) -> Result<ArrayRef> {
- let mut records_read = 0usize;
- while records_read < batch_size {
- let records_to_read = batch_size - records_read;
-
- // NB can be 0 if at end of page
- let records_read_once = self.record_reader.read_records(records_to_read)?;
- records_read += records_read_once;
-
- // Record reader exhausted
- if records_read_once < records_to_read {
- if let Some(page_reader) = self.pages.next() {
- // Read from new page reader
- self.record_reader.set_page_reader(page_reader?)?;
- } else {
- // Page reader also exhausted
- break;
- }
- }
- }
-
- let target_type = self.get_data_type().clone();
- let arrow_data_type = match T::get_physical_type() {
- PhysicalType::BOOLEAN => ArrowBooleanType::DATA_TYPE,
- PhysicalType::INT32 => {
- match target_type {
- ArrowType::UInt32 => {
- // follow C++ implementation and use overflow/reinterpret cast from i32 to u32 which will map
- // `i32::MIN..0` to `(i32::MAX as u32)..u32::MAX`
- ArrowUInt32Type::DATA_TYPE
- }
- _ => ArrowInt32Type::DATA_TYPE,
- }
- }
- PhysicalType::INT64 => {
- match target_type {
- ArrowType::UInt64 => {
- // follow C++ implementation and use overflow/reinterpret cast from i64 to u64 which will map
- // `i64::MIN..0` to `(i64::MAX as u64)..u64::MAX`
- ArrowUInt64Type::DATA_TYPE
- }
- _ => ArrowInt64Type::DATA_TYPE,
- }
- }
- PhysicalType::FLOAT => ArrowFloat32Type::DATA_TYPE,
- PhysicalType::DOUBLE => ArrowFloat64Type::DATA_TYPE,
- PhysicalType::INT96
- | PhysicalType::BYTE_ARRAY
- | PhysicalType::FIXED_LEN_BYTE_ARRAY => {
- unreachable!(
- "PrimitiveArrayReaders don't support complex physical types"
- );
- }
- };
-
- // Convert to arrays by using the Parquet phyisical type.
- // The physical types are then cast to Arrow types if necessary
-
- let mut record_data = self.record_reader.consume_record_data()?;
-
- if T::get_physical_type() == PhysicalType::BOOLEAN {
- let mut boolean_buffer = BooleanBufferBuilder::new(record_data.len());
-
- for e in record_data.as_slice() {
- boolean_buffer.append(*e > 0);
- }
- record_data = boolean_buffer.finish();
- }
-
- let mut array_data = ArrayDataBuilder::new(arrow_data_type)
- .len(self.record_reader.num_values())
- .add_buffer(record_data);
-
- if let Some(b) = self.record_reader.consume_bitmap_buffer()? {
- array_data = array_data.null_bit_buffer(b);
- }
-
- let array = match T::get_physical_type() {
- PhysicalType::BOOLEAN => {
- Arc::new(BooleanArray::from(array_data.build())) as ArrayRef
- }
- PhysicalType::INT32 => {
- Arc::new(PrimitiveArray::<ArrowInt32Type>::from(array_data.build()))
- as ArrayRef
- }
- PhysicalType::INT64 => {
- Arc::new(PrimitiveArray::<ArrowInt64Type>::from(array_data.build()))
- as ArrayRef
- }
- PhysicalType::FLOAT => {
- Arc::new(PrimitiveArray::<ArrowFloat32Type>::from(array_data.build()))
- as ArrayRef
- }
- PhysicalType::DOUBLE => {
- Arc::new(PrimitiveArray::<ArrowFloat64Type>::from(array_data.build()))
- as ArrayRef
- }
- PhysicalType::INT96
- | PhysicalType::BYTE_ARRAY
- | PhysicalType::FIXED_LEN_BYTE_ARRAY => {
- unreachable!(
- "PrimitiveArrayReaders don't support complex physical types"
- );
- }
- };
-
- // cast to Arrow type
- // We make a strong assumption here that the casts should be infallible.
- // If the cast fails because of incompatible datatypes, then there might
- // be a bigger problem with how Arrow schemas are converted to Parquet.
- //
- // As there is not always a 1:1 mapping between Arrow and Parquet, there
- // are datatypes which we must convert explicitly.
- // These are:
- // - date64: we should cast int32 to date32, then date32 to date64.
- let array = match target_type {
- ArrowType::Date64 => {
- // this is cheap as it internally reinterprets the data
- let a = arrow::compute::cast(&array, &ArrowType::Date32)?;
- arrow::compute::cast(&a, &target_type)?
- }
- ArrowType::Decimal(p, s) => {
- let mut builder = DecimalBuilder::new(array.len(), p, s);
- match array.data_type() {
- ArrowType::Int32 => {
- let values = array.as_any().downcast_ref::<Int32Array>().unwrap();
- for maybe_value in values.iter() {
- match maybe_value {
- Some(value) => builder.append_value(value as i128)?,
- None => builder.append_null()?,
- }
- }
- }
- ArrowType::Int64 => {
- let values = array.as_any().downcast_ref::<Int64Array>().unwrap();
- for maybe_value in values.iter() {
- match maybe_value {
- Some(value) => builder.append_value(value as i128)?,
- None => builder.append_null()?,
- }
- }
- }
- _ => {
- return Err(ArrowError(format!(
- "Cannot convert {:?} to decimal",
- array.data_type()
- )))
- }
- }
- Arc::new(builder.finish()) as ArrayRef
- }
- _ => arrow::compute::cast(&array, &target_type)?,
- };
-
- // save definition and repetition buffers
- self.def_levels_buffer = self.record_reader.consume_def_levels()?;
- self.rep_levels_buffer = self.record_reader.consume_rep_levels()?;
- self.record_reader.reset();
- Ok(array)
- }
-
- fn get_def_levels(&self) -> Option<&[i16]> {
- self.def_levels_buffer
- .as_ref()
- .map(|buf| unsafe { buf.typed_data() })
- }
-
- fn get_rep_levels(&self) -> Option<&[i16]> {
- self.rep_levels_buffer
- .as_ref()
- .map(|buf| unsafe { buf.typed_data() })
- }
-}
-
-/// Primitive array readers are leaves of array reader tree. They accept page iterator
-/// and read them into primitive arrays.
-pub struct ComplexObjectArrayReader<T, C>
-where
- T: DataType,
- C: Converter<Vec<Option<T::T>>, ArrayRef> + 'static,
-{
- data_type: ArrowType,
- pages: Box<dyn PageIterator>,
- def_levels_buffer: Option<Vec<i16>>,
- rep_levels_buffer: Option<Vec<i16>>,
- column_desc: ColumnDescPtr,
- column_reader: Option<ColumnReaderImpl<T>>,
- converter: C,
- _parquet_type_marker: PhantomData<T>,
- _converter_marker: PhantomData<C>,
-}
-
-impl<T, C> ArrayReader for ComplexObjectArrayReader<T, C>
-where
- T: DataType,
- C: Converter<Vec<Option<T::T>>, ArrayRef> + 'static,
-{
- fn as_any(&self) -> &dyn Any {
- self
- }
-
- fn get_data_type(&self) -> &ArrowType {
- &self.data_type
- }
-
- fn next_batch(&mut self, batch_size: usize) -> Result<ArrayRef> {
- // Try to initialize column reader
- if self.column_reader.is_none() {
- self.next_column_reader()?;
- }
-
- let mut data_buffer: Vec<T::T> = Vec::with_capacity(batch_size);
- data_buffer.resize_with(batch_size, T::T::default);
-
- let mut def_levels_buffer = if self.column_desc.max_def_level() > 0 {
- let mut buf: Vec<i16> = Vec::with_capacity(batch_size);
- buf.resize_with(batch_size, || 0);
- Some(buf)
- } else {
- None
- };
-
- let mut rep_levels_buffer = if self.column_desc.max_rep_level() > 0 {
- let mut buf: Vec<i16> = Vec::with_capacity(batch_size);
- buf.resize_with(batch_size, || 0);
- Some(buf)
- } else {
- None
- };
-
- let mut num_read = 0;
-
- while self.column_reader.is_some() && num_read < batch_size {
- let num_to_read = batch_size - num_read;
- let cur_data_buf = &mut data_buffer[num_read..];
- let cur_def_levels_buf =
- def_levels_buffer.as_mut().map(|b| &mut b[num_read..]);
- let cur_rep_levels_buf =
- rep_levels_buffer.as_mut().map(|b| &mut b[num_read..]);
- let (data_read, levels_read) =
- self.column_reader.as_mut().unwrap().read_batch(
- num_to_read,
- cur_def_levels_buf,
- cur_rep_levels_buf,
- cur_data_buf,
- )?;
-
- // Fill space
- if levels_read > data_read {
- def_levels_buffer.iter().for_each(|def_levels_buffer| {
- let (mut level_pos, mut data_pos) = (levels_read, data_read);
- while level_pos > 0 && data_pos > 0 {
- if def_levels_buffer[num_read + level_pos - 1]
- == self.column_desc.max_def_level()
- {
- cur_data_buf.swap(level_pos - 1, data_pos - 1);
- level_pos -= 1;
- data_pos -= 1;
- } else {
- level_pos -= 1;
- }
- }
- });
- }
-
- let values_read = max(levels_read, data_read);
- num_read += values_read;
- // current page exhausted && page iterator exhausted
- if values_read < num_to_read && !self.next_column_reader()? {
- break;
- }
- }
-
- data_buffer.truncate(num_read);
- def_levels_buffer
- .iter_mut()
- .for_each(|buf| buf.truncate(num_read));
- rep_levels_buffer
- .iter_mut()
- .for_each(|buf| buf.truncate(num_read));
-
- self.def_levels_buffer = def_levels_buffer;
- self.rep_levels_buffer = rep_levels_buffer;
-
- let data: Vec<Option<T::T>> = if self.def_levels_buffer.is_some() {
- data_buffer
- .into_iter()
- .zip(self.def_levels_buffer.as_ref().unwrap().iter())
- .map(|(t, def_level)| {
- if *def_level == self.column_desc.max_def_level() {
- Some(t)
- } else {
- None
- }
- })
- .collect()
- } else {
- data_buffer.into_iter().map(Some).collect()
- };
-
- let mut array = self.converter.convert(data)?;
-
- if let ArrowType::Dictionary(_, _) = self.data_type {
- array = arrow::compute::cast(&array, &self.data_type)?;
- }
-
- Ok(array)
- }
-
- fn get_def_levels(&self) -> Option<&[i16]> {
- self.def_levels_buffer.as_deref()
- }
-
- fn get_rep_levels(&self) -> Option<&[i16]> {
- self.rep_levels_buffer.as_deref()
- }
-}
-
-impl<T, C> ComplexObjectArrayReader<T, C>
-where
- T: DataType,
- C: Converter<Vec<Option<T::T>>, ArrayRef> + 'static,
-{
- pub fn new(
- pages: Box<dyn PageIterator>,
- column_desc: ColumnDescPtr,
- converter: C,
- arrow_type: Option<ArrowType>,
- ) -> Result<Self> {
- let data_type = match arrow_type {
- Some(t) => t,
- None => parquet_to_arrow_field(column_desc.as_ref())?
- .data_type()
- .clone(),
- };
-
- Ok(Self {
- data_type,
- pages,
- def_levels_buffer: None,
- rep_levels_buffer: None,
- column_desc,
- column_reader: None,
- converter,
- _parquet_type_marker: PhantomData,
- _converter_marker: PhantomData,
- })
- }
-
- fn next_column_reader(&mut self) -> Result<bool> {
- Ok(match self.pages.next() {
- Some(page) => {
- self.column_reader =
- Some(ColumnReaderImpl::<T>::new(self.column_desc.clone(), page?));
- true
- }
- None => false,
- })
- }
-}
-
-/// Implementation of list array reader.
-pub struct ListArrayReader<OffsetSize: OffsetSizeTrait> {
- item_reader: Box<dyn ArrayReader>,
- data_type: ArrowType,
- item_type: ArrowType,
- list_def_level: i16,
- list_rep_level: i16,
- list_empty_def_level: i16,
- list_null_def_level: i16,
- def_level_buffer: Option<Buffer>,
- rep_level_buffer: Option<Buffer>,
- _marker: PhantomData<OffsetSize>,
-}
-
-impl<OffsetSize: OffsetSizeTrait> ListArrayReader<OffsetSize> {
- /// Construct list array reader.
- pub fn new(
- item_reader: Box<dyn ArrayReader>,
- data_type: ArrowType,
- item_type: ArrowType,
- def_level: i16,
- rep_level: i16,
- list_null_def_level: i16,
- list_empty_def_level: i16,
- ) -> Self {
- Self {
- item_reader,
- data_type,
- item_type,
- list_def_level: def_level,
- list_rep_level: rep_level,
- list_null_def_level,
- list_empty_def_level,
- def_level_buffer: None,
- rep_level_buffer: None,
- _marker: PhantomData,
- }
- }
-}
-
-macro_rules! remove_primitive_array_indices {
- ($arr: expr, $item_type:ty, $indices:expr) => {{
- let array_data = match $arr.as_any().downcast_ref::<PrimitiveArray<$item_type>>() {
- Some(a) => a,
- _ => return Err(ParquetError::General(format!("Error generating next batch for ListArray: {:?} cannot be downcast to PrimitiveArray", $arr))),
- };
- let mut builder = PrimitiveBuilder::<$item_type>::new($arr.len());
- for i in 0..array_data.len() {
- if !$indices.contains(&i) {
- if array_data.is_null(i) {
- builder.append_null()?;
- } else {
- builder.append_value(array_data.value(i))?;
- }
- }
- }
- Ok(Arc::new(builder.finish()))
- }};
-}
-
-macro_rules! remove_array_indices_custom_builder {
- ($arr: expr, $array_type:ty, $item_builder:ident, $indices:expr) => {{
- let array_data = match $arr.as_any().downcast_ref::<$array_type>() {
- Some(a) => a,
- _ => return Err(ParquetError::General(format!("Error generating next batch for ListArray: {:?} cannot be downcast to PrimitiveArray", $arr))),
- };
- let mut builder = $item_builder::new(array_data.len());
-
- for i in 0..array_data.len() {
- if !$indices.contains(&i) {
- if array_data.is_null(i) {
- builder.append_null()?;
- } else {
- builder.append_value(array_data.value(i))?;
- }
- }
- }
- Ok(Arc::new(builder.finish()))
- }};
-}
-
-macro_rules! remove_fixed_size_binary_array_indices {
- ($arr: expr, $array_type:ty, $item_builder:ident, $indices:expr, $len:expr) => {{
- let array_data = match $arr.as_any().downcast_ref::<$array_type>() {
- Some(a) => a,
- _ => return Err(ParquetError::General(format!("Error generating next batch for ListArray: {:?} cannot be downcast to PrimitiveArray", $arr))),
- };
- let mut builder = FixedSizeBinaryBuilder::new(array_data.len(), $len);
- for i in 0..array_data.len() {
- if !$indices.contains(&i) {
- if array_data.is_null(i) {
- builder.append_null()?;
- } else {
- builder.append_value(array_data.value(i))?;
- }
- }
- }
- Ok(Arc::new(builder.finish()))
- }};
-}
-
-fn remove_indices(
- arr: ArrayRef,
- item_type: ArrowType,
- indices: Vec<usize>,
-) -> Result<ArrayRef> {
- match item_type {
- ArrowType::UInt8 => remove_primitive_array_indices!(arr, ArrowUInt8Type, indices),
- ArrowType::UInt16 => {
- remove_primitive_array_indices!(arr, ArrowUInt16Type, indices)
- }
- ArrowType::UInt32 => {
- remove_primitive_array_indices!(arr, ArrowUInt32Type, indices)
- }
- ArrowType::UInt64 => {
- remove_primitive_array_indices!(arr, ArrowUInt64Type, indices)
- }
- ArrowType::Int8 => remove_primitive_array_indices!(arr, ArrowInt8Type, indices),
- ArrowType::Int16 => remove_primitive_array_indices!(arr, ArrowInt16Type, indices),
- ArrowType::Int32 => remove_primitive_array_indices!(arr, ArrowInt32Type, indices),
- ArrowType::Int64 => remove_primitive_array_indices!(arr, ArrowInt64Type, indices),
- ArrowType::Float32 => {
- remove_primitive_array_indices!(arr, ArrowFloat32Type, indices)
- }
- ArrowType::Float64 => {
- remove_primitive_array_indices!(arr, ArrowFloat64Type, indices)
- }
- ArrowType::Boolean => {
- remove_array_indices_custom_builder!(
- arr,
- BooleanArray,
- BooleanBuilder,
- indices
- )
- }
- ArrowType::Date32 => {
- remove_primitive_array_indices!(arr, ArrowDate32Type, indices)
- }
- ArrowType::Date64 => {
- remove_primitive_array_indices!(arr, ArrowDate64Type, indices)
- }
- ArrowType::Time32(ArrowTimeUnit::Second) => {
- remove_primitive_array_indices!(arr, ArrowTime32SecondType, indices)
- }
- ArrowType::Time32(ArrowTimeUnit::Millisecond) => {
- remove_primitive_array_indices!(arr, ArrowTime32MillisecondType, indices)
- }
- ArrowType::Time64(ArrowTimeUnit::Microsecond) => {
- remove_primitive_array_indices!(arr, ArrowTime64MicrosecondType, indices)
- }
- ArrowType::Time64(ArrowTimeUnit::Nanosecond) => {
- remove_primitive_array_indices!(arr, ArrowTime64NanosecondType, indices)
- }
- ArrowType::Duration(ArrowTimeUnit::Second) => {
- remove_primitive_array_indices!(arr, ArrowDurationSecondType, indices)
- }
- ArrowType::Duration(ArrowTimeUnit::Millisecond) => {
- remove_primitive_array_indices!(arr, ArrowDurationMillisecondType, indices)
- }
- ArrowType::Duration(ArrowTimeUnit::Microsecond) => {
- remove_primitive_array_indices!(arr, ArrowDurationMicrosecondType, indices)
- }
- ArrowType::Duration(ArrowTimeUnit::Nanosecond) => {
- remove_primitive_array_indices!(arr, ArrowDurationNanosecondType, indices)
- }
- ArrowType::Timestamp(ArrowTimeUnit::Second, _) => {
- remove_primitive_array_indices!(arr, ArrowTimestampSecondType, indices)
- }
- ArrowType::Timestamp(ArrowTimeUnit::Millisecond, _) => {
- remove_primitive_array_indices!(arr, ArrowTimestampMillisecondType, indices)
- }
- ArrowType::Timestamp(ArrowTimeUnit::Microsecond, _) => {
- remove_primitive_array_indices!(arr, ArrowTimestampMicrosecondType, indices)
- }
- ArrowType::Timestamp(ArrowTimeUnit::Nanosecond, _) => {
- remove_primitive_array_indices!(arr, ArrowTimestampNanosecondType, indices)
- }
- ArrowType::Utf8 => {
- remove_array_indices_custom_builder!(arr, StringArray, StringBuilder, indices)
- }
- ArrowType::Binary => {
- remove_array_indices_custom_builder!(arr, BinaryArray, BinaryBuilder, indices)
- }
- ArrowType::FixedSizeBinary(size) => remove_fixed_size_binary_array_indices!(
- arr,
- FixedSizeBinaryArray,
- FixedSizeBinaryBuilder,
- indices,
- size
- ),
- _ => Err(ParquetError::General(format!(
- "ListArray of type List({:?}) is not supported by array_reader",
- item_type
- ))),
- }
-}
-
-/// Implementation of ListArrayReader. Nested lists and lists of structs are not yet supported.
-impl<OffsetSize: OffsetSizeTrait> ArrayReader for ListArrayReader<OffsetSize> {
- fn as_any(&self) -> &dyn Any {
- self
- }
-
- /// Returns data type.
- /// This must be a List.
- fn get_data_type(&self) -> &ArrowType {
- &self.data_type
- }
-
- fn next_batch(&mut self, batch_size: usize) -> Result<ArrayRef> {
- let next_batch_array = self.item_reader.next_batch(batch_size)?;
- let item_type = self.item_reader.get_data_type().clone();
-
- if next_batch_array.len() == 0 {
- return Ok(new_empty_array(&self.data_type));
- }
- let def_levels = self
- .item_reader
- .get_def_levels()
- .ok_or_else(|| ArrowError("item_reader def levels are None.".to_string()))?;
- let rep_levels = self
- .item_reader
- .get_rep_levels()
- .ok_or_else(|| ArrowError("item_reader rep levels are None.".to_string()))?;
-
- if !((def_levels.len() == rep_levels.len())
- && (rep_levels.len() == next_batch_array.len()))
- {
- return Err(ArrowError(
- "Expected item_reader def_levels and rep_levels to be same length as batch".to_string(),
- ));
- }
-
- // List definitions can be encoded as 4 values:
- // - n + 0: the list slot is null
- // - n + 1: the list slot is not null, but is empty (i.e. [])
- // - n + 2: the list slot is not null, but its child is empty (i.e. [ null ])
- // - n + 3: the list slot is not null, and its child is not empty
- // Where n is the max definition level of the list's parent.
- // If a Parquet schema's only leaf is the list, then n = 0.
-
- // If the list index is at empty definition, the child slot is null
- let null_list_indices: Vec<usize> = def_levels
- .iter()
- .enumerate()
- .filter_map(|(index, def)| {
- if *def <= self.list_empty_def_level {
- Some(index)
- } else {
- None
- }
- })
- .collect();
- let batch_values = match null_list_indices.len() {
- 0 => next_batch_array.clone(),
- _ => remove_indices(next_batch_array.clone(), item_type, null_list_indices)?,
- };
-
- // first item in each list has rep_level = 0, subsequent items have rep_level = 1
- let mut offsets: Vec<OffsetSize> = Vec::new();
- let mut cur_offset = OffsetSize::zero();
- def_levels.iter().zip(rep_levels).for_each(|(d, r)| {
- if *r == 0 || d == &self.list_empty_def_level {
- offsets.push(cur_offset);
- }
- if d > &self.list_empty_def_level {
- cur_offset += OffsetSize::one();
- }
- });
- offsets.push(cur_offset);
-
- let num_bytes = bit_util::ceil(offsets.len(), 8);
- // TODO: A useful optimization is to use the null count to fill with
- // 0 or null, to reduce individual bits set in a loop.
- // To favour dense data, set every slot to true, then unset
- let mut null_buf = MutableBuffer::new(num_bytes).with_bitset(num_bytes, true);
- let null_slice = null_buf.as_slice_mut();
- let mut list_index = 0;
- for i in 0..rep_levels.len() {
- // If the level is lower than empty, then the slot is null.
- // When a list is non-nullable, its empty level = null level,
- // so this automatically factors that in.
- if rep_levels[i] == 0 && def_levels[i] < self.list_empty_def_level {
- bit_util::unset_bit(null_slice, list_index);
- }
- if rep_levels[i] == 0 {
- list_index += 1;
- }
- }
- let value_offsets = Buffer::from(&offsets.to_byte_slice());
-
- let list_data = ArrayData::builder(self.get_data_type().clone())
- .len(offsets.len() - 1)
- .add_buffer(value_offsets)
- .add_child_data(batch_values.data().clone())
- .null_bit_buffer(null_buf.into())
- .offset(next_batch_array.offset())
- .build();
-
- let result_array = GenericListArray::<OffsetSize>::from(list_data);
- Ok(Arc::new(result_array))
- }
-
- fn get_def_levels(&self) -> Option<&[i16]> {
- self.def_level_buffer
- .as_ref()
- .map(|buf| unsafe { buf.typed_data() })
- }
-
- fn get_rep_levels(&self) -> Option<&[i16]> {
- self.rep_level_buffer
- .as_ref()
- .map(|buf| unsafe { buf.typed_data() })
- }
-}
-
-/// Implementation of struct array reader.
-pub struct StructArrayReader {
- children: Vec<Box<dyn ArrayReader>>,
- data_type: ArrowType,
- struct_def_level: i16,
- struct_rep_level: i16,
- def_level_buffer: Option<Buffer>,
- rep_level_buffer: Option<Buffer>,
-}
-
-impl StructArrayReader {
- /// Construct struct array reader.
- pub fn new(
- data_type: ArrowType,
- children: Vec<Box<dyn ArrayReader>>,
- def_level: i16,
- rep_level: i16,
- ) -> Self {
- Self {
- data_type,
- children,
- struct_def_level: def_level,
- struct_rep_level: rep_level,
- def_level_buffer: None,
- rep_level_buffer: None,
- }
- }
-}
-
-impl ArrayReader for StructArrayReader {
- fn as_any(&self) -> &dyn Any {
- self
- }
-
- /// Returns data type.
- /// This must be a struct.
- fn get_data_type(&self) -> &ArrowType {
- &self.data_type
- }
-
- /// Read `batch_size` struct records.
- ///
- /// Definition levels of struct array is calculated as following:
- /// ```ignore
- /// def_levels[i] = min(child1_def_levels[i], child2_def_levels[i], ...,
- /// childn_def_levels[i]);
- /// ```
- ///
- /// Repetition levels of struct array is calculated as following:
- /// ```ignore
- /// rep_levels[i] = child1_rep_levels[i];
- /// ```
- ///
- /// The null bitmap of struct array is calculated from def_levels:
- /// ```ignore
- /// null_bitmap[i] = (def_levels[i] >= self.def_level);
- /// ```
- fn next_batch(&mut self, batch_size: usize) -> Result<ArrayRef> {
- if self.children.is_empty() {
- self.def_level_buffer = None;
- self.rep_level_buffer = None;
- return Ok(Arc::new(StructArray::from(Vec::new())));
- }
-
- let children_array = self
- .children
- .iter_mut()
- .map(|reader| reader.next_batch(batch_size))
- .try_fold(
- Vec::new(),
- |mut result, child_array| -> Result<Vec<ArrayRef>> {
- result.push(child_array?);
- Ok(result)
- },
- )?;
-
- // check that array child data has same size
- let children_array_len =
- children_array.first().map(|arr| arr.len()).ok_or_else(|| {
- general_err!("Struct array reader should have at least one child!")
- })?;
-
- let all_children_len_eq = children_array
- .iter()
- .all(|arr| arr.len() == children_array_len);
- if !all_children_len_eq {
- return Err(general_err!("Not all children array length are the same!"));
- }
-
- // calculate struct def level data
- let buffer_size = children_array_len * size_of::<i16>();
- let mut def_level_data_buffer = MutableBuffer::new(buffer_size);
- def_level_data_buffer.resize(buffer_size, 0);
-
- let def_level_data = def_level_data_buffer.typed_data_mut();
-
- def_level_data
- .iter_mut()
- .for_each(|v| *v = self.struct_def_level);
-
- for child in &self.children {
- if let Some(current_child_def_levels) = child.get_def_levels() {
- if current_child_def_levels.len() != children_array_len {
- return Err(general_err!("Child array length are not equal!"));
- } else {
- for i in 0..children_array_len {
- def_level_data[i] =
- min(def_level_data[i], current_child_def_levels[i]);
- }
- }
- }
- }
-
- // calculate bitmap for current array
- let mut bitmap_builder = BooleanBufferBuilder::new(children_array_len);
- for def_level in def_level_data {
- let not_null = *def_level >= self.struct_def_level;
- bitmap_builder.append(not_null);
- }
-
- // Now we can build array data
- let array_data = ArrayDataBuilder::new(self.data_type.clone())
- .len(children_array_len)
- .null_bit_buffer(bitmap_builder.finish())
- .child_data(
- children_array
- .iter()
- .map(|x| x.data().clone())
- .collect::<Vec<ArrayData>>(),
- )
- .build();
-
- // calculate struct rep level data, since struct doesn't add to repetition
- // levels, here we just need to keep repetition levels of first array
- // TODO: Verify that all children array reader has same repetition levels
- let rep_level_data = self
- .children
- .first()
- .ok_or_else(|| {
- general_err!("Struct array reader should have at least one child!")
- })?
- .get_rep_levels()
- .map(|data| -> Result<Buffer> {
- let mut buffer = Int16BufferBuilder::new(children_array_len);
- buffer.append_slice(data);
- Ok(buffer.finish())
- })
- .transpose()?;
-
- self.def_level_buffer = Some(def_level_data_buffer.into());
- self.rep_level_buffer = rep_level_data;
- Ok(Arc::new(StructArray::from(array_data)))
- }
-
- fn get_def_levels(&self) -> Option<&[i16]> {
- self.def_level_buffer
- .as_ref()
- .map(|buf| unsafe { buf.typed_data() })
- }
-
- fn get_rep_levels(&self) -> Option<&[i16]> {
- self.rep_level_buffer
- .as_ref()
- .map(|buf| unsafe { buf.typed_data() })
- }
-}
-
-/// Create array reader from parquet schema, column indices, and parquet file reader.
-pub fn build_array_reader<T>(
- parquet_schema: SchemaDescPtr,
- arrow_schema: Schema,
- column_indices: T,
- file_reader: Arc<dyn FileReader>,
-) -> Result<Box<dyn ArrayReader>>
-where
- T: IntoIterator<Item = usize>,
-{
- let mut leaves = HashMap::<*const Type, usize>::new();
-
- let mut filtered_root_names = HashSet::<String>::new();
-
- for c in column_indices {
- let column = parquet_schema.column(c).self_type() as *const Type;
-
- leaves.insert(column, c);
-
- let root = parquet_schema.get_column_root_ptr(c);
- filtered_root_names.insert(root.name().to_string());
- }
-
- if leaves.is_empty() {
- return Err(general_err!("Can't build array reader without columns!"));
- }
-
- // Only pass root fields that take part in the projection
- // to avoid traversal of columns that are not read.
- // TODO: also prune unread parts of the tree in child structures
- let filtered_root_fields = parquet_schema
- .root_schema()
- .get_fields()
- .iter()
- .filter(|field| filtered_root_names.contains(field.name()))
- .cloned()
- .collect::<Vec<_>>();
-
- let proj = Type::GroupType {
- basic_info: parquet_schema.root_schema().get_basic_info().clone(),
- fields: filtered_root_fields,
- };
-
- ArrayReaderBuilder::new(
- Arc::new(proj),
- Arc::new(arrow_schema),
- Arc::new(leaves),
- file_reader,
- )
- .build_array_reader()
-}
-
-/// Used to build array reader.
-struct ArrayReaderBuilder {
- root_schema: TypePtr,
- arrow_schema: Arc<Schema>,
- // Key: columns that need to be included in final array builder
- // Value: column index in schema
- columns_included: Arc<HashMap<*const Type, usize>>,
- file_reader: Arc<dyn FileReader>,
-}
-
-/// Used in type visitor.
-#[derive(Clone)]
-struct ArrayReaderBuilderContext {
- def_level: i16,
- rep_level: i16,
- path: ColumnPath,
-}
-
-impl Default for ArrayReaderBuilderContext {
- fn default() -> Self {
- Self {
- def_level: 0i16,
- rep_level: 0i16,
- path: ColumnPath::new(Vec::new()),
- }
- }
-}
-
-/// Create array reader by visiting schema.
-impl<'a> TypeVisitor<Option<Box<dyn ArrayReader>>, &'a ArrayReaderBuilderContext>
- for ArrayReaderBuilder
-{
- /// Build array reader for primitive type.
- /// Currently we don't have a list reader implementation, so repeated type is not
- /// supported yet.
- fn visit_primitive(
- &mut self,
- cur_type: TypePtr,
- context: &'a ArrayReaderBuilderContext,
- ) -> Result<Option<Box<dyn ArrayReader>>> {
- if self.is_included(cur_type.as_ref()) {
- let mut new_context = context.clone();
- new_context.path.append(vec![cur_type.name().to_string()]);
-
- match cur_type.get_basic_info().repetition() {
- Repetition::REPEATED => {
- new_context.def_level += 1;
- new_context.rep_level += 1;
- }
- Repetition::OPTIONAL => {
- new_context.def_level += 1;
- }
- _ => (),
- }
-
- let reader =
- self.build_for_primitive_type_inner(cur_type.clone(), &new_context)?;
-
- if cur_type.get_basic_info().repetition() == Repetition::REPEATED {
- Err(ArrowError(
- "Reading repeated field is not supported yet!".to_string(),
- ))
- } else {
- Ok(Some(reader))
- }
- } else {
- Ok(None)
- }
- }
-
- /// Build array reader for struct type.
- fn visit_struct(
- &mut self,
- cur_type: Arc<Type>,
- context: &'a ArrayReaderBuilderContext,
- ) -> Result<Option<Box<dyn ArrayReader>>> {
- let mut new_context = context.clone();
- new_context.path.append(vec![cur_type.name().to_string()]);
-
- if cur_type.get_basic_info().has_repetition() {
- match cur_type.get_basic_info().repetition() {
- Repetition::REPEATED => {
- new_context.def_level += 1;
- new_context.rep_level += 1;
- }
- Repetition::OPTIONAL => {
- new_context.def_level += 1;
- }
- _ => (),
- }
- }
-
- if let Some(reader) = self.build_for_struct_type_inner(&cur_type, &new_context)? {
- if cur_type.get_basic_info().has_repetition()
- && cur_type.get_basic_info().repetition() == Repetition::REPEATED
- {
- Err(ArrowError(
- "Reading repeated field is not supported yet!".to_string(),
- ))
- } else {
- Ok(Some(reader))
- }
- } else {
- Ok(None)
- }
- }
-
- /// Build array reader for map type.
- /// Currently this is not supported.
- fn visit_map(
- &mut self,
- _cur_type: Arc<Type>,
- _context: &'a ArrayReaderBuilderContext,
- ) -> Result<Option<Box<dyn ArrayReader>>> {
- Err(ArrowError(
- "Reading parquet map array into arrow is not supported yet!".to_string(),
- ))
- }
-
- /// Build array reader for list type.
- fn visit_list_with_item(
- &mut self,
- list_type: Arc<Type>,
- item_type: Arc<Type>,
- context: &'a ArrayReaderBuilderContext,
- ) -> Result<Option<Box<dyn ArrayReader>>> {
- let list_child = &list_type
- .get_fields()
- .first()
- .ok_or_else(|| ArrowError("List field must have a child.".to_string()))?;
- let mut new_context = context.clone();
-
- new_context.path.append(vec![list_type.name().to_string()]);
- // We need to know at what definition a list or its child is null
- let list_null_def = new_context.def_level;
- let mut list_empty_def = new_context.def_level;
-
- // If the list's root is nullable
- if let Repetition::OPTIONAL = list_type.get_basic_info().repetition() {
- new_context.def_level += 1;
- // current level is nullable, increment to get level for empty list slot
- list_empty_def += 1;
- }
-
- match list_child.get_basic_info().repetition() {
- Repetition::REPEATED => {
- new_context.def_level += 1;
- new_context.rep_level += 1;
- }
- Repetition::OPTIONAL => {
- new_context.def_level += 1;
- }
- _ => (),
- }
-
- let item_reader = self
- .dispatch(item_type.clone(), &new_context)
- .unwrap()
- .unwrap();
-
- let item_reader_type = item_reader.get_data_type().clone();
-
- match item_reader_type {
- ArrowType::List(_)
- | ArrowType::FixedSizeList(_, _)
- | ArrowType::Struct(_)
- | ArrowType::Dictionary(_, _) => Err(ArrowError(format!(
- "reading List({:?}) into arrow not supported yet",
- item_type
- ))),
- _ => {
- // a list is a group type with a single child. The list child's
- // name comes from the child's field name.
- let mut list_child = list_type.get_fields().first().ok_or(ArrowError(
- "List GroupType should have a field".to_string(),
- ))?;
- // if the child's name is "list" and it has a child, then use this child
- if list_child.name() == "list" && !list_child.get_fields().is_empty() {
- list_child = list_child.get_fields().first().unwrap();
- }
- let arrow_type = self
- .arrow_schema
- .field_with_name(list_type.name())
- .ok()
- .map(|f| f.data_type().to_owned())
- .unwrap_or_else(|| {
- ArrowType::List(Box::new(Field::new(
- list_child.name(),
- item_reader_type.clone(),
- list_child.is_optional(),
- )))
- });
-
- let list_array_reader: Box<dyn ArrayReader> = match arrow_type {
- ArrowType::List(_) => Box::new(ListArrayReader::<i32>::new(
- item_reader,
- arrow_type,
- item_reader_type,
- new_context.def_level,
- new_context.rep_level,
- list_null_def,
- list_empty_def,
- )),
- ArrowType::LargeList(_) => Box::new(ListArrayReader::<i64>::new(
- item_reader,
- arrow_type,
- item_reader_type,
- new_context.def_level,
- new_context.rep_level,
- list_null_def,
- list_empty_def,
- )),
-
- _ => {
- return Err(ArrowError(format!(
- "creating ListArrayReader with type {:?} should be unreachable",
- arrow_type
- )))
- }
- };
-
- Ok(Some(list_array_reader))
- }
- }
- }
-}
-
-impl<'a> ArrayReaderBuilder {
- /// Construct array reader builder.
- fn new(
- root_schema: TypePtr,
- arrow_schema: Arc<Schema>,
- columns_included: Arc<HashMap<*const Type, usize>>,
- file_reader: Arc<dyn FileReader>,
- ) -> Self {
- Self {
- root_schema,
- arrow_schema,
- columns_included,
- file_reader,
- }
- }
-
- /// Main entry point.
- fn build_array_reader(&mut self) -> Result<Box<dyn ArrayReader>> {
- let context = ArrayReaderBuilderContext::default();
-
- self.visit_struct(self.root_schema.clone(), &context)
- .and_then(|reader_opt| {
- reader_opt.ok_or_else(|| general_err!("Failed to build array reader!"))
- })
- }
-
- // Utility functions
-
- /// Check whether one column in included in this array reader builder.
- fn is_included(&self, t: &Type) -> bool {
- self.columns_included.contains_key(&(t as *const Type))
- }
-
- /// Creates primitive array reader for each primitive type.
- fn build_for_primitive_type_inner(
- &self,
- cur_type: TypePtr,
- context: &'a ArrayReaderBuilderContext,
- ) -> Result<Box<dyn ArrayReader>> {
- let column_desc = Arc::new(ColumnDescriptor::new(
- cur_type.clone(),
- context.def_level,
- context.rep_level,
- context.path.clone(),
- ));
- let page_iterator = Box::new(FilePageIterator::new(
- self.columns_included[&(cur_type.as_ref() as *const Type)],
- self.file_reader.clone(),
- )?);
-
- let arrow_type: Option<ArrowType> = self
- .get_arrow_field(&cur_type, context)
- .map(|f| f.data_type().clone());
-
- match cur_type.get_physical_type() {
- PhysicalType::BOOLEAN => Ok(Box::new(PrimitiveArrayReader::<BoolType>::new(
- page_iterator,
- column_desc,
- arrow_type,
- )?)),
- PhysicalType::INT32 => {
- if let Some(ArrowType::Null) = arrow_type {
- Ok(Box::new(NullArrayReader::<Int32Type>::new(
- page_iterator,
- column_desc,
- )?))
- } else {
- Ok(Box::new(PrimitiveArrayReader::<Int32Type>::new(
- page_iterator,
- column_desc,
- arrow_type,
- )?))
- }
- }
- PhysicalType::INT64 => Ok(Box::new(PrimitiveArrayReader::<Int64Type>::new(
- page_iterator,
- column_desc,
- arrow_type,
- )?)),
- PhysicalType::INT96 => {
- // get the optional timezone information from arrow type
- let timezone = arrow_type
- .as_ref()
- .map(|data_type| {
- if let ArrowType::Timestamp(_, tz) = data_type {
- tz.clone()
- } else {
- None
- }
- })
- .flatten();
- let converter = Int96Converter::new(Int96ArrayConverter { timezone });
- Ok(Box::new(ComplexObjectArrayReader::<
- Int96Type,
- Int96Converter,
- >::new(
- page_iterator,
- column_desc,
- converter,
- arrow_type,
- )?))
- }
- PhysicalType::FLOAT => Ok(Box::new(PrimitiveArrayReader::<FloatType>::new(
- page_iterator,
- column_desc,
- arrow_type,
- )?)),
- PhysicalType::DOUBLE => {
- Ok(Box::new(PrimitiveArrayReader::<DoubleType>::new(
- page_iterator,
- column_desc,
- arrow_type,
- )?))
- }
- PhysicalType::BYTE_ARRAY => {
- if cur_type.get_basic_info().converted_type() == ConvertedType::UTF8 {
- if let Some(ArrowType::LargeUtf8) = arrow_type {
- let converter =
- LargeUtf8Converter::new(LargeUtf8ArrayConverter {});
- Ok(Box::new(ComplexObjectArrayReader::<
- ByteArrayType,
- LargeUtf8Converter,
- >::new(
- page_iterator,
- column_desc,
- converter,
- arrow_type,
- )?))
- } else {
- use crate::arrow::arrow_array_reader::{
- ArrowArrayReader, StringArrayConverter,
- };
- let converter = StringArrayConverter::new();
- Ok(Box::new(ArrowArrayReader::try_new(
- *page_iterator,
- column_desc,
- converter,
- arrow_type,
- )?))
- }
- } else if let Some(ArrowType::LargeBinary) = arrow_type {
- let converter =
- LargeBinaryConverter::new(LargeBinaryArrayConverter {});
- Ok(Box::new(ComplexObjectArrayReader::<
- ByteArrayType,
- LargeBinaryConverter,
- >::new(
- page_iterator,
- column_desc,
- converter,
- arrow_type,
- )?))
- } else {
- let converter = BinaryConverter::new(BinaryArrayConverter {});
- Ok(Box::new(ComplexObjectArrayReader::<
- ByteArrayType,
- BinaryConverter,
- >::new(
- page_iterator,
- column_desc,
- converter,
- arrow_type,
- )?))
- }
- }
- PhysicalType::FIXED_LEN_BYTE_ARRAY
- if cur_type.get_basic_info().converted_type()
- == ConvertedType::DECIMAL =>
- {
- let converter = DecimalConverter::new(DecimalArrayConverter::new(
- cur_type.get_precision(),
- cur_type.get_scale(),
- ));
- Ok(Box::new(ComplexObjectArrayReader::<
- FixedLenByteArrayType,
- DecimalConverter,
- >::new(
- page_iterator,
- column_desc,
- converter,
- arrow_type,
- )?))
- }
- PhysicalType::FIXED_LEN_BYTE_ARRAY => {
- if cur_type.get_basic_info().converted_type() == ConvertedType::INTERVAL {
- let byte_width = match *cur_type {
- Type::PrimitiveType {
- ref type_length, ..
- } => *type_length,
- _ => {
- return Err(ArrowError(
- "Expected a physical type, not a group type".to_string(),
- ))
- }
- };
- if byte_width != 12 {
- return Err(ArrowError(format!(
- "Parquet interval type should have length of 12, found {}",
- byte_width
- )));
- }
- match arrow_type {
- Some(ArrowType::Interval(IntervalUnit::DayTime)) => {
- let converter = IntervalDayTimeConverter::new(
- IntervalDayTimeArrayConverter {},
- );
- Ok(Box::new(ComplexObjectArrayReader::<
- FixedLenByteArrayType,
- IntervalDayTimeConverter,
- >::new(
- page_iterator,
- column_desc,
- converter,
- arrow_type,
- )?))
- }
- Some(ArrowType::Interval(IntervalUnit::YearMonth)) => {
- let converter = IntervalYearMonthConverter::new(
- IntervalYearMonthArrayConverter {},
- );
- Ok(Box::new(ComplexObjectArrayReader::<
- FixedLenByteArrayType,
- IntervalYearMonthConverter,
- >::new(
- page_iterator,
- column_desc,
- converter,
- arrow_type,
- )?))
- }
- Some(t) => Err(ArrowError(format!(
- "Cannot write a Parquet interval to {:?}",
- t
- ))),
- None => {
- // we do not support an interval not matched to an Arrow type,
- // because we risk data loss as we won't know which of the 12 bytes
- // are or should be populated
- Err(ArrowError(
- "Cannot write a Parquet interval with no Arrow type specified.
- There is a risk of data loss as Arrow either supports YearMonth or
- DayTime precision. Without the Arrow type, we cannot infer the type.
- ".to_string()
- ))
- }
- }
- } else {
- let byte_width = match *cur_type {
- Type::PrimitiveType {
- ref type_length, ..
- } => *type_length,
- _ => {
- return Err(ArrowError(
- "Expected a physical type, not a group type".to_string(),
- ))
- }
- };
- let converter = FixedLenBinaryConverter::new(
- FixedSizeArrayConverter::new(byte_width),
- );
- Ok(Box::new(ComplexObjectArrayReader::<
- FixedLenByteArrayType,
- FixedLenBinaryConverter,
- >::new(
- page_iterator,
- column_desc,
- converter,
- arrow_type,
- )?))
- }
- }
- }
- }
-
- /// Constructs struct array reader without considering repetition.
- fn build_for_struct_type_inner(
- &mut self,
- cur_type: &Type,
- context: &'a ArrayReaderBuilderContext,
- ) -> Result<Option<Box<dyn ArrayReader>>> {
- let mut fields = Vec::with_capacity(cur_type.get_fields().len());
- let mut children_reader = Vec::with_capacity(cur_type.get_fields().len());
-
- for child in cur_type.get_fields() {
- let mut struct_context = context.clone();
- if let Some(child_reader) = self.dispatch(child.clone(), context)? {
- // TODO: this results in calling get_arrow_field twice, it could be reused
- // from child_reader above, by making child_reader carry its `Field`
- struct_context.path.append(vec![child.name().to_string()]);
- let field = match self.get_arrow_field(child, &struct_context) {
- Some(f) => f.clone(),
- _ => Field::new(
- child.name(),
- child_reader.get_data_type().clone(),
- child.is_optional(),
- ),
- };
- fields.push(field);
- children_reader.push(child_reader);
- }
- }
-
- if !fields.is_empty() {
- let arrow_type = ArrowType::Struct(fields);
- Ok(Some(Box::new(StructArrayReader::new(
- arrow_type,
- children_reader,
- context.def_level,
- context.rep_level,
- ))))
- } else {
- Ok(None)
- }
- }
-
- fn get_arrow_field(
- &self,
- cur_type: &Type,
- context: &'a ArrayReaderBuilderContext,
- ) -> Option<&Field> {
- let parts: Vec<&str> = context
- .path
- .parts()
- .iter()
- .map(|x| -> &str { x })
- .collect::<Vec<&str>>();
-
- // If the parts length is one it'll have the top level "schema" type. If
- // it's two then it'll be a top-level type that we can get from the arrow
- // schema directly.
- if parts.len() <= 2 {
- self.arrow_schema.field_with_name(cur_type.name()).ok()
- } else {
- // If it's greater than two then we need to traverse the type path
- // until we find the actual field we're looking for.
- let mut field: Option<&Field> = None;
-
- for (i, part) in parts.iter().enumerate().skip(1) {
- if i == 1 {
- field = self.arrow_schema.field_with_name(part).ok();
- } else if let Some(f) = field {
- if let ArrowType::Struct(fields) = f.data_type() {
- field = fields.iter().find(|f| f.name() == part)
- } else {
- field = None
- }
- } else {
- field = None
- }
- }
- field
- }
- }
-}
-
-#[cfg(test)]
-mod tests {
- use super::*;
- use crate::arrow::converter::{Utf8ArrayConverter, Utf8Converter};
- use crate::arrow::schema::parquet_to_arrow_schema;
- use crate::basic::{Encoding, Type as PhysicalType};
- use crate::column::page::{Page, PageReader};
- use crate::data_type::{ByteArray, DataType, Int32Type, Int64Type};
- use crate::errors::Result;
- use crate::file::reader::{FileReader, SerializedFileReader};
- use crate::schema::parser::parse_message_type;
- use crate::schema::types::{ColumnDescPtr, SchemaDescriptor};
- use crate::util::test_common::page_util::{
- DataPageBuilder, DataPageBuilderImpl, InMemoryPageIterator,
- };
- use crate::util::test_common::{get_test_file, make_pages};
- use arrow::array::{
- Array, ArrayRef, LargeListArray, ListArray, PrimitiveArray, StringArray,
- StructArray,
- };
- use arrow::datatypes::{
- ArrowPrimitiveType, DataType as ArrowType, Date32Type as ArrowDate32, Field,
- Int32Type as ArrowInt32, Int64Type as ArrowInt64,
- Time32MillisecondType as ArrowTime32MillisecondArray,
- Time64MicrosecondType as ArrowTime64MicrosecondArray,
- TimestampMicrosecondType as ArrowTimestampMicrosecondType,
- TimestampMillisecondType as ArrowTimestampMillisecondType,
- };
- use rand::distributions::uniform::SampleUniform;
- use rand::{thread_rng, Rng};
- use std::any::Any;
- use std::collections::VecDeque;
- use std::sync::Arc;
-
- fn make_column_chunks<T: DataType>(
- column_desc: ColumnDescPtr,
- encoding: Encoding,
- num_levels: usize,
- min_value: T::T,
- max_value: T::T,
- def_levels: &mut Vec<i16>,
- rep_levels: &mut Vec<i16>,
- values: &mut Vec<T::T>,
- page_lists: &mut Vec<Vec<Page>>,
- use_v2: bool,
- num_chunks: usize,
- ) where
- T::T: PartialOrd + SampleUniform + Copy,
- {
- for _i in 0..num_chunks {
- let mut pages = VecDeque::new();
- let mut data = Vec::new();
- let mut page_def_levels = Vec::new();
- let mut page_rep_levels = Vec::new();
-
- make_pages::<T>(
- column_desc.clone(),
- encoding,
- 1,
- num_levels,
- min_value,
- max_value,
- &mut page_def_levels,
- &mut page_rep_levels,
- &mut data,
- &mut pages,
- use_v2,
- );
-
- def_levels.append(&mut page_def_levels);
- rep_levels.append(&mut page_rep_levels);
- values.append(&mut data);
- page_lists.push(Vec::from(pages));
- }
- }
-
- #[test]
- fn test_primitive_array_reader_empty_pages() {
- // Construct column schema
- let message_type = "
- message test_schema {
- REQUIRED INT32 leaf;
- }
- ";
-
- let schema = parse_message_type(message_type)
- .map(|t| Arc::new(SchemaDescriptor::new(Arc::new(t))))
- .unwrap();
-
- let column_desc = schema.column(0);
- let page_iterator = EmptyPageIterator::new(schema);
-
- let mut array_reader = PrimitiveArrayReader::<Int32Type>::new(
- Box::new(page_iterator),
- column_desc,
- None,
- )
- .unwrap();
-
- // expect no values to be read
- let array = array_reader.next_batch(50).unwrap();
- assert!(array.is_empty());
- }
-
- #[test]
- fn test_primitive_array_reader_data() {
- // Construct column schema
- let message_type = "
- message test_schema {
- REQUIRED INT32 leaf;
- }
- ";
-
- let schema = parse_message_type(message_type)
- .map(|t| Arc::new(SchemaDescriptor::new(Arc::new(t))))
- .unwrap();
-
- let column_desc = schema.column(0);
-
- // Construct page iterator
- {
- let mut data = Vec::new();
- let mut page_lists = Vec::new();
- make_column_chunks::<Int32Type>(
- column_desc.clone(),
- Encoding::PLAIN,
- 100,
- 1,
- 200,
- &mut Vec::new(),
- &mut Vec::new(),
- &mut data,
- &mut page_lists,
- true,
- 2,
- );
- let page_iterator =
- InMemoryPageIterator::new(schema, column_desc.clone(), page_lists);
-
- let mut array_reader = PrimitiveArrayReader::<Int32Type>::new(
- Box::new(page_iterator),
- column_desc,
- None,
- )
- .unwrap();
-
- // Read first 50 values, which are all from the first column chunk
- let array = array_reader.next_batch(50).unwrap();
- let array = array
- .as_any()
- .downcast_ref::<PrimitiveArray<ArrowInt32>>()
- .unwrap();
-
- assert_eq!(
- &PrimitiveArray::<ArrowInt32>::from(data[0..50].to_vec()),
- array
- );
-
- // Read next 100 values, the first 50 ones are from the first column chunk,
- // and the last 50 ones are from the second column chunk
- let array = array_reader.next_batch(100).unwrap();
- let array = array
- .as_any()
- .downcast_ref::<PrimitiveArray<ArrowInt32>>()
- .unwrap();
-
- assert_eq!(
- &PrimitiveArray::<ArrowInt32>::from(data[50..150].to_vec()),
- array
- );
-
- // Try to read 100 values, however there are only 50 values
- let array = array_reader.next_batch(100).unwrap();
- let array = array
- .as_any()
- .downcast_ref::<PrimitiveArray<ArrowInt32>>()
- .unwrap();
-
- assert_eq!(
- &PrimitiveArray::<ArrowInt32>::from(data[150..200].to_vec()),
- array
- );
- }
- }
-
- macro_rules! test_primitive_array_reader_one_type {
- ($arrow_parquet_type:ty, $physical_type:expr, $converted_type_str:expr, $result_arrow_type:ty, $result_arrow_cast_type:ty, $result_primitive_type:ty) => {{
- let message_type = format!(
- "
- message test_schema {{
- REQUIRED {:?} leaf ({});
- }}
- ",
- $physical_type, $converted_type_str
- );
- let schema = parse_message_type(&message_type)
- .map(|t| Arc::new(SchemaDescriptor::new(Arc::new(t))))
- .unwrap();
-
- let column_desc = schema.column(0);
-
- // Construct page iterator
- {
- let mut data = Vec::new();
- let mut page_lists = Vec::new();
- make_column_chunks::<$arrow_parquet_type>(
- column_desc.clone(),
- Encoding::PLAIN,
- 100,
- 1,
- 200,
- &mut Vec::new(),
- &mut Vec::new(),
- &mut data,
- &mut page_lists,
- true,
- 2,
- );
- let page_iterator = InMemoryPageIterator::new(
- schema.clone(),
- column_desc.clone(),
- page_lists,
- );
- let mut array_reader = PrimitiveArrayReader::<$arrow_parquet_type>::new(
- Box::new(page_iterator),
- column_desc.clone(),
- None,
- )
- .expect("Unable to get array reader");
-
- let array = array_reader
- .next_batch(50)
- .expect("Unable to get batch from reader");
-
- let result_data_type = <$result_arrow_type>::DATA_TYPE;
- let array = array
- .as_any()
- .downcast_ref::<PrimitiveArray<$result_arrow_type>>()
- .expect(
- format!(
- "Unable to downcast {:?} to {:?}",
- array.data_type(),
- result_data_type
- )
- .as_str(),
- );
-
- // create expected array as primitive, and cast to result type
- let expected = PrimitiveArray::<$result_arrow_cast_type>::from(
- data[0..50]
- .iter()
- .map(|x| *x as $result_primitive_type)
- .collect::<Vec<$result_primitive_type>>(),
- );
- let expected = Arc::new(expected) as ArrayRef;
- let expected = arrow::compute::cast(&expected, &result_data_type)
- .expect("Unable to cast expected array");
- assert_eq!(expected.data_type(), &result_data_type);
- let expected = expected
- .as_any()
- .downcast_ref::<PrimitiveArray<$result_arrow_type>>()
- .expect(
- format!(
- "Unable to downcast expected {:?} to {:?}",
- expected.data_type(),
- result_data_type
- )
- .as_str(),
- );
- assert_eq!(expected, array);
- }
- }};
- }
-
- #[test]
- fn test_primitive_array_reader_temporal_types() {
- test_primitive_array_reader_one_type!(
- Int32Type,
- PhysicalType::INT32,
- "DATE",
- ArrowDate32,
- ArrowInt32,
- i32
- );
- test_primitive_array_reader_one_type!(
- Int32Type,
- PhysicalType::INT32,
- "TIME_MILLIS",
- ArrowTime32MillisecondArray,
- ArrowInt32,
- i32
- );
- test_primitive_array_reader_one_type!(
- Int64Type,
- PhysicalType::INT64,
- "TIME_MICROS",
- ArrowTime64MicrosecondArray,
- ArrowInt64,
- i64
- );
- test_primitive_array_reader_one_type!(
- Int64Type,
- PhysicalType::INT64,
- "TIMESTAMP_MILLIS",
- ArrowTimestampMillisecondType,
- ArrowInt64,
- i64
- );
- test_primitive_array_reader_one_type!(
- Int64Type,
- PhysicalType::INT64,
- "TIMESTAMP_MICROS",
- ArrowTimestampMicrosecondType,
- ArrowInt64,
- i64
- );
- }
-
- #[test]
- fn test_primitive_array_reader_def_and_rep_levels() {
- // Construct column schema
- let message_type = "
- message test_schema {
- REPEATED Group test_mid {
- OPTIONAL INT32 leaf;
- }
- }
- ";
-
- let schema = parse_message_type(message_type)
- .map(|t| Arc::new(SchemaDescriptor::new(Arc::new(t))))
- .unwrap();
-
- let column_desc = schema.column(0);
-
- // Construct page iterator
- {
- let mut def_levels = Vec::new();
- let mut rep_levels = Vec::new();
- let mut page_lists = Vec::new();
- make_column_chunks::<Int32Type>(
- column_desc.clone(),
- Encoding::PLAIN,
- 100,
- 1,
- 200,
- &mut def_levels,
- &mut rep_levels,
- &mut Vec::new(),
- &mut page_lists,
- true,
- 2,
- );
-
- let page_iterator =
- InMemoryPageIterator::new(schema, column_desc.clone(), page_lists);
-
- let mut array_reader = PrimitiveArrayReader::<Int32Type>::new(
- Box::new(page_iterator),
- column_desc,
- None,
- )
- .unwrap();
-
- let mut accu_len: usize = 0;
-
- // Read first 50 values, which are all from the first column chunk
- let array = array_reader.next_batch(50).unwrap();
- assert_eq!(
- Some(&def_levels[accu_len..(accu_len + array.len())]),
- array_reader.get_def_levels()
- );
- assert_eq!(
- Some(&rep_levels[accu_len..(accu_len + array.len())]),
- array_reader.get_rep_levels()
- );
- accu_len += array.len();
-
- // Read next 100 values, the first 50 ones are from the first column chunk,
- // and the last 50 ones are from the second column chunk
- let array = array_reader.next_batch(100).unwrap();
- assert_eq!(
- Some(&def_levels[accu_len..(accu_len + array.len())]),
- array_reader.get_def_levels()
- );
- assert_eq!(
- Some(&rep_levels[accu_len..(accu_len + array.len())]),
- array_reader.get_rep_levels()
- );
- accu_len += array.len();
-
- // Try to read 100 values, however there are only 50 values
- let array = array_reader.next_batch(100).unwrap();
- assert_eq!(
- Some(&def_levels[accu_len..(accu_len + array.len())]),
- array_reader.get_def_levels()
- );
- assert_eq!(
- Some(&rep_levels[accu_len..(accu_len + array.len())]),
- array_reader.get_rep_levels()
- );
- }
- }
-
- #[test]
- fn test_complex_array_reader_no_pages() {
- let message_type = "
- message test_schema {
- REPEATED Group test_mid {
- OPTIONAL BYTE_ARRAY leaf (UTF8);
- }
- }
- ";
- let schema = parse_message_type(message_type)
- .map(|t| Arc::new(SchemaDescriptor::new(Arc::new(t))))
- .unwrap();
- let column_desc = schema.column(0);
- let pages: Vec<Vec<Page>> = Vec::new();
- let page_iterator = InMemoryPageIterator::new(schema, column_desc.clone(), pages);
-
- let converter = Utf8Converter::new(Utf8ArrayConverter {});
- let mut array_reader =
- ComplexObjectArrayReader::<ByteArrayType, Utf8Converter>::new(
- Box::new(page_iterator),
- column_desc,
- converter,
- None,
- )
- .unwrap();
-
- let values_per_page = 100; // this value is arbitrary in this test - the result should always be an array of 0 length
- let array = array_reader.next_batch(values_per_page).unwrap();
- assert_eq!(array.len(), 0);
- }
-
- #[test]
- fn test_complex_array_reader_def_and_rep_levels() {
- // Construct column schema
- let message_type = "
- message test_schema {
- REPEATED Group test_mid {
- OPTIONAL BYTE_ARRAY leaf (UTF8);
- }
- }
- ";
- let num_pages = 2;
- let values_per_page = 100;
- let str_base = "Hello World";
-
- let schema = parse_message_type(message_type)
- .map(|t| Arc::new(SchemaDescriptor::new(Arc::new(t))))
- .unwrap();
-
- let max_def_level = schema.column(0).max_def_level();
- let max_rep_level = schema.column(0).max_rep_level();
-
- assert_eq!(max_def_level, 2);
- assert_eq!(max_rep_level, 1);
-
- let mut rng = thread_rng();
- let column_desc = schema.column(0);
- let mut pages: Vec<Vec<Page>> = Vec::new();
-
- let mut rep_levels = Vec::with_capacity(num_pages * values_per_page);
- let mut def_levels = Vec::with_capacity(num_pages * values_per_page);
- let mut all_values = Vec::with_capacity(num_pages * values_per_page);
-
- for i in 0..num_pages {
- let mut values = Vec::with_capacity(values_per_page);
-
- for _ in 0..values_per_page {
- let def_level = rng.gen_range(0..max_def_level + 1);
- let rep_level = rng.gen_range(0..max_rep_level + 1);
- if def_level == max_def_level {
- let len = rng.gen_range(1..str_base.len());
- let slice = &str_base[..len];
- values.push(ByteArray::from(slice));
- all_values.push(Some(slice.to_string()));
- } else {
- all_values.push(None)
- }
- rep_levels.push(rep_level);
- def_levels.push(def_level)
- }
-
- let range = i * values_per_page..(i + 1) * values_per_page;
- let mut pb =
- DataPageBuilderImpl::new(column_desc.clone(), values.len() as u32, true);
-
- pb.add_rep_levels(max_rep_level, &rep_levels.as_slice()[range.clone()]);
- pb.add_def_levels(max_def_level, &def_levels.as_slice()[range]);
- pb.add_values::<ByteArrayType>(Encoding::PLAIN, values.as_slice());
-
- let data_page = pb.consume();
- pages.push(vec![data_page]);
- }
-
- let page_iterator = InMemoryPageIterator::new(schema, column_desc.clone(), pages);
-
- let converter = Utf8Converter::new(Utf8ArrayConverter {});
- let mut array_reader =
- ComplexObjectArrayReader::<ByteArrayType, Utf8Converter>::new(
- Box::new(page_iterator),
- column_desc,
- converter,
- None,
- )
- .unwrap();
-
- let mut accu_len: usize = 0;
-
- let array = array_reader.next_batch(values_per_page / 2).unwrap();
- assert_eq!(array.len(), values_per_page / 2);
- assert_eq!(
- Some(&def_levels[accu_len..(accu_len + array.len())]),
- array_reader.get_def_levels()
- );
- assert_eq!(
- Some(&rep_levels[accu_len..(accu_len + array.len())]),
- array_reader.get_rep_levels()
- );
- accu_len += array.len();
-
- // Read next values_per_page values, the first values_per_page/2 ones are from the first column chunk,
- // and the last values_per_page/2 ones are from the second column chunk
- let array = array_reader.next_batch(values_per_page).unwrap();
- assert_eq!(array.len(), values_per_page);
- assert_eq!(
- Some(&def_levels[accu_len..(accu_len + array.len())]),
- array_reader.get_def_levels()
- );
- assert_eq!(
- Some(&rep_levels[accu_len..(accu_len + array.len())]),
- array_reader.get_rep_levels()
- );
- let strings = array.as_any().downcast_ref::<StringArray>().unwrap();
- for i in 0..array.len() {
- if array.is_valid(i) {
- assert_eq!(
- all_values[i + accu_len].as_ref().unwrap().as_str(),
- strings.value(i)
- )
- } else {
- assert_eq!(all_values[i + accu_len], None)
- }
- }
- accu_len += array.len();
-
- // Try to read values_per_page values, however there are only values_per_page/2 values
- let array = array_reader.next_batch(values_per_page).unwrap();
- assert_eq!(array.len(), values_per_page / 2);
- assert_eq!(
- Some(&def_levels[accu_len..(accu_len + array.len())]),
- array_reader.get_def_levels()
- );
- assert_eq!(
- Some(&rep_levels[accu_len..(accu_len + array.len())]),
- array_reader.get_rep_levels()
- );
- }
-
- /// Array reader for test.
- struct InMemoryArrayReader {
- data_type: ArrowType,
- array: ArrayRef,
- def_levels: Option<Vec<i16>>,
- rep_levels: Option<Vec<i16>>,
- }
-
- impl InMemoryArrayReader {
- pub fn new(
- data_type: ArrowType,
- array: ArrayRef,
- def_levels: Option<Vec<i16>>,
- rep_levels: Option<Vec<i16>>,
- ) -> Self {
- Self {
- data_type,
- array,
- def_levels,
- rep_levels,
- }
- }
- }
-
- impl ArrayReader for InMemoryArrayReader {
- fn as_any(&self) -> &dyn Any {
- self
- }
-
- fn get_data_type(&self) -> &ArrowType {
- &self.data_type
- }
-
- fn next_batch(&mut self, _batch_size: usize) -> Result<ArrayRef> {
- Ok(self.array.clone())
- }
-
- fn get_def_levels(&self) -> Option<&[i16]> {
- self.def_levels.as_deref()
- }
-
- fn get_rep_levels(&self) -> Option<&[i16]> {
- self.rep_levels.as_deref()
- }
- }
-
- /// Iterator for testing reading empty columns
- struct EmptyPageIterator {
- schema: SchemaDescPtr,
- }
-
- impl EmptyPageIterator {
- fn new(schema: SchemaDescPtr) -> Self {
- EmptyPageIterator { schema }
- }
- }
-
- impl Iterator for EmptyPageIterator {
- type Item = Result<Box<dyn PageReader>>;
-
- fn next(&mut self) -> Option<Self::Item> {
- None
- }
- }
-
- impl PageIterator for EmptyPageIterator {
- fn schema(&mut self) -> Result<SchemaDescPtr> {
- Ok(self.schema.clone())
- }
-
- fn column_schema(&mut self) -> Result<ColumnDescPtr> {
- Ok(self.schema.column(0))
- }
- }
-
- #[test]
- fn test_struct_array_reader() {
- let array_1 = Arc::new(PrimitiveArray::<ArrowInt32>::from(vec![1, 2, 3, 4, 5]));
- let array_reader_1 = InMemoryArrayReader::new(
- ArrowType::Int32,
- array_1.clone(),
- Some(vec![0, 1, 2, 3, 1]),
- Some(vec![1, 1, 1, 1, 1]),
- );
-
- let array_2 = Arc::new(PrimitiveArray::<ArrowInt32>::from(vec![5, 4, 3, 2, 1]));
- let array_reader_2 = InMemoryArrayReader::new(
- ArrowType::Int32,
- array_2.clone(),
- Some(vec![0, 1, 3, 1, 2]),
- Some(vec![1, 1, 1, 1, 1]),
- );
-
- let struct_type = ArrowType::Struct(vec![
- Field::new("f1", array_1.data_type().clone(), true),
- Field::new("f2", array_2.data_type().clone(), true),
- ]);
-
- let mut struct_array_reader = StructArrayReader::new(
- struct_type,
- vec![Box::new(array_reader_1), Box::new(array_reader_2)],
- 1,
- 1,
- );
-
- let struct_array = struct_array_reader.next_batch(5).unwrap();
- let struct_array = struct_array.as_any().downcast_ref::<StructArray>().unwrap();
-
- assert_eq!(5, struct_array.len());
- assert_eq!(
- vec![true, false, false, false, false],
- (0..5)
- .map(|idx| struct_array.data_ref().is_null(idx))
- .collect::<Vec<bool>>()
- );
- assert_eq!(
- Some(vec![0, 1, 1, 1, 1].as_slice()),
- struct_array_reader.get_def_levels()
- );
- assert_eq!(
- Some(vec![1, 1, 1, 1, 1].as_slice()),
- struct_array_reader.get_rep_levels()
- );
- }
-
- #[test]
- fn test_create_array_reader() {
- let file = get_test_file("nulls.snappy.parquet");
- let file_reader = Arc::new(SerializedFileReader::new(file).unwrap());
-
- let file_metadata = file_reader.metadata().file_metadata();
- let arrow_schema = parquet_to_arrow_schema(
- file_metadata.schema_descr(),
- file_metadata.key_value_metadata(),
- )
- .unwrap();
-
- let array_reader = build_array_reader(
- file_reader.metadata().file_metadata().schema_descr_ptr(),
- arrow_schema,
- vec![0usize].into_iter(),
- file_reader,
- )
- .unwrap();
-
- // Create arrow types
- let arrow_type = ArrowType::Struct(vec![Field::new(
- "b_struct",
- ArrowType::Struct(vec![Field::new("b_c_int", ArrowType::Int32, true)]),
- true,
- )]);
-
- assert_eq!(array_reader.get_data_type(), &arrow_type);
- }
-
- #[test]
- fn test_list_array_reader() {
- // [[1, null, 2], null, [3, 4]]
- let array = Arc::new(PrimitiveArray::<ArrowInt32>::from(vec![
- Some(1),
- None,
- Some(2),
- None,
- Some(3),
- Some(4),
- ]));
- let item_array_reader = InMemoryArrayReader::new(
- ArrowType::Int32,
- array,
- Some(vec![3, 2, 3, 0, 3, 3]),
- Some(vec![0, 1, 1, 0, 0, 1]),
- );
-
- let mut list_array_reader = ListArrayReader::<i32>::new(
- Box::new(item_array_reader),
- ArrowType::List(Box::new(Field::new("item", ArrowType::Int32, true))),
- ArrowType::Int32,
- 1,
- 1,
- 0,
- 1,
- );
-
- let next_batch = list_array_reader.next_batch(1024).unwrap();
- let list_array = next_batch.as_any().downcast_ref::<ListArray>().unwrap();
-
- assert_eq!(3, list_array.len());
- // This passes as I expect
- assert_eq!(1, list_array.null_count());
-
- assert_eq!(
- list_array
- .value(0)
- .as_any()
- .downcast_ref::<PrimitiveArray<ArrowInt32>>()
- .unwrap(),
- &PrimitiveArray::<ArrowInt32>::from(vec![Some(1), None, Some(2)])
- );
-
- assert!(list_array.is_null(1));
-
- assert_eq!(
- list_array
- .value(2)
- .as_any()
- .downcast_ref::<PrimitiveArray<ArrowInt32>>()
- .unwrap(),
- &PrimitiveArray::<ArrowInt32>::from(vec![Some(3), Some(4)])
- );
- }
-
- #[test]
- fn test_large_list_array_reader() {
- // [[1, null, 2], null, [3, 4]]
- let array = Arc::new(PrimitiveArray::<ArrowInt32>::from(vec![
- Some(1),
- None,
- Some(2),
- None,
- Some(3),
- Some(4),
- ]));
- let item_array_reader = InMemoryArrayReader::new(
- ArrowType::Int32,
- array,
- Some(vec![3, 2, 3, 0, 3, 3]),
- Some(vec![0, 1, 1, 0, 0, 1]),
- );
-
- let mut list_array_reader = ListArrayReader::<i64>::new(
- Box::new(item_array_reader),
- ArrowType::LargeList(Box::new(Field::new("item", ArrowType::Int32, true))),
- ArrowType::Int32,
- 1,
- 1,
- 0,
- 1,
- );
-
- let next_batch = list_array_reader.next_batch(1024).unwrap();
- let list_array = next_batch
- .as_any()
- .downcast_ref::<LargeListArray>()
- .unwrap();
-
- assert_eq!(3, list_array.len());
-
- assert_eq!(
- list_array
- .value(0)
- .as_any()
- .downcast_ref::<PrimitiveArray<ArrowInt32>>()
- .unwrap(),
- &PrimitiveArray::<ArrowInt32>::from(vec![Some(1), None, Some(2)])
- );
-
- assert!(list_array.is_null(1));
-
- assert_eq!(
- list_array
- .value(2)
- .as_any()
- .downcast_ref::<PrimitiveArray<ArrowInt32>>()
- .unwrap(),
- &PrimitiveArray::<ArrowInt32>::from(vec![Some(3), Some(4)])
- );
- }
-}
diff --git a/parquet/src/arrow/arrow_array_reader.rs b/parquet/src/arrow/arrow_array_reader.rs
deleted file mode 100644
index c06d872..0000000
--- a/parquet/src/arrow/arrow_array_reader.rs
+++ /dev/null
@@ -1,1562 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-use super::array_reader::ArrayReader;
-use crate::arrow::schema::parquet_to_arrow_field;
-use crate::basic::Encoding;
-use crate::errors::{ParquetError, Result};
-use crate::{
- column::page::{Page, PageIterator},
- memory::ByteBufferPtr,
- schema::types::{ColumnDescPtr, ColumnDescriptor},
-};
-use arrow::{
- array::{ArrayRef, Int16Array},
- buffer::MutableBuffer,
- datatypes::{DataType as ArrowType, ToByteSlice},
-};
-use std::{any::Any, collections::VecDeque, marker::PhantomData};
-use std::{cell::RefCell, rc::Rc};
-
-struct UnzipIter<Source, Target, State> {
- shared_state: Rc<RefCell<State>>,
- select_item_buffer: fn(&mut State) -> &mut VecDeque<Target>,
- consume_source_item: fn(source_item: Source, state: &mut State) -> Target,
-}
-
-impl<Source, Target, State> UnzipIter<Source, Target, State> {
- fn new(
- shared_state: Rc<RefCell<State>>,
- item_buffer_selector: fn(&mut State) -> &mut VecDeque<Target>,
- source_item_consumer: fn(source_item: Source, state: &mut State) -> Target,
- ) -> Self {
- Self {
- shared_state,
- select_item_buffer: item_buffer_selector,
- consume_source_item: source_item_consumer,
- }
- }
-}
-
-trait UnzipIterState<T> {
- type SourceIter: Iterator<Item = T>;
- fn source_iter(&mut self) -> &mut Self::SourceIter;
-}
-
-impl<Source, Target, State: UnzipIterState<Source>> Iterator
- for UnzipIter<Source, Target, State>
-{
- type Item = Target;
-
- fn next(&mut self) -> Option<Self::Item> {
- let mut inner = self.shared_state.borrow_mut();
- // try to get one from the stored data
- (self.select_item_buffer)(&mut *inner)
- .pop_front()
- .or_else(||
- // nothing stored, we need a new element.
- inner.source_iter().next().map(|s| {
- (self.consume_source_item)(s, &mut inner)
- }))
- }
-}
-
-struct PageBufferUnzipIterState<V, L, It> {
- iter: It,
- value_iter_buffer: VecDeque<V>,
- def_level_iter_buffer: VecDeque<L>,
- rep_level_iter_buffer: VecDeque<L>,
-}
-
-impl<V, L, It: Iterator<Item = (V, L, L)>> UnzipIterState<(V, L, L)>
- for PageBufferUnzipIterState<V, L, It>
-{
- type SourceIter = It;
-
- #[inline]
- fn source_iter(&mut self) -> &mut Self::SourceIter {
- &mut self.iter
- }
-}
-
-type ValueUnzipIter<V, L, It> =
- UnzipIter<(V, L, L), V, PageBufferUnzipIterState<V, L, It>>;
-type LevelUnzipIter<V, L, It> =
- UnzipIter<(V, L, L), L, PageBufferUnzipIterState<V, L, It>>;
-type PageUnzipResult<V, L, It> = (
- ValueUnzipIter<V, L, It>,
- LevelUnzipIter<V, L, It>,
- LevelUnzipIter<V, L, It>,
-);
-
-fn unzip_iter<V, L, It: Iterator<Item = (V, L, L)>>(it: It) -> PageUnzipResult<V, L, It> {
- let shared_data = Rc::new(RefCell::new(PageBufferUnzipIterState {
- iter: it,
- value_iter_buffer: VecDeque::new(),
- def_level_iter_buffer: VecDeque::new(),
- rep_level_iter_buffer: VecDeque::new(),
- }));
-
- let value_iter = UnzipIter::new(
- shared_data.clone(),
- |state| &mut state.value_iter_buffer,
- |(v, d, r), state| {
- state.def_level_iter_buffer.push_back(d);
- state.rep_level_iter_buffer.push_back(r);
- v
- },
- );
-
- let def_level_iter = UnzipIter::new(
- shared_data.clone(),
- |state| &mut state.def_level_iter_buffer,
- |(v, d, r), state| {
- state.value_iter_buffer.push_back(v);
- state.rep_level_iter_buffer.push_back(r);
- d
- },
- );
-
- let rep_level_iter = UnzipIter::new(
- shared_data,
- |state| &mut state.rep_level_iter_buffer,
- |(v, d, r), state| {
- state.value_iter_buffer.push_back(v);
- state.def_level_iter_buffer.push_back(d);
- r
- },
- );
-
- (value_iter, def_level_iter, rep_level_iter)
-}
-
-pub trait ArrayConverter {
- fn convert_value_bytes(
- &self,
- value_decoder: &mut impl ValueDecoder,
- num_values: usize,
- ) -> Result<arrow::array::ArrayData>;
-}
-
-pub struct ArrowArrayReader<'a, C: ArrayConverter + 'a> {
- column_desc: ColumnDescPtr,
- data_type: ArrowType,
- def_level_decoder: Box<dyn ValueDecoder + 'a>,
- rep_level_decoder: Box<dyn ValueDecoder + 'a>,
- value_decoder: Box<dyn ValueDecoder + 'a>,
- last_def_levels: Option<Int16Array>,
- last_rep_levels: Option<Int16Array>,
- array_converter: C,
-}
-
-pub(crate) struct ColumnChunkContext {
- dictionary_values: Option<Vec<ByteBufferPtr>>,
-}
-
-impl ColumnChunkContext {
- fn new() -> Self {
- Self {
- dictionary_values: None,
- }
- }
-
- fn set_dictionary(&mut self, dictionary_values: Vec<ByteBufferPtr>) {
- self.dictionary_values = Some(dictionary_values);
- }
-}
-
-type PageDecoderTuple = (
- Box<dyn ValueDecoder>,
- Box<dyn ValueDecoder>,
- Box<dyn ValueDecoder>,
-);
-
-impl<'a, C: ArrayConverter + 'a> ArrowArrayReader<'a, C> {
- pub fn try_new<P: PageIterator + 'a>(
- column_chunk_iterator: P,
- column_desc: ColumnDescPtr,
- array_converter: C,
- arrow_type: Option<ArrowType>,
- ) -> Result<Self> {
- let data_type = match arrow_type {
- Some(t) => t,
- None => parquet_to_arrow_field(column_desc.as_ref())?
- .data_type()
- .clone(),
- };
- type PageIteratorItem = Result<(Page, Rc<RefCell<ColumnChunkContext>>)>;
- let page_iter = column_chunk_iterator
- // build iterator of pages across column chunks
- .flat_map(|x| -> Box<dyn Iterator<Item = PageIteratorItem>> {
- // attach column chunk context
- let context = Rc::new(RefCell::new(ColumnChunkContext::new()));
- match x {
- Ok(page_reader) => Box::new(
- page_reader.map(move |pr| pr.map(|p| (p, context.clone()))),
- ),
- // errors from reading column chunks / row groups are propagated to page level
- Err(e) => Box::new(std::iter::once(Err(e))),
- }
- });
- // capture a clone of column_desc in closure so that it can outlive current function
- let map_page_fn_factory = |column_desc: ColumnDescPtr| {
- move |x: Result<(Page, Rc<RefCell<ColumnChunkContext>>)>| {
- x.and_then(|(page, context)| {
- Self::map_page(page, context, column_desc.as_ref())
- })
- }
- };
- let map_page_fn = map_page_fn_factory(column_desc.clone());
- // map page iterator into tuple of buffer iterators for (values, def levels, rep levels)
- // errors from lower levels are surfaced through the value decoder iterator
- let decoder_iter = page_iter.map(map_page_fn).map(|x| match x {
- Ok(iter_tuple) => iter_tuple,
- // errors from reading pages are propagated to decoder iterator level
- Err(e) => Self::map_page_error(e),
- });
- // split tuple iterator into separate iterators for (values, def levels, rep levels)
- let (value_iter, def_level_iter, rep_level_iter) = unzip_iter(decoder_iter);
-
- Ok(Self {
- column_desc,
- data_type,
- def_level_decoder: Box::new(CompositeValueDecoder::new(def_level_iter)),
- rep_level_decoder: Box::new(CompositeValueDecoder::new(rep_level_iter)),
- value_decoder: Box::new(CompositeValueDecoder::new(value_iter)),
- last_def_levels: None,
- last_rep_levels: None,
- array_converter,
- })
- }
-
- #[inline]
- fn def_levels_available(column_desc: &ColumnDescriptor) -> bool {
- column_desc.max_def_level() > 0
- }
-
- #[inline]
- fn rep_levels_available(column_desc: &ColumnDescriptor) -> bool {
- column_desc.max_rep_level() > 0
- }
-
- fn map_page_error(err: ParquetError) -> PageDecoderTuple {
- (
- Box::new(<dyn ValueDecoder>::once(Err(err.clone()))),
- Box::new(<dyn ValueDecoder>::once(Err(err.clone()))),
- Box::new(<dyn ValueDecoder>::once(Err(err))),
- )
- }
-
- // Split Result<Page> into Result<(Iterator<Values>, Iterator<DefLevels>, Iterator<RepLevels>)>
- // this method could fail, e.g. if the page encoding is not supported
- fn map_page(
- page: Page,
- column_chunk_context: Rc<RefCell<ColumnChunkContext>>,
- column_desc: &ColumnDescriptor,
- ) -> Result<PageDecoderTuple> {
- use crate::encodings::levels::LevelDecoder;
- match page {
- Page::DictionaryPage {
- buf,
- num_values,
- encoding,
- ..
- } => {
- let mut column_chunk_context = column_chunk_context.borrow_mut();
- if column_chunk_context.dictionary_values.is_some() {
- return Err(general_err!(
- "Column chunk cannot have more than one dictionary"
- ));
- }
- // create plain decoder for dictionary values
- let mut dict_decoder = Self::get_dictionary_page_decoder(
- buf,
- num_values as usize,
- encoding,
- column_desc,
- )?;
- // decode and cache dictionary values
- let dictionary_values = dict_decoder.read_dictionary_values()?;
- column_chunk_context.set_dictionary(dictionary_values);
-
- // a dictionary page doesn't return any values
- Ok((
- Box::new(<dyn ValueDecoder>::empty()),
- Box::new(<dyn ValueDecoder>::empty()),
- Box::new(<dyn ValueDecoder>::empty()),
- ))
- }
- Page::DataPage {
- buf,
- num_values,
- encoding,
- def_level_encoding,
- rep_level_encoding,
- statistics: _,
- } => {
- let mut buffer_ptr = buf;
- // create rep level decoder iterator
- let rep_level_iter: Box<dyn ValueDecoder> =
- if Self::rep_levels_available(&column_desc) {
- let mut rep_decoder = LevelDecoder::v1(
- rep_level_encoding,
- column_desc.max_rep_level(),
- );
- let rep_level_byte_len =
- rep_decoder.set_data(num_values as usize, buffer_ptr.all());
- // advance buffer pointer
- buffer_ptr = buffer_ptr.start_from(rep_level_byte_len);
- Box::new(LevelValueDecoder::new(rep_decoder))
- } else {
- Box::new(<dyn ValueDecoder>::once(Err(ParquetError::General(
- "rep levels are not available".to_string(),
- ))))
- };
- // create def level decoder iterator
- let def_level_iter: Box<dyn ValueDecoder> =
- if Self::def_levels_available(&column_desc) {
- let mut def_decoder = LevelDecoder::v1(
- def_level_encoding,
- column_desc.max_def_level(),
- );
- let def_levels_byte_len =
- def_decoder.set_data(num_values as usize, buffer_ptr.all());
- // advance buffer pointer
- buffer_ptr = buffer_ptr.start_from(def_levels_byte_len);
- Box::new(LevelValueDecoder::new(def_decoder))
- } else {
- Box::new(<dyn ValueDecoder>::once(Err(ParquetError::General(
- "def levels are not available".to_string(),
- ))))
- };
- // create value decoder iterator
- let value_iter = Self::get_value_decoder(
- buffer_ptr,
- num_values as usize,
- encoding,
- column_desc,
- column_chunk_context,
- )?;
- Ok((value_iter, def_level_iter, rep_level_iter))
- }
- Page::DataPageV2 {
- buf,
- num_values,
- encoding,
- num_nulls: _,
- num_rows: _,
- def_levels_byte_len,
- rep_levels_byte_len,
- is_compressed: _,
- statistics: _,
- } => {
- let mut offset = 0;
- // create rep level decoder iterator
- let rep_level_iter: Box<dyn ValueDecoder> =
- if Self::rep_levels_available(&column_desc) {
- let rep_levels_byte_len = rep_levels_byte_len as usize;
- let mut rep_decoder =
- LevelDecoder::v2(column_desc.max_rep_level());
- rep_decoder.set_data_range(
- num_values as usize,
- &buf,
- offset,
- rep_levels_byte_len,
- );
- offset += rep_levels_byte_len;
- Box::new(LevelValueDecoder::new(rep_decoder))
- } else {
- Box::new(<dyn ValueDecoder>::once(Err(ParquetError::General(
- "rep levels are not available".to_string(),
- ))))
- };
- // create def level decoder iterator
- let def_level_iter: Box<dyn ValueDecoder> =
- if Self::def_levels_available(&column_desc) {
- let def_levels_byte_len = def_levels_byte_len as usize;
- let mut def_decoder =
- LevelDecoder::v2(column_desc.max_def_level());
- def_decoder.set_data_range(
- num_values as usize,
- &buf,
- offset,
- def_levels_byte_len,
- );
- offset += def_levels_byte_len;
- Box::new(LevelValueDecoder::new(def_decoder))
- } else {
- Box::new(<dyn ValueDecoder>::once(Err(ParquetError::General(
- "def levels are not available".to_string(),
- ))))
- };
-
- // create value decoder iterator
- let values_buffer = buf.start_from(offset);
- let value_iter = Self::get_value_decoder(
- values_buffer,
- num_values as usize,
- encoding,
- column_desc,
- column_chunk_context,
- )?;
- Ok((value_iter, def_level_iter, rep_level_iter))
- }
- }
- }
-
- fn get_dictionary_page_decoder(
- values_buffer: ByteBufferPtr,
- num_values: usize,
- mut encoding: Encoding,
- column_desc: &ColumnDescriptor,
- ) -> Result<Box<dyn DictionaryValueDecoder>> {
- if encoding == Encoding::PLAIN || encoding == Encoding::PLAIN_DICTIONARY {
- encoding = Encoding::RLE_DICTIONARY
- }
-
- if encoding == Encoding::RLE_DICTIONARY {
- Ok(
- Self::get_plain_value_decoder(values_buffer, num_values, column_desc)
- .into_dictionary_decoder(),
- )
- } else {
- Err(nyi_err!(
- "Invalid/Unsupported encoding type for dictionary: {}",
- encoding
- ))
- }
- }
-
- fn get_value_decoder(
- values_buffer: ByteBufferPtr,
- num_values: usize,
- mut encoding: Encoding,
- column_desc: &ColumnDescriptor,
- column_chunk_context: Rc<RefCell<ColumnChunkContext>>,
- ) -> Result<Box<dyn ValueDecoder>> {
- if encoding == Encoding::PLAIN_DICTIONARY {
- encoding = Encoding::RLE_DICTIONARY;
- }
-
- match encoding {
- Encoding::PLAIN => {
- Ok(
- Self::get_plain_value_decoder(values_buffer, num_values, column_desc)
- .into_value_decoder(),
- )
- }
- Encoding::RLE_DICTIONARY => {
- if column_chunk_context.borrow().dictionary_values.is_some() {
- let value_bit_len = Self::get_column_physical_bit_len(column_desc);
- let dictionary_decoder: Box<dyn ValueDecoder> = if value_bit_len == 0
- {
- Box::new(VariableLenDictionaryDecoder::new(
- column_chunk_context,
- values_buffer,
- num_values,
- ))
- } else {
- Box::new(FixedLenDictionaryDecoder::new(
- column_chunk_context,
- values_buffer,
- num_values,
- value_bit_len,
- ))
- };
- Ok(dictionary_decoder)
- } else {
- Err(general_err!("Dictionary values have not been initialized."))
- }
- }
- // Encoding::RLE => Box::new(RleValueDecoder::new()),
- // Encoding::DELTA_BINARY_PACKED => Box::new(DeltaBitPackDecoder::new()),
- // Encoding::DELTA_LENGTH_BYTE_ARRAY => Box::new(DeltaLengthByteArrayDecoder::new()),
- // Encoding::DELTA_BYTE_ARRAY => Box::new(DeltaByteArrayDecoder::new()),
- e => return Err(nyi_err!("Encoding {} is not supported", e)),
- }
- }
-
- fn get_column_physical_bit_len(column_desc: &ColumnDescriptor) -> usize {
- use crate::basic::Type as PhysicalType;
- // parquet only supports a limited number of physical types
- // later converters cast to a more specific arrow / logical type if necessary
- match column_desc.physical_type() {
- PhysicalType::BOOLEAN => 1,
- PhysicalType::INT32 | PhysicalType::FLOAT => 32,
- PhysicalType::INT64 | PhysicalType::DOUBLE => 64,
- PhysicalType::INT96 => 96,
- PhysicalType::BYTE_ARRAY => 0,
- PhysicalType::FIXED_LEN_BYTE_ARRAY => column_desc.type_length() as usize * 8,
- }
- }
-
- fn get_plain_value_decoder(
- values_buffer: ByteBufferPtr,
- num_values: usize,
- column_desc: &ColumnDescriptor,
- ) -> Box<dyn PlainValueDecoder> {
- let value_bit_len = Self::get_column_physical_bit_len(column_desc);
- if value_bit_len == 0 {
- Box::new(VariableLenPlainDecoder::new(values_buffer, num_values))
- } else {
- Box::new(FixedLenPlainDecoder::new(
- values_buffer,
- num_values,
- value_bit_len,
- ))
- }
- }
-
- fn build_level_array(
- level_decoder: &mut impl ValueDecoder,
- batch_size: usize,
- ) -> Result<Int16Array> {
- use arrow::datatypes::Int16Type;
- let level_converter = PrimitiveArrayConverter::<Int16Type>::new();
- let array_data =
- level_converter.convert_value_bytes(level_decoder, batch_size)?;
- Ok(Int16Array::from(array_data))
- }
-}
-
-impl<C: ArrayConverter> ArrayReader for ArrowArrayReader<'static, C> {
- fn as_any(&self) -> &dyn Any {
- self
- }
-
- fn get_data_type(&self) -> &ArrowType {
- &self.data_type
- }
-
- fn next_batch(&mut self, batch_size: usize) -> Result<ArrayRef> {
- if Self::rep_levels_available(&self.column_desc) {
- // read rep levels if available
- let rep_level_array =
- Self::build_level_array(&mut self.rep_level_decoder, batch_size)?;
- self.last_rep_levels = Some(rep_level_array);
- }
-
- // check if def levels are available
- let (values_to_read, null_bitmap_array) =
- if !Self::def_levels_available(&self.column_desc) {
- // if no def levels - just read (up to) batch_size values
- (batch_size, None)
- } else {
- // if def levels are available - they determine how many values will be read
- // decode def levels, return first error if any
- let def_level_array =
- Self::build_level_array(&mut self.def_level_decoder, batch_size)?;
- let def_level_count = def_level_array.len();
- // use eq_scalar to efficiently build null bitmap array from def levels
- let null_bitmap_array = arrow::compute::eq_scalar(
- &def_level_array,
- self.column_desc.max_def_level(),
- )?;
- self.last_def_levels = Some(def_level_array);
- // efficiently calculate values to read
- let values_to_read = null_bitmap_array
- .values()
- .count_set_bits_offset(0, def_level_count);
- let maybe_null_bitmap = if values_to_read != null_bitmap_array.len() {
- Some(null_bitmap_array)
- } else {
- // shortcut if no NULLs
- None
- };
- (values_to_read, maybe_null_bitmap)
- };
-
- // read a batch of values
- // converter only creates a no-null / all value array data
- let mut value_array_data = self
- .array_converter
- .convert_value_bytes(&mut self.value_decoder, values_to_read)?;
-
- if let Some(null_bitmap_array) = null_bitmap_array {
- // Only if def levels are available - insert null values efficiently using MutableArrayData.
- // This will require value bytes to be copied again, but converter requirements are reduced.
- // With a small number of NULLs, this will only be a few copies of large byte sequences.
- let actual_batch_size = null_bitmap_array.len();
- // use_nulls is false, because null_bitmap_array is already calculated and re-used
- let mut mutable = arrow::array::MutableArrayData::new(
- vec![&value_array_data],
- false,
- actual_batch_size,
- );
- // SlicesIterator slices only the true values, NULLs are inserted to fill any gaps
- arrow::compute::SlicesIterator::new(&null_bitmap_array).for_each(
- |(start, end)| {
- // the gap needs to be filled with NULLs
- if start > mutable.len() {
- let nulls_to_add = start - mutable.len();
- mutable.extend_nulls(nulls_to_add);
- }
- // fill values, adjust start and end with NULL count so far
- let nulls_added = mutable.null_count();
- mutable.extend(0, start - nulls_added, end - nulls_added);
- },
- );
- // any remaining part is NULLs
- if mutable.len() < actual_batch_size {
- let nulls_to_add = actual_batch_size - mutable.len();
- mutable.extend_nulls(nulls_to_add);
- }
-
- value_array_data = mutable
- .into_builder()
- .null_bit_buffer(null_bitmap_array.values().clone())
- .build();
- }
- let mut array = arrow::array::make_array(value_array_data);
- if array.data_type() != &self.data_type {
- // cast array to self.data_type if necessary
- array = arrow::compute::cast(&array, &self.data_type)?
- }
- Ok(array)
- }
-
- fn get_def_levels(&self) -> Option<&[i16]> {
- self.last_def_levels.as_ref().map(|x| x.values())
- }
-
- fn get_rep_levels(&self) -> Option<&[i16]> {
- self.last_rep_levels.as_ref().map(|x| x.values())
- }
-}
-
-use crate::encodings::rle::RleDecoder;
-
-pub trait ValueDecoder {
- fn read_value_bytes(
- &mut self,
- num_values: usize,
- read_bytes: &mut dyn FnMut(&[u8], usize),
- ) -> Result<usize>;
-}
-
-trait DictionaryValueDecoder {
- fn read_dictionary_values(&mut self) -> Result<Vec<ByteBufferPtr>>;
-}
-
-trait PlainValueDecoder: ValueDecoder + DictionaryValueDecoder {
- fn into_value_decoder(self: Box<Self>) -> Box<dyn ValueDecoder>;
- fn into_dictionary_decoder(self: Box<Self>) -> Box<dyn DictionaryValueDecoder>;
-}
-
-impl<T> PlainValueDecoder for T
-where
- T: ValueDecoder + DictionaryValueDecoder + 'static,
-{
- fn into_value_decoder(self: Box<T>) -> Box<dyn ValueDecoder> {
- self
- }
-
- fn into_dictionary_decoder(self: Box<T>) -> Box<dyn DictionaryValueDecoder> {
- self
- }
-}
-
-impl dyn ValueDecoder {
- fn empty() -> impl ValueDecoder {
- SingleValueDecoder::new(Ok(0))
- }
-
- fn once(value: Result<usize>) -> impl ValueDecoder {
- SingleValueDecoder::new(value)
- }
-}
-
-impl ValueDecoder for Box<dyn ValueDecoder> {
- #[inline]
- fn read_value_bytes(
- &mut self,
- num_values: usize,
- read_bytes: &mut dyn FnMut(&[u8], usize),
- ) -> Result<usize> {
- self.as_mut().read_value_bytes(num_values, read_bytes)
- }
-}
-
-struct SingleValueDecoder {
- value: Result<usize>,
-}
-
-impl SingleValueDecoder {
- fn new(value: Result<usize>) -> Self {
- Self { value }
- }
-}
-
-impl ValueDecoder for SingleValueDecoder {
- fn read_value_bytes(
- &mut self,
- _num_values: usize,
- _read_bytes: &mut dyn FnMut(&[u8], usize),
- ) -> Result<usize> {
- self.value.clone()
- }
-}
-
-struct CompositeValueDecoder<I: Iterator<Item = Box<dyn ValueDecoder>>> {
- current_decoder: Option<Box<dyn ValueDecoder>>,
- decoder_iter: I,
-}
-
-impl<I: Iterator<Item = Box<dyn ValueDecoder>>> CompositeValueDecoder<I> {
- fn new(mut decoder_iter: I) -> Self {
- let current_decoder = decoder_iter.next();
- Self {
- current_decoder,
- decoder_iter,
- }
- }
-}
-
-impl<I: Iterator<Item = Box<dyn ValueDecoder>>> ValueDecoder
- for CompositeValueDecoder<I>
-{
- fn read_value_bytes(
- &mut self,
- num_values: usize,
- read_bytes: &mut dyn FnMut(&[u8], usize),
- ) -> Result<usize> {
- let mut values_to_read = num_values;
- while values_to_read > 0 {
- let value_decoder = match self.current_decoder.as_mut() {
- Some(d) => d,
- // no more decoders
- None => break,
- };
- while values_to_read > 0 {
- let values_read =
- value_decoder.read_value_bytes(values_to_read, read_bytes)?;
- if values_read > 0 {
- values_to_read -= values_read;
- } else {
- // no more values in current decoder
- self.current_decoder = self.decoder_iter.next();
- break;
- }
- }
- }
-
- Ok(num_values - values_to_read)
- }
-}
-
-struct LevelValueDecoder {
- level_decoder: crate::encodings::levels::LevelDecoder,
- level_value_buffer: Vec<i16>,
-}
-
-impl LevelValueDecoder {
- fn new(level_decoder: crate::encodings::levels::LevelDecoder) -> Self {
- Self {
- level_decoder,
- level_value_buffer: vec![0i16; 2048],
- }
- }
-}
-
-impl ValueDecoder for LevelValueDecoder {
- fn read_value_bytes(
- &mut self,
- num_values: usize,
- read_bytes: &mut dyn FnMut(&[u8], usize),
- ) -> Result<usize> {
- let value_size = std::mem::size_of::<i16>();
- let mut total_values_read = 0;
- while total_values_read < num_values {
- let values_to_read = std::cmp::min(
- num_values - total_values_read,
- self.level_value_buffer.len(),
- );
- let values_read = match self
- .level_decoder
- .get(&mut self.level_value_buffer[..values_to_read])
- {
- Ok(values_read) => values_read,
- Err(e) => return Err(e),
- };
- if values_read > 0 {
- let level_value_bytes =
- &self.level_value_buffer.to_byte_slice()[..values_read * value_size];
- read_bytes(level_value_bytes, values_read);
- total_values_read += values_read;
- } else {
- break;
- }
- }
- Ok(total_values_read)
- }
-}
-
-pub(crate) struct FixedLenPlainDecoder {
- data: ByteBufferPtr,
- num_values: usize,
- value_bit_len: usize,
-}
-
-impl FixedLenPlainDecoder {
- pub(crate) fn new(
- data: ByteBufferPtr,
- num_values: usize,
- value_bit_len: usize,
- ) -> Self {
- Self {
- data,
- num_values,
- value_bit_len,
- }
- }
-}
-
-impl DictionaryValueDecoder for FixedLenPlainDecoder {
- fn read_dictionary_values(&mut self) -> Result<Vec<ByteBufferPtr>> {
- let value_byte_len = self.value_bit_len / 8;
- let available_values = self.data.len() / value_byte_len;
- let values_to_read = std::cmp::min(available_values, self.num_values);
- let byte_len = values_to_read * value_byte_len;
- let values = vec![self.data.range(0, byte_len)];
- self.num_values = 0;
- self.data.set_range(self.data.start(), 0);
- Ok(values)
- }
-}
-
-impl ValueDecoder for FixedLenPlainDecoder {
- fn read_value_bytes(
- &mut self,
- num_values: usize,
- read_bytes: &mut dyn FnMut(&[u8], usize),
- ) -> Result<usize> {
- let available_values = self.data.len() * 8 / self.value_bit_len;
- if available_values > 0 {
- let values_to_read = std::cmp::min(available_values, num_values);
- let byte_len = values_to_read * self.value_bit_len / 8;
- read_bytes(&self.data.data()[..byte_len], values_to_read);
- self.data
- .set_range(self.data.start() + byte_len, self.data.len() - byte_len);
- Ok(values_to_read)
- } else {
- Ok(0)
- }
- }
-}
-
-pub(crate) struct VariableLenPlainDecoder {
- data: ByteBufferPtr,
- num_values: usize,
- position: usize,
-}
-
-impl VariableLenPlainDecoder {
- pub(crate) fn new(data: ByteBufferPtr, num_values: usize) -> Self {
- Self {
- data,
- num_values,
- position: 0,
- }
- }
-}
-
-impl DictionaryValueDecoder for VariableLenPlainDecoder {
- fn read_dictionary_values(&mut self) -> Result<Vec<ByteBufferPtr>> {
- const LEN_SIZE: usize = std::mem::size_of::<u32>();
- let data = self.data.data();
- let data_len = data.len();
- let values_to_read = self.num_values;
- let mut values = Vec::with_capacity(values_to_read);
- let mut values_read = 0;
- while self.position < data_len && values_read < values_to_read {
- let len: usize =
- read_num_bytes!(u32, LEN_SIZE, data[self.position..]) as usize;
- self.position += LEN_SIZE;
- if data_len < self.position + len {
- return Err(eof_err!("Not enough bytes to decode"));
- }
- values.push(self.data.range(self.position, len));
- self.position += len;
- values_read += 1;
- }
- self.num_values -= values_read;
- Ok(values)
- }
-}
-
-impl ValueDecoder for VariableLenPlainDecoder {
- fn read_value_bytes(
- &mut self,
- num_values: usize,
- read_bytes: &mut dyn FnMut(&[u8], usize),
- ) -> Result<usize> {
- const LEN_SIZE: usize = std::mem::size_of::<u32>();
- let data = self.data.data();
- let data_len = data.len();
- let values_to_read = std::cmp::min(self.num_values, num_values);
- let mut values_read = 0;
- while self.position < data_len && values_read < values_to_read {
- let len: usize =
- read_num_bytes!(u32, LEN_SIZE, data[self.position..]) as usize;
- self.position += LEN_SIZE;
- if data_len < self.position + len {
- return Err(eof_err!("Not enough bytes to decode"));
- }
- read_bytes(&data[self.position..][..len], 1);
- self.position += len;
- values_read += 1;
- }
- self.num_values -= values_read;
- Ok(values_read)
- }
-}
-
-pub(crate) struct FixedLenDictionaryDecoder {
- context_ref: Rc<RefCell<ColumnChunkContext>>,
- key_data_bufer: ByteBufferPtr,
- num_values: usize,
- rle_decoder: RleDecoder,
- value_byte_len: usize,
- keys_buffer: Vec<i32>,
-}
-
-impl FixedLenDictionaryDecoder {
- pub(crate) fn new(
- column_chunk_context: Rc<RefCell<ColumnChunkContext>>,
- key_data_bufer: ByteBufferPtr,
- num_values: usize,
- value_bit_len: usize,
- ) -> Self {
- assert!(
- value_bit_len % 8 == 0,
- "value_bit_size must be a multiple of 8"
- );
- // First byte in `data` is bit width
- let bit_width = key_data_bufer.data()[0];
- let mut rle_decoder = RleDecoder::new(bit_width);
- rle_decoder.set_data(key_data_bufer.start_from(1));
-
- Self {
- context_ref: column_chunk_context,
- key_data_bufer,
- num_values,
- rle_decoder,
- value_byte_len: value_bit_len / 8,
- keys_buffer: vec![0; 2048],
- }
- }
-}
-
-impl ValueDecoder for FixedLenDictionaryDecoder {
- fn read_value_bytes(
- &mut self,
- num_values: usize,
- read_bytes: &mut dyn FnMut(&[u8], usize),
- ) -> Result<usize> {
- if self.num_values == 0 {
- return Ok(0);
- }
- let context = self.context_ref.borrow();
- let values = context.dictionary_values.as_ref().unwrap();
- let input_value_bytes = values[0].data();
- // read no more than available values or requested values
- let values_to_read = std::cmp::min(self.num_values, num_values);
- let mut values_read = 0;
- while values_read < values_to_read {
- // read values in batches of up to self.keys_buffer.len()
- let keys_to_read =
- std::cmp::min(values_to_read - values_read, self.keys_buffer.len());
- let keys_read = match self
- .rle_decoder
- .get_batch(&mut self.keys_buffer[..keys_to_read])
- {
- Ok(keys_read) => keys_read,
- Err(e) => return Err(e),
- };
- if keys_read == 0 {
- self.num_values = 0;
- return Ok(values_read);
- }
- for i in 0..keys_read {
- let key = self.keys_buffer[i] as usize;
- read_bytes(
- &input_value_bytes[key * self.value_byte_len..]
- [..self.value_byte_len],
- 1,
- );
- }
- values_read += keys_read;
- }
- self.num_values -= values_read;
- Ok(values_read)
- }
-}
-
-pub(crate) struct VariableLenDictionaryDecoder {
- context_ref: Rc<RefCell<ColumnChunkContext>>,
- key_data_bufer: ByteBufferPtr,
- num_values: usize,
- rle_decoder: RleDecoder,
- keys_buffer: Vec<i32>,
-}
-
-impl VariableLenDictionaryDecoder {
- pub(crate) fn new(
- column_chunk_context: Rc<RefCell<ColumnChunkContext>>,
- key_data_bufer: ByteBufferPtr,
- num_values: usize,
- ) -> Self {
- // First byte in `data` is bit width
- let bit_width = key_data_bufer.data()[0];
- let mut rle_decoder = RleDecoder::new(bit_width);
- rle_decoder.set_data(key_data_bufer.start_from(1));
-
- Self {
- context_ref: column_chunk_context,
- key_data_bufer,
- num_values,
- rle_decoder,
- keys_buffer: vec![0; 2048],
- }
- }
-}
-
-impl ValueDecoder for VariableLenDictionaryDecoder {
- fn read_value_bytes(
- &mut self,
- num_values: usize,
- read_bytes: &mut dyn FnMut(&[u8], usize),
- ) -> Result<usize> {
- if self.num_values == 0 {
- return Ok(0);
- }
- let context = self.context_ref.borrow();
- let values = context.dictionary_values.as_ref().unwrap();
- let values_to_read = std::cmp::min(self.num_values, num_values);
- let mut values_read = 0;
- while values_read < values_to_read {
- // read values in batches of up to self.keys_buffer.len()
- let keys_to_read =
- std::cmp::min(values_to_read - values_read, self.keys_buffer.len());
- let keys_read = match self
- .rle_decoder
- .get_batch(&mut self.keys_buffer[..keys_to_read])
- {
- Ok(keys_read) => keys_read,
- Err(e) => return Err(e),
- };
- if keys_read == 0 {
- self.num_values = 0;
- return Ok(values_read);
- }
- for i in 0..keys_read {
- let key = self.keys_buffer[i] as usize;
- read_bytes(values[key].data(), 1);
- }
- values_read += keys_read;
- }
- self.num_values -= values_read;
- Ok(values_read)
- }
-}
-
-use arrow::datatypes::ArrowPrimitiveType;
-
-pub struct PrimitiveArrayConverter<T: ArrowPrimitiveType> {
- _phantom_data: PhantomData<T>,
-}
-
-impl<T: ArrowPrimitiveType> PrimitiveArrayConverter<T> {
- pub fn new() -> Self {
- Self {
- _phantom_data: PhantomData,
- }
- }
-}
-
-impl<T: ArrowPrimitiveType> ArrayConverter for PrimitiveArrayConverter<T> {
- fn convert_value_bytes(
- &self,
- value_decoder: &mut impl ValueDecoder,
- num_values: usize,
- ) -> Result<arrow::array::ArrayData> {
- let value_size = T::get_byte_width();
- let values_byte_capacity = num_values * value_size;
- let mut values_buffer = MutableBuffer::new(values_byte_capacity);
-
- value_decoder.read_value_bytes(num_values, &mut |value_bytes, _| {
- values_buffer.extend_from_slice(value_bytes);
- })?;
-
- // calculate actual data_len, which may be different from the iterator's upper bound
- let value_count = values_buffer.len() / value_size;
- let array_data = arrow::array::ArrayData::builder(T::DATA_TYPE)
- .len(value_count)
- .add_buffer(values_buffer.into())
- .build();
- Ok(array_data)
- }
-}
-
-pub struct StringArrayConverter {}
-
-impl StringArrayConverter {
- pub fn new() -> Self {
- Self {}
- }
-}
-
-impl ArrayConverter for StringArrayConverter {
- fn convert_value_bytes(
- &self,
- value_decoder: &mut impl ValueDecoder,
- num_values: usize,
- ) -> Result<arrow::array::ArrayData> {
- use arrow::datatypes::ArrowNativeType;
- let offset_size = std::mem::size_of::<i32>();
- let mut offsets_buffer = MutableBuffer::new((num_values + 1) * offset_size);
- // allocate initial capacity of 1 byte for each item
- let values_byte_capacity = num_values;
- let mut values_buffer = MutableBuffer::new(values_byte_capacity);
-
- let mut length_so_far = i32::default();
- offsets_buffer.push(length_so_far);
-
- value_decoder.read_value_bytes(num_values, &mut |value_bytes, values_read| {
- debug_assert_eq!(
- values_read, 1,
- "offset length value buffers can only contain bytes for a single value"
- );
- length_so_far +=
- <i32 as ArrowNativeType>::from_usize(value_bytes.len()).unwrap();
- // this should be safe because a ValueDecoder should not read more than num_values
- unsafe {
- offsets_buffer.push_unchecked(length_so_far);
- }
- values_buffer.extend_from_slice(value_bytes);
- })?;
- // calculate actual data_len, which may be different from the iterator's upper bound
- let data_len = (offsets_buffer.len() / offset_size) - 1;
- let array_data = arrow::array::ArrayData::builder(ArrowType::Utf8)
- .len(data_len)
- .add_buffer(offsets_buffer.into())
- .add_buffer(values_buffer.into())
- .build();
- Ok(array_data)
- }
-}
-
-#[cfg(test)]
-mod tests {
- use super::*;
- use crate::column::page::Page;
- use crate::data_type::ByteArray;
- use crate::data_type::ByteArrayType;
- use crate::schema::parser::parse_message_type;
- use crate::schema::types::SchemaDescriptor;
- use crate::util::test_common::page_util::{
- DataPageBuilder, DataPageBuilderImpl, InMemoryPageIterator,
- };
- use crate::{
- basic::Encoding, column::page::PageReader, schema::types::SchemaDescPtr,
- };
- use arrow::array::{PrimitiveArray, StringArray};
- use arrow::datatypes::Int32Type as ArrowInt32;
- use rand::{distributions::uniform::SampleUniform, thread_rng, Rng};
- use std::sync::Arc;
-
- /// Iterator for testing reading empty columns
- struct EmptyPageIterator {
- schema: SchemaDescPtr,
- }
-
- impl EmptyPageIterator {
- fn new(schema: SchemaDescPtr) -> Self {
- EmptyPageIterator { schema }
- }
- }
-
- impl Iterator for EmptyPageIterator {
- type Item = Result<Box<dyn PageReader>>;
-
- fn next(&mut self) -> Option<Self::Item> {
- None
- }
- }
-
- impl PageIterator for EmptyPageIterator {
- fn schema(&mut self) -> Result<SchemaDescPtr> {
- Ok(self.schema.clone())
- }
-
- fn column_schema(&mut self) -> Result<ColumnDescPtr> {
- Ok(self.schema.column(0))
- }
- }
-
- #[test]
- fn test_array_reader_empty_pages() {
- // Construct column schema
- let message_type = "
- message test_schema {
- REQUIRED INT32 leaf;
- }
- ";
-
- let schema = parse_message_type(message_type)
- .map(|t| Arc::new(SchemaDescriptor::new(Arc::new(t))))
- .unwrap();
-
- let column_desc = schema.column(0);
- let page_iterator = EmptyPageIterator::new(schema);
-
- let converter = PrimitiveArrayConverter::<arrow::datatypes::Int32Type>::new();
- let mut array_reader =
- ArrowArrayReader::try_new(page_iterator, column_desc, converter, None)
- .unwrap();
-
- // expect no values to be read
- let array = array_reader.next_batch(50).unwrap();
- assert!(array.is_empty());
- }
-
- fn make_column_chunks<T: crate::data_type::DataType>(
- column_desc: ColumnDescPtr,
- encoding: Encoding,
- num_levels: usize,
- min_value: T::T,
- max_value: T::T,
- def_levels: &mut Vec<i16>,
- rep_levels: &mut Vec<i16>,
- values: &mut Vec<T::T>,
- page_lists: &mut Vec<Vec<Page>>,
- use_v2: bool,
- num_chunks: usize,
- ) where
- T::T: PartialOrd + SampleUniform + Copy,
- {
- for _i in 0..num_chunks {
- let mut pages = VecDeque::new();
- let mut data = Vec::new();
- let mut page_def_levels = Vec::new();
- let mut page_rep_levels = Vec::new();
-
- crate::util::test_common::make_pages::<T>(
- column_desc.clone(),
- encoding,
- 1,
- num_levels,
- min_value,
- max_value,
- &mut page_def_levels,
- &mut page_rep_levels,
- &mut data,
- &mut pages,
- use_v2,
- );
-
- def_levels.append(&mut page_def_levels);
- rep_levels.append(&mut page_rep_levels);
- values.append(&mut data);
- page_lists.push(Vec::from(pages));
- }
- }
-
- #[test]
- fn test_primitive_array_reader_data() {
- // Construct column schema
- let message_type = "
- message test_schema {
- REQUIRED INT32 leaf;
- }
- ";
-
- let schema = parse_message_type(message_type)
- .map(|t| Arc::new(SchemaDescriptor::new(Arc::new(t))))
- .unwrap();
-
- let column_desc = schema.column(0);
-
- // Construct page iterator
- {
- let mut data = Vec::new();
- let mut page_lists = Vec::new();
- make_column_chunks::<crate::data_type::Int32Type>(
- column_desc.clone(),
- Encoding::PLAIN,
- 100,
- 1,
- 200,
- &mut Vec::new(),
- &mut Vec::new(),
- &mut data,
- &mut page_lists,
- true,
- 2,
- );
- let page_iterator =
- InMemoryPageIterator::new(schema, column_desc.clone(), page_lists);
-
- let converter = PrimitiveArrayConverter::<arrow::datatypes::Int32Type>::new();
- let mut array_reader =
- ArrowArrayReader::try_new(page_iterator, column_desc, converter, None)
- .unwrap();
-
- // Read first 50 values, which are all from the first column chunk
- let array = array_reader.next_batch(50).unwrap();
- let array = array
- .as_any()
- .downcast_ref::<PrimitiveArray<ArrowInt32>>()
- .unwrap();
-
- assert_eq!(
- &PrimitiveArray::<ArrowInt32>::from(data[0..50].to_vec()),
- array
- );
-
- // Read next 100 values, the first 50 ones are from the first column chunk,
- // and the last 50 ones are from the second column chunk
- let array = array_reader.next_batch(100).unwrap();
- let array = array
- .as_any()
- .downcast_ref::<PrimitiveArray<ArrowInt32>>()
- .unwrap();
-
- assert_eq!(
- &PrimitiveArray::<ArrowInt32>::from(data[50..150].to_vec()),
- array
- );
-
- // Try to read 100 values, however there are only 50 values
- let array = array_reader.next_batch(100).unwrap();
- let array = array
- .as_any()
- .downcast_ref::<PrimitiveArray<ArrowInt32>>()
- .unwrap();
-
- assert_eq!(
- &PrimitiveArray::<ArrowInt32>::from(data[150..200].to_vec()),
- array
- );
- }
- }
-
- #[test]
- fn test_primitive_array_reader_def_and_rep_levels() {
- // Construct column schema
- let message_type = "
- message test_schema {
- REPEATED Group test_mid {
- OPTIONAL INT32 leaf;
- }
- }
- ";
-
- let schema = parse_message_type(message_type)
- .map(|t| Arc::new(SchemaDescriptor::new(Arc::new(t))))
- .unwrap();
-
- let column_desc = schema.column(0);
-
- // Construct page iterator
- {
- let mut def_levels = Vec::new();
- let mut rep_levels = Vec::new();
- let mut page_lists = Vec::new();
- make_column_chunks::<crate::data_type::Int32Type>(
- column_desc.clone(),
- Encoding::PLAIN,
- 100,
- 1,
- 200,
- &mut def_levels,
- &mut rep_levels,
- &mut Vec::new(),
- &mut page_lists,
- true,
- 2,
- );
-
- let page_iterator =
- InMemoryPageIterator::new(schema, column_desc.clone(), page_lists);
-
- let converter = PrimitiveArrayConverter::<arrow::datatypes::Int32Type>::new();
- let mut array_reader =
- ArrowArrayReader::try_new(page_iterator, column_desc, converter, None)
- .unwrap();
-
- let mut accu_len: usize = 0;
-
- // Read first 50 values, which are all from the first column chunk
- let array = array_reader.next_batch(50).unwrap();
- assert_eq!(
- Some(&def_levels[accu_len..(accu_len + array.len())]),
- array_reader.get_def_levels()
- );
- assert_eq!(
- Some(&rep_levels[accu_len..(accu_len + array.len())]),
- array_reader.get_rep_levels()
- );
- accu_len += array.len();
-
- // Read next 100 values, the first 50 ones are from the first column chunk,
- // and the last 50 ones are from the second column chunk
- let array = array_reader.next_batch(100).unwrap();
- assert_eq!(
- Some(&def_levels[accu_len..(accu_len + array.len())]),
- array_reader.get_def_levels()
- );
- assert_eq!(
- Some(&rep_levels[accu_len..(accu_len + array.len())]),
- array_reader.get_rep_levels()
- );
- accu_len += array.len();
-
- // Try to read 100 values, however there are only 50 values
- let array = array_reader.next_batch(100).unwrap();
- assert_eq!(
- Some(&def_levels[accu_len..(accu_len + array.len())]),
- array_reader.get_def_levels()
- );
- assert_eq!(
- Some(&rep_levels[accu_len..(accu_len + array.len())]),
- array_reader.get_rep_levels()
- );
-
- assert_eq!(accu_len + array.len(), 200);
- }
- }
-
- #[test]
- fn test_arrow_array_reader_string() {
- // Construct column schema
- let message_type = "
- message test_schema {
- REPEATED Group test_mid {
- OPTIONAL BYTE_ARRAY leaf (UTF8);
- }
- }
- ";
- let num_pages = 2;
- let values_per_page = 100;
- let str_base = "Hello World";
-
- let schema = parse_message_type(message_type)
- .map(|t| Arc::new(SchemaDescriptor::new(Arc::new(t))))
- .unwrap();
- let column_desc = schema.column(0);
- let max_def_level = column_desc.max_def_level();
- let max_rep_level = column_desc.max_rep_level();
-
- assert_eq!(max_def_level, 2);
- assert_eq!(max_rep_level, 1);
-
- let mut rng = thread_rng();
- let mut pages: Vec<Vec<Page>> = Vec::new();
-
- let mut rep_levels = Vec::with_capacity(num_pages * values_per_page);
- let mut def_levels = Vec::with_capacity(num_pages * values_per_page);
- let mut all_values = Vec::with_capacity(num_pages * values_per_page);
-
- for i in 0..num_pages {
- let mut values = Vec::with_capacity(values_per_page);
-
- for _ in 0..values_per_page {
- let def_level = rng.gen_range(0..max_def_level + 1);
- let rep_level = rng.gen_range(0..max_rep_level + 1);
- if def_level == max_def_level {
- let len = rng.gen_range(1..str_base.len());
- let slice = &str_base[..len];
- values.push(ByteArray::from(slice));
- all_values.push(Some(slice.to_string()));
- } else {
- all_values.push(None)
- }
- rep_levels.push(rep_level);
- def_levels.push(def_level)
- }
-
- let range = i * values_per_page..(i + 1) * values_per_page;
- let mut pb =
- DataPageBuilderImpl::new(column_desc.clone(), values.len() as u32, true);
-
- pb.add_rep_levels(max_rep_level, &rep_levels.as_slice()[range.clone()]);
- pb.add_def_levels(max_def_level, &def_levels.as_slice()[range]);
- pb.add_values::<ByteArrayType>(Encoding::PLAIN, values.as_slice());
-
- let data_page = pb.consume();
- pages.push(vec![data_page]);
- }
-
- let page_iterator = InMemoryPageIterator::new(schema, column_desc.clone(), pages);
- let converter = StringArrayConverter::new();
- let mut array_reader =
- ArrowArrayReader::try_new(page_iterator, column_desc, converter, None)
- .unwrap();
-
- let mut accu_len: usize = 0;
-
- let array = array_reader.next_batch(values_per_page / 2).unwrap();
- assert_eq!(array.len(), values_per_page / 2);
- assert_eq!(
- Some(&def_levels[accu_len..(accu_len + array.len())]),
- array_reader.get_def_levels()
- );
- assert_eq!(
- Some(&rep_levels[accu_len..(accu_len + array.len())]),
- array_reader.get_rep_levels()
- );
- accu_len += array.len();
-
- // Read next values_per_page values, the first values_per_page/2 ones are from the first column chunk,
- // and the last values_per_page/2 ones are from the second column chunk
- let array = array_reader.next_batch(values_per_page).unwrap();
- assert_eq!(array.len(), values_per_page);
- assert_eq!(
- Some(&def_levels[accu_len..(accu_len + array.len())]),
- array_reader.get_def_levels()
- );
- assert_eq!(
- Some(&rep_levels[accu_len..(accu_len + array.len())]),
- array_reader.get_rep_levels()
- );
- let strings = array.as_any().downcast_ref::<StringArray>().unwrap();
- for i in 0..array.len() {
- if array.is_valid(i) {
- assert_eq!(
- all_values[i + accu_len].as_ref().unwrap().as_str(),
- strings.value(i)
- )
- } else {
- assert_eq!(all_values[i + accu_len], None)
- }
- }
- accu_len += array.len();
-
- // Try to read values_per_page values, however there are only values_per_page/2 values
- let array = array_reader.next_batch(values_per_page).unwrap();
- assert_eq!(array.len(), values_per_page / 2);
- assert_eq!(
- Some(&def_levels[accu_len..(accu_len + array.len())]),
- array_reader.get_def_levels()
- );
- assert_eq!(
- Some(&rep_levels[accu_len..(accu_len + array.len())]),
- array_reader.get_rep_levels()
- );
- }
-}
diff --git a/parquet/src/arrow/arrow_reader.rs b/parquet/src/arrow/arrow_reader.rs
deleted file mode 100644
index 83fb0a2..0000000
--- a/parquet/src/arrow/arrow_reader.rs
+++ /dev/null
@@ -1,671 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-//! Contains reader which reads parquet data into arrow array.
-
-use crate::arrow::array_reader::{build_array_reader, ArrayReader, StructArrayReader};
-use crate::arrow::schema::parquet_to_arrow_schema;
-use crate::arrow::schema::{
- parquet_to_arrow_schema_by_columns, parquet_to_arrow_schema_by_root_columns,
-};
-use crate::errors::{ParquetError, Result};
-use crate::file::metadata::ParquetMetaData;
-use crate::file::reader::FileReader;
-use arrow::datatypes::{DataType as ArrowType, Schema, SchemaRef};
-use arrow::error::Result as ArrowResult;
-use arrow::record_batch::{RecordBatch, RecordBatchReader};
-use arrow::{array::StructArray, error::ArrowError};
-use std::sync::Arc;
-
-/// Arrow reader api.
-/// With this api, user can get arrow schema from parquet file, and read parquet data
-/// into arrow arrays.
-pub trait ArrowReader {
- type RecordReader: RecordBatchReader;
-
- /// Read parquet schema and convert it into arrow schema.
- fn get_schema(&mut self) -> Result<Schema>;
-
- /// Read parquet schema and convert it into arrow schema.
- /// This schema only includes columns identified by `column_indices`.
- /// To select leaf columns (i.e. `a.b.c` instead of `a`), set `leaf_columns = true`
- fn get_schema_by_columns<T>(
- &mut self,
- column_indices: T,
- leaf_columns: bool,
- ) -> Result<Schema>
- where
- T: IntoIterator<Item = usize>;
-
- /// Returns record batch reader from whole parquet file.
- ///
- /// # Arguments
- ///
- /// `batch_size`: The size of each record batch returned from this reader. Only the
- /// last batch may contain records less than this size, otherwise record batches
- /// returned from this reader should contains exactly `batch_size` elements.
- fn get_record_reader(&mut self, batch_size: usize) -> Result<Self::RecordReader>;
-
- /// Returns record batch reader whose record batch contains columns identified by
- /// `column_indices`.
- ///
- /// # Arguments
- ///
- /// `column_indices`: The columns that should be included in record batches.
- /// `batch_size`: Please refer to `get_record_reader`.
- fn get_record_reader_by_columns<T>(
- &mut self,
- column_indices: T,
- batch_size: usize,
- ) -> Result<Self::RecordReader>
- where
- T: IntoIterator<Item = usize>;
-}
-
-pub struct ParquetFileArrowReader {
- file_reader: Arc<dyn FileReader>,
-}
-
-impl ArrowReader for ParquetFileArrowReader {
- type RecordReader = ParquetRecordBatchReader;
-
- fn get_schema(&mut self) -> Result<Schema> {
- let file_metadata = self.file_reader.metadata().file_metadata();
- parquet_to_arrow_schema(
- file_metadata.schema_descr(),
- file_metadata.key_value_metadata(),
- )
- }
-
- fn get_schema_by_columns<T>(
- &mut self,
- column_indices: T,
- leaf_columns: bool,
- ) -> Result<Schema>
- where
- T: IntoIterator<Item = usize>,
- {
- let file_metadata = self.file_reader.metadata().file_metadata();
- if leaf_columns {
- parquet_to_arrow_schema_by_columns(
- file_metadata.schema_descr(),
- column_indices,
- file_metadata.key_value_metadata(),
- )
- } else {
- parquet_to_arrow_schema_by_root_columns(
- file_metadata.schema_descr(),
- column_indices,
- file_metadata.key_value_metadata(),
- )
- }
- }
-
- fn get_record_reader(
- &mut self,
- batch_size: usize,
- ) -> Result<ParquetRecordBatchReader> {
- let column_indices = 0..self
- .file_reader
- .metadata()
- .file_metadata()
- .schema_descr()
- .num_columns();
-
- self.get_record_reader_by_columns(column_indices, batch_size)
- }
-
- fn get_record_reader_by_columns<T>(
- &mut self,
- column_indices: T,
- batch_size: usize,
- ) -> Result<ParquetRecordBatchReader>
- where
- T: IntoIterator<Item = usize>,
- {
- let array_reader = build_array_reader(
- self.file_reader
- .metadata()
- .file_metadata()
- .schema_descr_ptr(),
- self.get_schema()?,
- column_indices,
- self.file_reader.clone(),
- )?;
-
- ParquetRecordBatchReader::try_new(batch_size, array_reader)
- }
-}
-
-impl ParquetFileArrowReader {
- pub fn new(file_reader: Arc<dyn FileReader>) -> Self {
- Self { file_reader }
- }
-
- // Expose the reader metadata
- pub fn get_metadata(&mut self) -> ParquetMetaData {
- self.file_reader.metadata().clone()
- }
-}
-
-pub struct ParquetRecordBatchReader {
- batch_size: usize,
- array_reader: Box<dyn ArrayReader>,
- schema: SchemaRef,
-}
-
-impl Iterator for ParquetRecordBatchReader {
- type Item = ArrowResult<RecordBatch>;
-
- fn next(&mut self) -> Option<Self::Item> {
- match self.array_reader.next_batch(self.batch_size) {
- Err(error) => Some(Err(error.into())),
- Ok(array) => {
- let struct_array =
- array.as_any().downcast_ref::<StructArray>().ok_or_else(|| {
- ArrowError::ParquetError(
- "Struct array reader should return struct array".to_string(),
- )
- });
- match struct_array {
- Err(err) => Some(Err(err)),
- Ok(e) => {
- match RecordBatch::try_new(self.schema.clone(), e.columns_ref()) {
- Err(err) => Some(Err(err)),
- Ok(record_batch) => {
- if record_batch.num_rows() > 0 {
- Some(Ok(record_batch))
- } else {
- None
- }
- }
- }
- }
- }
- }
- }
- }
-}
-
-impl RecordBatchReader for ParquetRecordBatchReader {
- fn schema(&self) -> SchemaRef {
- self.schema.clone()
- }
-}
-
-impl ParquetRecordBatchReader {
- pub fn try_new(
- batch_size: usize,
- array_reader: Box<dyn ArrayReader>,
- ) -> Result<Self> {
- // Check that array reader is struct array reader
- array_reader
- .as_any()
- .downcast_ref::<StructArrayReader>()
- .ok_or_else(|| general_err!("The input must be struct array reader!"))?;
-
- let schema = match array_reader.get_data_type() {
- ArrowType::Struct(ref fields) => Schema::new(fields.clone()),
- _ => unreachable!("Struct array reader's data type is not struct!"),
- };
-
- Ok(Self {
- batch_size,
- array_reader,
- schema: Arc::new(schema),
- })
- }
-}
-
-#[cfg(test)]
-mod tests {
- use crate::arrow::arrow_reader::{ArrowReader, ParquetFileArrowReader};
- use crate::arrow::converter::{
- Converter, FixedSizeArrayConverter, FromConverter, IntervalDayTimeArrayConverter,
- Utf8ArrayConverter,
- };
- use crate::column::writer::get_typed_column_writer_mut;
- use crate::data_type::{
- BoolType, ByteArray, ByteArrayType, DataType, FixedLenByteArray,
- FixedLenByteArrayType, Int32Type,
- };
- use crate::errors::Result;
- use crate::file::properties::WriterProperties;
- use crate::file::reader::{FileReader, SerializedFileReader};
- use crate::file::writer::{FileWriter, SerializedFileWriter};
- use crate::schema::parser::parse_message_type;
- use crate::schema::types::TypePtr;
- use crate::util::test_common::{get_temp_filename, RandGen};
- use arrow::array::*;
- use arrow::record_batch::RecordBatchReader;
- use rand::RngCore;
- use serde_json::json;
- use serde_json::Value::{Array as JArray, Null as JNull, Object as JObject};
- use std::cmp::min;
- use std::convert::TryFrom;
- use std::fs::File;
- use std::path::{Path, PathBuf};
- use std::sync::Arc;
-
- #[test]
- fn test_arrow_reader_all_columns() {
- let json_values = get_json_array("parquet/generated_simple_numerics/blogs.json");
-
- let parquet_file_reader =
- get_test_reader("parquet/generated_simple_numerics/blogs.parquet");
-
- let max_len = parquet_file_reader.metadata().file_metadata().num_rows() as usize;
-
- let mut arrow_reader = ParquetFileArrowReader::new(parquet_file_reader);
-
- let mut record_batch_reader = arrow_reader
- .get_record_reader(60)
- .expect("Failed to read into array!");
-
- // Verify that the schema was correctly parsed
- let original_schema = arrow_reader.get_schema().unwrap().fields().clone();
- assert_eq!(original_schema, *record_batch_reader.schema().fields());
-
- compare_batch_json(&mut record_batch_reader, json_values, max_len);
- }
-
- #[test]
- fn test_arrow_reader_single_column() {
- let json_values = get_json_array("parquet/generated_simple_numerics/blogs.json");
-
- let projected_json_values = json_values
- .into_iter()
- .map(|value| match value {
- JObject(fields) => {
- json!({ "blog_id": fields.get("blog_id").unwrap_or(&JNull).clone()})
- }
- _ => panic!("Input should be json object array!"),
- })
- .collect::<Vec<_>>();
-
- let parquet_file_reader =
- get_test_reader("parquet/generated_simple_numerics/blogs.parquet");
-
- let max_len = parquet_file_reader.metadata().file_metadata().num_rows() as usize;
-
- let mut arrow_reader = ParquetFileArrowReader::new(parquet_file_reader);
-
- let mut record_batch_reader = arrow_reader
- .get_record_reader_by_columns(vec![2], 60)
- .expect("Failed to read into array!");
-
- // Verify that the schema was correctly parsed
- let original_schema = arrow_reader.get_schema().unwrap().fields().clone();
- assert_eq!(1, record_batch_reader.schema().fields().len());
- assert_eq!(original_schema[1], record_batch_reader.schema().fields()[0]);
-
- compare_batch_json(&mut record_batch_reader, projected_json_values, max_len);
- }
-
- #[test]
- fn test_bool_single_column_reader_test() {
- let message_type = "
- message test_schema {
- REQUIRED BOOLEAN leaf;
- }
- ";
-
- let converter = FromConverter::new();
- run_single_column_reader_tests::<
- BoolType,
- BooleanArray,
- FromConverter<Vec<Option<bool>>, BooleanArray>,
- BoolType,
- >(2, message_type, &converter);
- }
-
- struct RandFixedLenGen {}
-
- impl RandGen<FixedLenByteArrayType> for RandFixedLenGen {
- fn gen(len: i32) -> FixedLenByteArray {
- let mut v = vec![0u8; len as usize];
- rand::thread_rng().fill_bytes(&mut v);
- ByteArray::from(v).into()
- }
- }
-
- #[test]
- fn test_fixed_length_binary_column_reader() {
- let message_type = "
- message test_schema {
- REQUIRED FIXED_LEN_BYTE_ARRAY (20) leaf;
- }
- ";
-
- let converter = FixedSizeArrayConverter::new(20);
- run_single_column_reader_tests::<
- FixedLenByteArrayType,
- FixedSizeBinaryArray,
- FixedSizeArrayConverter,
- RandFixedLenGen,
- >(20, message_type, &converter);
- }
-
- #[test]
- fn test_interval_day_time_column_reader() {
- let message_type = "
- message test_schema {
- REQUIRED FIXED_LEN_BYTE_ARRAY (12) leaf (INTERVAL);
- }
- ";
-
- let converter = IntervalDayTimeArrayConverter {};
- run_single_column_reader_tests::<
- FixedLenByteArrayType,
- IntervalDayTimeArray,
- IntervalDayTimeArrayConverter,
- RandFixedLenGen,
- >(12, message_type, &converter);
- }
-
- struct RandUtf8Gen {}
-
- impl RandGen<ByteArrayType> for RandUtf8Gen {
- fn gen(len: i32) -> ByteArray {
- Int32Type::gen(len).to_string().as_str().into()
- }
- }
-
- #[test]
- fn test_utf8_single_column_reader_test() {
- let message_type = "
- message test_schema {
- REQUIRED BINARY leaf (UTF8);
- }
- ";
-
- let converter = Utf8ArrayConverter {};
- run_single_column_reader_tests::<
- ByteArrayType,
- StringArray,
- Utf8ArrayConverter,
- RandUtf8Gen,
- >(2, message_type, &converter);
- }
-
- #[test]
- fn test_read_decimal_file() {
- use arrow::array::DecimalArray;
- let testdata = arrow::util::test_util::parquet_test_data();
- let file_variants = vec![("fixed_length", 25), ("int32", 4), ("int64", 10)];
- for (prefix, target_precision) in file_variants {
- let path = format!("{}/{}_decimal.parquet", testdata, prefix);
- let parquet_reader =
- SerializedFileReader::try_from(File::open(&path).unwrap()).unwrap();
- let mut arrow_reader = ParquetFileArrowReader::new(Arc::new(parquet_reader));
-
- let mut record_reader = arrow_reader.get_record_reader(32).unwrap();
-
- let batch = record_reader.next().unwrap().unwrap();
- assert_eq!(batch.num_rows(), 24);
- let col = batch
- .column(0)
- .as_any()
- .downcast_ref::<DecimalArray>()
- .unwrap();
-
- let expected = 1..25;
-
- assert_eq!(col.precision(), target_precision);
- assert_eq!(col.scale(), 2);
-
- for (i, v) in expected.enumerate() {
- assert_eq!(col.value(i), v * 100_i128);
- }
- }
- }
-
- /// Parameters for single_column_reader_test
- #[derive(Debug)]
- struct TestOptions {
- /// Number of row group to write to parquet (row group size =
- /// num_row_groups / num_rows)
- num_row_groups: usize,
- /// Total number of rows
- num_rows: usize,
- /// Size of batches to read back
- record_batch_size: usize,
- /// Total number of batches to attempt to read.
- /// `record_batch_size` * `num_iterations` should be greater
- /// than `num_rows` to ensure the data can be read back completely
- num_iterations: usize,
- }
-
- /// Create a parquet file and then read it using
- /// `ParquetFileArrowReader` using a standard set of parameters
- /// `opts`.
- ///
- /// `rand_max` represents the maximum size of value to pass to to
- /// value generator
- fn run_single_column_reader_tests<T, A, C, G>(
- rand_max: i32,
- message_type: &str,
- converter: &C,
- ) where
- T: DataType,
- G: RandGen<T>,
- A: PartialEq + Array + 'static,
- C: Converter<Vec<Option<T::T>>, A> + 'static,
- {
- let all_options = vec![
- // choose record_batch_batch (15) so batches cross row
- // group boundaries (50 rows in 2 row groups) cases.
- TestOptions {
- num_row_groups: 2,
- num_rows: 100,
- record_batch_size: 15,
- num_iterations: 50,
- },
- // choose record_batch_batch (5) so batches sometime fall
- // on row group boundaries and (25 rows in 3 row groups
- // --> row groups of 10, 10, and 5). Tests buffer
- // refilling edge cases.
- TestOptions {
- num_row_groups: 3,
- num_rows: 25,
- record_batch_size: 5,
- num_iterations: 50,
- },
- // Choose record_batch_size (25) so all batches fall
- // exactly on row group boundary (25). Tests buffer
- // refilling edge cases.
- TestOptions {
- num_row_groups: 4,
- num_rows: 100,
- record_batch_size: 25,
- num_iterations: 50,
- },
- ];
-
- all_options.into_iter().for_each(|opts| {
- // Print out options to facilitate debugging failures on CI
- println!("Running with Test Options: {:?}", opts);
- single_column_reader_test::<T, A, C, G>(
- opts,
- rand_max,
- message_type,
- converter,
- )
- });
- }
-
- /// Create a parquet file and then read it using
- /// `ParquetFileArrowReader` using the parameters described in
- /// `opts`.
- fn single_column_reader_test<T, A, C, G>(
- opts: TestOptions,
- rand_max: i32,
- message_type: &str,
- converter: &C,
- ) where
- T: DataType,
- G: RandGen<T>,
- A: PartialEq + Array + 'static,
- C: Converter<Vec<Option<T::T>>, A> + 'static,
- {
- let values: Vec<Vec<T::T>> = (0..opts.num_row_groups)
- .map(|_| G::gen_vec(rand_max, opts.num_rows))
- .collect();
-
- let path = get_temp_filename();
-
- let schema = parse_message_type(message_type).map(Arc::new).unwrap();
-
- generate_single_column_file_with_data::<T>(&values, path.as_path(), schema)
- .unwrap();
-
- let parquet_reader =
- SerializedFileReader::try_from(File::open(&path).unwrap()).unwrap();
- let mut arrow_reader = ParquetFileArrowReader::new(Arc::new(parquet_reader));
-
- let mut record_reader = arrow_reader
- .get_record_reader(opts.record_batch_size)
- .unwrap();
-
- let expected_data: Vec<Option<T::T>> = values
- .iter()
- .flat_map(|v| v.iter())
- .map(|b| Some(b.clone()))
- .collect();
-
- for i in 0..opts.num_iterations {
- let start = i * opts.record_batch_size;
-
- let batch = record_reader.next();
- if start < expected_data.len() {
- let end = min(start + opts.record_batch_size, expected_data.len());
- assert!(batch.is_some());
-
- let mut data = vec![];
- data.extend_from_slice(&expected_data[start..end]);
-
- assert_eq!(
- &converter.convert(data).unwrap(),
- batch
- .unwrap()
- .unwrap()
- .column(0)
- .as_any()
- .downcast_ref::<A>()
- .unwrap()
- );
- } else {
- assert!(batch.is_none());
- }
- }
- }
-
- fn generate_single_column_file_with_data<T: DataType>(
- values: &[Vec<T::T>],
- path: &Path,
- schema: TypePtr,
- ) -> Result<parquet_format::FileMetaData> {
- let file = File::create(path)?;
- let writer_props = Arc::new(WriterProperties::builder().build());
-
- let mut writer = SerializedFileWriter::new(file, schema, writer_props)?;
-
- for v in values {
- let mut row_group_writer = writer.next_row_group()?;
- let mut column_writer = row_group_writer
- .next_column()?
- .expect("Column writer is none!");
-
- get_typed_column_writer_mut::<T>(&mut column_writer)
- .write_batch(v, None, None)?;
-
- row_group_writer.close_column(column_writer)?;
- writer.close_row_group(row_group_writer)?
- }
-
- writer.close()
- }
-
- fn get_test_reader(file_name: &str) -> Arc<dyn FileReader> {
- let file = get_test_file(file_name);
-
- let reader =
- SerializedFileReader::new(file).expect("Failed to create serialized reader");
-
- Arc::new(reader)
- }
-
- fn get_test_file(file_name: &str) -> File {
- let mut path = PathBuf::new();
- path.push(arrow::util::test_util::arrow_test_data());
- path.push(file_name);
-
- File::open(path.as_path()).expect("File not found!")
- }
-
- fn get_json_array(filename: &str) -> Vec<serde_json::Value> {
- match serde_json::from_reader(get_test_file(filename))
- .expect("Failed to read json value from file!")
- {
- JArray(values) => values,
- _ => panic!("Input should be json array!"),
- }
- }
-
- fn compare_batch_json(
- record_batch_reader: &mut dyn RecordBatchReader,
- json_values: Vec<serde_json::Value>,
- max_len: usize,
- ) {
- for i in 0..20 {
- let array: Option<StructArray> = record_batch_reader
- .next()
- .map(|r| r.expect("Failed to read record batch!").into());
-
- let (start, end) = (i * 60_usize, (i + 1) * 60_usize);
-
- if start < max_len {
- assert!(array.is_some());
- assert_ne!(0, array.as_ref().unwrap().len());
- let end = min(end, max_len);
- let json = JArray(Vec::from(&json_values[start..end]));
- assert_eq!(array.unwrap(), json)
- } else {
- assert!(array.is_none());
- }
- }
- }
-
- #[test]
- fn test_read_structs() {
- // This particular test file has columns of struct types where there is
- // a column that has the same name as one of the struct fields
- // (see: ARROW-11452)
- let testdata = arrow::util::test_util::parquet_test_data();
- let path = format!("{}/nested_structs.rust.parquet", testdata);
- let parquet_file_reader =
- SerializedFileReader::try_from(File::open(&path).unwrap()).unwrap();
- let mut arrow_reader = ParquetFileArrowReader::new(Arc::new(parquet_file_reader));
- let record_batch_reader = arrow_reader
- .get_record_reader(60)
- .expect("Failed to read into array!");
-
- for batch in record_batch_reader {
- batch.unwrap();
- }
- }
-}
diff --git a/parquet/src/arrow/arrow_writer.rs b/parquet/src/arrow/arrow_writer.rs
deleted file mode 100644
index 69ebce6..0000000
--- a/parquet/src/arrow/arrow_writer.rs
+++ /dev/null
@@ -1,1678 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-//! Contains writer which writes arrow data into parquet data.
-
-use std::sync::Arc;
-
-use arrow::array as arrow_array;
-use arrow::datatypes::{DataType as ArrowDataType, IntervalUnit, SchemaRef};
-use arrow::record_batch::RecordBatch;
-use arrow_array::Array;
-
-use super::levels::LevelInfo;
-use super::schema::{
- add_encoded_arrow_schema_to_metadata, decimal_length_from_precision,
-};
-
-use crate::column::writer::ColumnWriter;
-use crate::errors::{ParquetError, Result};
-use crate::file::properties::WriterProperties;
-use crate::{
- data_type::*,
- file::writer::{FileWriter, ParquetWriter, RowGroupWriter, SerializedFileWriter},
-};
-
-/// Arrow writer
-///
-/// Writes Arrow `RecordBatch`es to a Parquet writer
-pub struct ArrowWriter<W: ParquetWriter> {
- /// Underlying Parquet writer
- writer: SerializedFileWriter<W>,
- /// A copy of the Arrow schema.
- ///
- /// The schema is used to verify that each record batch written has the correct schema
- arrow_schema: SchemaRef,
- /// The length of arrays to write to each row group
- max_row_group_size: usize,
-}
-
-impl<W: 'static + ParquetWriter> ArrowWriter<W> {
- /// Try to create a new Arrow writer
- ///
- /// The writer will fail if:
- /// * a `SerializedFileWriter` cannot be created from the ParquetWriter
- /// * the Arrow schema contains unsupported datatypes such as Unions
- pub fn try_new(
- writer: W,
- arrow_schema: SchemaRef,
- props: Option<WriterProperties>,
- ) -> Result<Self> {
- let schema = crate::arrow::arrow_to_parquet_schema(&arrow_schema)?;
- // add serialized arrow schema
- let mut props = props.unwrap_or_else(|| WriterProperties::builder().build());
- add_encoded_arrow_schema_to_metadata(&arrow_schema, &mut props);
-
- let max_row_group_size = props.max_row_group_size();
-
- let file_writer = SerializedFileWriter::new(
- writer.try_clone()?,
- schema.root_schema_ptr(),
- Arc::new(props),
- )?;
-
- Ok(Self {
- writer: file_writer,
- arrow_schema,
- max_row_group_size,
- })
- }
-
- /// Write a RecordBatch to writer
- ///
- /// The writer will slice the `batch` into `max_row_group_size`,
- /// but if a batch has left-over rows less than the row group size,
- /// the last row group will have fewer records.
- /// This is currently a limitation because we close the row group
- /// instead of keeping it open for the next batch.
- pub fn write(&mut self, batch: &RecordBatch) -> Result<()> {
- // validate batch schema against writer's supplied schema
- if self.arrow_schema != batch.schema() {
- return Err(ParquetError::ArrowError(
- "Record batch schema does not match writer schema".to_string(),
- ));
- }
- // Track the number of rows being written in the batch.
- // We currently do not have a way of slicing nested arrays, thus we
- // track this manually.
- let num_rows = batch.num_rows();
- let batches = (num_rows + self.max_row_group_size - 1) / self.max_row_group_size;
- let min_batch = num_rows.min(self.max_row_group_size);
- for batch_index in 0..batches {
- // Determine the offset and length of arrays
- let offset = batch_index * min_batch;
- let length = (num_rows - offset).min(self.max_row_group_size);
-
- // Compute the definition and repetition levels of the batch
- let batch_level = LevelInfo::new(offset, length);
- let mut row_group_writer = self.writer.next_row_group()?;
- for (array, field) in batch.columns().iter().zip(batch.schema().fields()) {
- let mut levels = batch_level.calculate_array_levels(array, field);
- // Reverse levels as we pop() them when writing arrays
- levels.reverse();
- write_leaves(&mut row_group_writer, array, &mut levels)?;
- }
-
- self.writer.close_row_group(row_group_writer)?;
- }
-
- Ok(())
- }
-
- /// Close and finalize the underlying Parquet writer
- pub fn close(&mut self) -> Result<parquet_format::FileMetaData> {
- self.writer.close()
- }
-}
-
-/// Convenience method to get the next ColumnWriter from the RowGroupWriter
-#[inline]
-#[allow(clippy::borrowed_box)]
-fn get_col_writer(
- row_group_writer: &mut Box<dyn RowGroupWriter>,
-) -> Result<ColumnWriter> {
- let col_writer = row_group_writer
- .next_column()?
- .expect("Unable to get column writer");
- Ok(col_writer)
-}
-
-#[allow(clippy::borrowed_box)]
-fn write_leaves(
- mut row_group_writer: &mut Box<dyn RowGroupWriter>,
- array: &arrow_array::ArrayRef,
- mut levels: &mut Vec<LevelInfo>,
-) -> Result<()> {
- match array.data_type() {
- ArrowDataType::Null
- | ArrowDataType::Boolean
- | ArrowDataType::Int8
- | ArrowDataType::Int16
- | ArrowDataType::Int32
- | ArrowDataType::Int64
- | ArrowDataType::UInt8
- | ArrowDataType::UInt16
- | ArrowDataType::UInt32
- | ArrowDataType::UInt64
- | ArrowDataType::Float32
- | ArrowDataType::Float64
- | ArrowDataType::Timestamp(_, _)
- | ArrowDataType::Date32
- | ArrowDataType::Date64
- | ArrowDataType::Time32(_)
- | ArrowDataType::Time64(_)
- | ArrowDataType::Duration(_)
- | ArrowDataType::Interval(_)
- | ArrowDataType::LargeBinary
- | ArrowDataType::Binary
- | ArrowDataType::Utf8
- | ArrowDataType::LargeUtf8
- | ArrowDataType::Decimal(_, _)
- | ArrowDataType::FixedSizeBinary(_) => {
- let mut col_writer = get_col_writer(&mut row_group_writer)?;
- write_leaf(
- &mut col_writer,
- array,
- levels.pop().expect("Levels exhausted"),
- )?;
- row_group_writer.close_column(col_writer)?;
- Ok(())
- }
- ArrowDataType::List(_) | ArrowDataType::LargeList(_) => {
- // write the child list
- let data = array.data();
- let child_array = arrow_array::make_array(data.child_data()[0].clone());
- write_leaves(&mut row_group_writer, &child_array, &mut levels)?;
- Ok(())
- }
- ArrowDataType::Struct(_) => {
- let struct_array: &arrow_array::StructArray = array
- .as_any()
- .downcast_ref::<arrow_array::StructArray>()
- .expect("Unable to get struct array");
- for field in struct_array.columns() {
- write_leaves(&mut row_group_writer, field, &mut levels)?;
- }
- Ok(())
- }
- ArrowDataType::Dictionary(_, value_type) => {
- // cast dictionary to a primitive
- let array = arrow::compute::cast(array, value_type)?;
-
- let mut col_writer = get_col_writer(&mut row_group_writer)?;
- write_leaf(
- &mut col_writer,
- &array,
- levels.pop().expect("Levels exhausted"),
- )?;
- row_group_writer.close_column(col_writer)?;
- Ok(())
- }
- ArrowDataType::Float16 => Err(ParquetError::ArrowError(
- "Float16 arrays not supported".to_string(),
- )),
- ArrowDataType::FixedSizeList(_, _) | ArrowDataType::Union(_) => {
- Err(ParquetError::NYI(
- format!(
- "Attempting to write an Arrow type {:?} to parquet that is not yet implemented",
- array.data_type()
- )
- ))
- }
- }
-}
-
-fn write_leaf(
- writer: &mut ColumnWriter,
- column: &arrow_array::ArrayRef,
- levels: LevelInfo,
-) -> Result<i64> {
- let indices = levels.filter_array_indices();
- // Slice array according to computed offset and length
- let column = column.slice(levels.offset, levels.length);
- let written = match writer {
- ColumnWriter::Int32ColumnWriter(ref mut typed) => {
- let values = match column.data_type() {
- ArrowDataType::Date64 => {
- // If the column is a Date64, we cast it to a Date32, and then interpret that as Int32
- let array = if let ArrowDataType::Date64 = column.data_type() {
- let array =
- arrow::compute::cast(&column, &ArrowDataType::Date32)?;
- arrow::compute::cast(&array, &ArrowDataType::Int32)?
- } else {
- arrow::compute::cast(&column, &ArrowDataType::Int32)?
- };
- let array = array
- .as_any()
- .downcast_ref::<arrow_array::Int32Array>()
- .expect("Unable to get int32 array");
- get_numeric_array_slice::<Int32Type, _>(&array, &indices)
- }
- ArrowDataType::UInt32 => {
- // follow C++ implementation and use overflow/reinterpret cast from u32 to i32 which will map
- // `(i32::MAX as u32)..u32::MAX` to `i32::MIN..0`
- let array = column
- .as_any()
- .downcast_ref::<arrow_array::UInt32Array>()
- .expect("Unable to get u32 array");
- let array = arrow::compute::unary::<_, _, arrow::datatypes::Int32Type>(
- array,
- |x| x as i32,
- );
- get_numeric_array_slice::<Int32Type, _>(&array, &indices)
- }
- _ => {
- let array = arrow::compute::cast(&column, &ArrowDataType::Int32)?;
- let array = array
- .as_any()
- .downcast_ref::<arrow_array::Int32Array>()
- .expect("Unable to get i32 array");
- get_numeric_array_slice::<Int32Type, _>(&array, &indices)
- }
- };
- typed.write_batch(
- values.as_slice(),
- Some(levels.definition.as_slice()),
- levels.repetition.as_deref(),
- )?
- }
- ColumnWriter::BoolColumnWriter(ref mut typed) => {
- let array = column
- .as_any()
- .downcast_ref::<arrow_array::BooleanArray>()
- .expect("Unable to get boolean array");
- typed.write_batch(
- get_bool_array_slice(&array, &indices).as_slice(),
- Some(levels.definition.as_slice()),
- levels.repetition.as_deref(),
- )?
- }
- ColumnWriter::Int64ColumnWriter(ref mut typed) => {
- let values = match column.data_type() {
- ArrowDataType::Int64 => {
- let array = column
- .as_any()
- .downcast_ref::<arrow_array::Int64Array>()
- .expect("Unable to get i64 array");
- get_numeric_array_slice::<Int64Type, _>(&array, &indices)
- }
- ArrowDataType::UInt64 => {
- // follow C++ implementation and use overflow/reinterpret cast from u64 to i64 which will map
- // `(i64::MAX as u64)..u64::MAX` to `i64::MIN..0`
- let array = column
- .as_any()
- .downcast_ref::<arrow_array::UInt64Array>()
- .expect("Unable to get u64 array");
- let array = arrow::compute::unary::<_, _, arrow::datatypes::Int64Type>(
- array,
- |x| x as i64,
- );
- get_numeric_array_slice::<Int64Type, _>(&array, &indices)
- }
- _ => {
- let array = arrow::compute::cast(&column, &ArrowDataType::Int64)?;
- let array = array
- .as_any()
- .downcast_ref::<arrow_array::Int64Array>()
- .expect("Unable to get i64 array");
- get_numeric_array_slice::<Int64Type, _>(&array, &indices)
- }
- };
- typed.write_batch(
- values.as_slice(),
- Some(levels.definition.as_slice()),
- levels.repetition.as_deref(),
- )?
- }
- ColumnWriter::Int96ColumnWriter(ref mut _typed) => {
- unreachable!("Currently unreachable because data type not supported")
- }
- ColumnWriter::FloatColumnWriter(ref mut typed) => {
- let array = column
- .as_any()
- .downcast_ref::<arrow_array::Float32Array>()
- .expect("Unable to get Float32 array");
- typed.write_batch(
- get_numeric_array_slice::<FloatType, _>(&array, &indices).as_slice(),
- Some(levels.definition.as_slice()),
- levels.repetition.as_deref(),
- )?
- }
- ColumnWriter::DoubleColumnWriter(ref mut typed) => {
- let array = column
- .as_any()
- .downcast_ref::<arrow_array::Float64Array>()
- .expect("Unable to get Float64 array");
- typed.write_batch(
- get_numeric_array_slice::<DoubleType, _>(&array, &indices).as_slice(),
- Some(levels.definition.as_slice()),
- levels.repetition.as_deref(),
- )?
- }
- ColumnWriter::ByteArrayColumnWriter(ref mut typed) => match column.data_type() {
- ArrowDataType::Binary => {
- let array = column
- .as_any()
- .downcast_ref::<arrow_array::BinaryArray>()
- .expect("Unable to get BinaryArray array");
- typed.write_batch(
- get_binary_array(&array).as_slice(),
- Some(levels.definition.as_slice()),
- levels.repetition.as_deref(),
- )?
- }
- ArrowDataType::Utf8 => {
- let array = column
- .as_any()
- .downcast_ref::<arrow_array::StringArray>()
- .expect("Unable to get LargeBinaryArray array");
- typed.write_batch(
- get_string_array(&array).as_slice(),
- Some(levels.definition.as_slice()),
- levels.repetition.as_deref(),
- )?
- }
- ArrowDataType::LargeBinary => {
- let array = column
- .as_any()
- .downcast_ref::<arrow_array::LargeBinaryArray>()
- .expect("Unable to get LargeBinaryArray array");
- typed.write_batch(
- get_large_binary_array(&array).as_slice(),
- Some(levels.definition.as_slice()),
- levels.repetition.as_deref(),
- )?
- }
- ArrowDataType::LargeUtf8 => {
- let array = column
- .as_any()
- .downcast_ref::<arrow_array::LargeStringArray>()
- .expect("Unable to get LargeUtf8 array");
- typed.write_batch(
- get_large_string_array(&array).as_slice(),
- Some(levels.definition.as_slice()),
- levels.repetition.as_deref(),
- )?
- }
- _ => unreachable!("Currently unreachable because data type not supported"),
- },
- ColumnWriter::FixedLenByteArrayColumnWriter(ref mut typed) => {
- let bytes = match column.data_type() {
- ArrowDataType::Interval(interval_unit) => match interval_unit {
- IntervalUnit::YearMonth => {
- let array = column
- .as_any()
- .downcast_ref::<arrow_array::IntervalYearMonthArray>()
- .unwrap();
- get_interval_ym_array_slice(&array, &indices)
- }
- IntervalUnit::DayTime => {
- let array = column
- .as_any()
- .downcast_ref::<arrow_array::IntervalDayTimeArray>()
- .unwrap();
- get_interval_dt_array_slice(&array, &indices)
- }
- },
- ArrowDataType::FixedSizeBinary(_) => {
- let array = column
- .as_any()
- .downcast_ref::<arrow_array::FixedSizeBinaryArray>()
- .unwrap();
- get_fsb_array_slice(&array, &indices)
- }
- ArrowDataType::Decimal(_, _) => {
- let array = column
- .as_any()
- .downcast_ref::<arrow_array::DecimalArray>()
- .unwrap();
- get_decimal_array_slice(&array, &indices)
- }
- _ => {
- return Err(ParquetError::NYI(
- "Attempting to write an Arrow type that is not yet implemented"
- .to_string(),
- ));
- }
- };
- typed.write_batch(
- bytes.as_slice(),
- Some(levels.definition.as_slice()),
- levels.repetition.as_deref(),
- )?
- }
- };
- Ok(written as i64)
-}
-
-macro_rules! def_get_binary_array_fn {
- ($name:ident, $ty:ty) => {
- fn $name(array: &$ty) -> Vec<ByteArray> {
- let mut values = Vec::with_capacity(array.len() - array.null_count());
- for i in 0..array.len() {
- if array.is_valid(i) {
- let bytes: Vec<u8> = array.value(i).into();
- let bytes = ByteArray::from(bytes);
- values.push(bytes);
- }
- }
- values
- }
- };
-}
-
-def_get_binary_array_fn!(get_binary_array, arrow_array::BinaryArray);
-def_get_binary_array_fn!(get_string_array, arrow_array::StringArray);
-def_get_binary_array_fn!(get_large_binary_array, arrow_array::LargeBinaryArray);
-def_get_binary_array_fn!(get_large_string_array, arrow_array::LargeStringArray);
-
-/// Get the underlying numeric array slice, skipping any null values.
-/// If there are no null values, it might be quicker to get the slice directly instead of
-/// calling this function.
-fn get_numeric_array_slice<T, A>(
- array: &arrow_array::PrimitiveArray<A>,
- indices: &[usize],
-) -> Vec<T::T>
-where
- T: DataType,
- A: arrow::datatypes::ArrowNumericType,
- T::T: From<A::Native>,
-{
- let mut values = Vec::with_capacity(indices.len());
- for i in indices {
- values.push(array.value(*i).into())
- }
- values
-}
-
-fn get_bool_array_slice(
- array: &arrow_array::BooleanArray,
- indices: &[usize],
-) -> Vec<bool> {
- let mut values = Vec::with_capacity(indices.len());
- for i in indices {
- values.push(array.value(*i))
- }
- values
-}
-
-/// Returns 12-byte values representing 3 values of months, days and milliseconds (4-bytes each).
-/// An Arrow YearMonth interval only stores months, thus only the first 4 bytes are populated.
-fn get_interval_ym_array_slice(
- array: &arrow_array::IntervalYearMonthArray,
- indices: &[usize],
-) -> Vec<FixedLenByteArray> {
- let mut values = Vec::with_capacity(indices.len());
- for i in indices {
- let mut value = array.value(*i).to_le_bytes().to_vec();
- let mut suffix = vec![0; 8];
- value.append(&mut suffix);
- values.push(FixedLenByteArray::from(ByteArray::from(value)))
- }
- values
-}
-
-/// Returns 12-byte values representing 3 values of months, days and milliseconds (4-bytes each).
-/// An Arrow DayTime interval only stores days and millis, thus the first 4 bytes are not populated.
-fn get_interval_dt_array_slice(
- array: &arrow_array::IntervalDayTimeArray,
- indices: &[usize],
-) -> Vec<FixedLenByteArray> {
- let mut values = Vec::with_capacity(indices.len());
- for i in indices {
- let mut prefix = vec![0; 4];
- let mut value = array.value(*i).to_le_bytes().to_vec();
- prefix.append(&mut value);
- debug_assert_eq!(prefix.len(), 12);
- values.push(FixedLenByteArray::from(ByteArray::from(prefix)));
- }
- values
-}
-
-fn get_decimal_array_slice(
- array: &arrow_array::DecimalArray,
- indices: &[usize],
-) -> Vec<FixedLenByteArray> {
- let mut values = Vec::with_capacity(indices.len());
- let size = decimal_length_from_precision(array.precision());
- for i in indices {
- let as_be_bytes = array.value(*i).to_be_bytes();
- let resized_value = as_be_bytes[(16 - size)..].to_vec();
- values.push(FixedLenByteArray::from(ByteArray::from(resized_value)));
- }
- values
-}
-
-fn get_fsb_array_slice(
- array: &arrow_array::FixedSizeBinaryArray,
- indices: &[usize],
-) -> Vec<FixedLenByteArray> {
- let mut values = Vec::with_capacity(indices.len());
- for i in indices {
- let value = array.value(*i).to_vec();
- values.push(FixedLenByteArray::from(ByteArray::from(value)))
- }
- values
-}
-
-#[cfg(test)]
-mod tests {
- use super::*;
-
- use std::sync::Arc;
- use std::{fs::File, io::Seek};
-
- use arrow::datatypes::ToByteSlice;
- use arrow::datatypes::{DataType, Field, Schema, UInt32Type, UInt8Type};
- use arrow::record_batch::RecordBatch;
- use arrow::{array::*, buffer::Buffer};
-
- use crate::arrow::{ArrowReader, ParquetFileArrowReader};
- use crate::file::{
- reader::{FileReader, SerializedFileReader},
- statistics::Statistics,
- writer::InMemoryWriteableCursor,
- };
- use crate::util::test_common::get_temp_file;
-
- #[test]
- fn arrow_writer() {
- // define schema
- let schema = Schema::new(vec![
- Field::new("a", DataType::Int32, false),
- Field::new("b", DataType::Int32, true),
- ]);
-
- // create some data
- let a = Int32Array::from(vec![1, 2, 3, 4, 5]);
- let b = Int32Array::from(vec![Some(1), None, None, Some(4), Some(5)]);
-
- // build a record batch
- let batch = RecordBatch::try_new(
- Arc::new(schema.clone()),
- vec![Arc::new(a), Arc::new(b)],
- )
- .unwrap();
-
- let file = get_temp_file("test_arrow_writer.parquet", &[]);
- let mut writer = ArrowWriter::try_new(file, Arc::new(schema), None).unwrap();
- writer.write(&batch).unwrap();
- writer.close().unwrap();
- }
-
- #[test]
- fn roundtrip_bytes() {
- // define schema
- let schema = Arc::new(Schema::new(vec![
- Field::new("a", DataType::Int32, false),
- Field::new("b", DataType::Int32, true),
- ]));
-
- // create some data
- let a = Int32Array::from(vec![1, 2, 3, 4, 5]);
- let b = Int32Array::from(vec![Some(1), None, None, Some(4), Some(5)]);
-
- // build a record batch
- let expected_batch =
- RecordBatch::try_new(schema.clone(), vec![Arc::new(a), Arc::new(b)]).unwrap();
-
- let cursor = InMemoryWriteableCursor::default();
-
- {
- let mut writer = ArrowWriter::try_new(cursor.clone(), schema, None).unwrap();
- writer.write(&expected_batch).unwrap();
- writer.close().unwrap();
- }
-
- let buffer = cursor.into_inner().unwrap();
-
- let cursor = crate::file::serialized_reader::SliceableCursor::new(buffer);
- let reader = SerializedFileReader::new(cursor).unwrap();
- let mut arrow_reader = ParquetFileArrowReader::new(Arc::new(reader));
- let mut record_batch_reader = arrow_reader.get_record_reader(1024).unwrap();
-
- let actual_batch = record_batch_reader
- .next()
- .expect("No batch found")
- .expect("Unable to get batch");
-
- assert_eq!(expected_batch.schema(), actual_batch.schema());
- assert_eq!(expected_batch.num_columns(), actual_batch.num_columns());
- assert_eq!(expected_batch.num_rows(), actual_batch.num_rows());
- for i in 0..expected_batch.num_columns() {
- let expected_data = expected_batch.column(i).data().clone();
- let actual_data = actual_batch.column(i).data().clone();
-
- assert_eq!(expected_data, actual_data);
- }
- }
-
- #[test]
- fn arrow_writer_non_null() {
- // define schema
- let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]);
-
- // create some data
- let a = Int32Array::from(vec![1, 2, 3, 4, 5]);
-
- // build a record batch
- let batch =
- RecordBatch::try_new(Arc::new(schema.clone()), vec![Arc::new(a)]).unwrap();
-
- let file = get_temp_file("test_arrow_writer_non_null.parquet", &[]);
- let mut writer = ArrowWriter::try_new(file, Arc::new(schema), None).unwrap();
- writer.write(&batch).unwrap();
- writer.close().unwrap();
- }
-
- #[test]
- fn arrow_writer_list() {
- // define schema
- let schema = Schema::new(vec![Field::new(
- "a",
- DataType::List(Box::new(Field::new("item", DataType::Int32, true))),
- false,
- )]);
-
- // create some data
- let a_values = Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10]);
-
- // Construct a buffer for value offsets, for the nested array:
- // [[1], [2, 3], null, [4, 5, 6], [7, 8, 9, 10]]
- let a_value_offsets =
- arrow::buffer::Buffer::from(&[0, 1, 3, 3, 6, 10].to_byte_slice());
-
- // Construct a list array from the above two
- let a_list_data = ArrayData::builder(DataType::List(Box::new(Field::new(
- "item",
- DataType::Int32,
- true,
- ))))
- .len(5)
- .add_buffer(a_value_offsets)
- .add_child_data(a_values.data().clone())
- .null_bit_buffer(Buffer::from(vec![0b00011011]))
- .build();
- let a = ListArray::from(a_list_data);
-
- // build a record batch
- let batch =
- RecordBatch::try_new(Arc::new(schema.clone()), vec![Arc::new(a)]).unwrap();
-
- // I think this setup is incorrect because this should pass
- assert_eq!(batch.column(0).data().null_count(), 1);
-
- let file = get_temp_file("test_arrow_writer_list.parquet", &[]);
- let mut writer = ArrowWriter::try_new(file, Arc::new(schema), None).unwrap();
- writer.write(&batch).unwrap();
- writer.close().unwrap();
- }
-
- #[test]
- fn arrow_writer_binary() {
- let string_field = Field::new("a", DataType::Utf8, false);
- let binary_field = Field::new("b", DataType::Binary, false);
- let schema = Schema::new(vec![string_field, binary_field]);
-
- let raw_string_values = vec!["foo", "bar", "baz", "quux"];
- let raw_binary_values = vec![
- b"foo".to_vec(),
- b"bar".to_vec(),
- b"baz".to_vec(),
- b"quux".to_vec(),
- ];
- let raw_binary_value_refs = raw_binary_values
- .iter()
- .map(|x| x.as_slice())
- .collect::<Vec<_>>();
-
- let string_values = StringArray::from(raw_string_values.clone());
- let binary_values = BinaryArray::from(raw_binary_value_refs);
- let batch = RecordBatch::try_new(
- Arc::new(schema.clone()),
- vec![Arc::new(string_values), Arc::new(binary_values)],
- )
- .unwrap();
-
- let mut file = get_temp_file("test_arrow_writer_binary.parquet", &[]);
- let mut writer =
- ArrowWriter::try_new(file.try_clone().unwrap(), Arc::new(schema), None)
- .unwrap();
- writer.write(&batch).unwrap();
- writer.close().unwrap();
-
- file.seek(std::io::SeekFrom::Start(0)).unwrap();
- let file_reader = SerializedFileReader::new(file).unwrap();
- let mut arrow_reader = ParquetFileArrowReader::new(Arc::new(file_reader));
- let mut record_batch_reader = arrow_reader.get_record_reader(1024).unwrap();
-
- let batch = record_batch_reader.next().unwrap().unwrap();
- let string_col = batch
- .column(0)
- .as_any()
- .downcast_ref::<StringArray>()
- .unwrap();
- let binary_col = batch
- .column(1)
- .as_any()
- .downcast_ref::<BinaryArray>()
- .unwrap();
-
- for i in 0..batch.num_rows() {
- assert_eq!(string_col.value(i), raw_string_values[i]);
- assert_eq!(binary_col.value(i), raw_binary_values[i].as_slice());
- }
- }
-
- #[test]
- fn arrow_writer_decimal() {
- let decimal_field = Field::new("a", DataType::Decimal(5, 2), false);
- let schema = Schema::new(vec![decimal_field]);
-
- let mut dec_builder = DecimalBuilder::new(4, 5, 2);
- dec_builder.append_value(10_000).unwrap();
- dec_builder.append_value(50_000).unwrap();
- dec_builder.append_value(0).unwrap();
- dec_builder.append_value(-100).unwrap();
-
- let raw_decimal_i128_values: Vec<i128> = vec![10_000, 50_000, 0, -100];
- let decimal_values = dec_builder.finish();
- let batch = RecordBatch::try_new(
- Arc::new(schema.clone()),
- vec![Arc::new(decimal_values)],
- )
- .unwrap();
-
- let mut file = get_temp_file("test_arrow_writer_decimal.parquet", &[]);
- let mut writer =
- ArrowWriter::try_new(file.try_clone().unwrap(), Arc::new(schema), None)
- .unwrap();
- writer.write(&batch).unwrap();
- writer.close().unwrap();
-
- file.seek(std::io::SeekFrom::Start(0)).unwrap();
- let file_reader = SerializedFileReader::new(file).unwrap();
- let mut arrow_reader = ParquetFileArrowReader::new(Arc::new(file_reader));
- let mut record_batch_reader = arrow_reader.get_record_reader(1024).unwrap();
-
- let batch = record_batch_reader.next().unwrap().unwrap();
- let decimal_col = batch
- .column(0)
- .as_any()
- .downcast_ref::<DecimalArray>()
- .unwrap();
-
- for i in 0..batch.num_rows() {
- assert_eq!(decimal_col.value(i), raw_decimal_i128_values[i]);
- }
- }
-
- #[test]
- fn arrow_writer_complex() {
- // define schema
- let struct_field_d = Field::new("d", DataType::Float64, true);
- let struct_field_f = Field::new("f", DataType::Float32, true);
- let struct_field_g = Field::new(
- "g",
- DataType::List(Box::new(Field::new("item", DataType::Int16, true))),
- false,
- );
- let struct_field_h = Field::new(
- "h",
- DataType::List(Box::new(Field::new("item", DataType::Int16, false))),
- true,
- );
- let struct_field_e = Field::new(
- "e",
- DataType::Struct(vec![
- struct_field_f.clone(),
- struct_field_g.clone(),
- struct_field_h.clone(),
- ]),
- false,
- );
- let schema = Schema::new(vec![
- Field::new("a", DataType::Int32, false),
- Field::new("b", DataType::Int32, true),
- Field::new(
- "c",
- DataType::Struct(vec![struct_field_d.clone(), struct_field_e.clone()]),
- false,
- ),
- ]);
-
- // create some data
- let a = Int32Array::from(vec![1, 2, 3, 4, 5]);
- let b = Int32Array::from(vec![Some(1), None, None, Some(4), Some(5)]);
- let d = Float64Array::from(vec![None, None, None, Some(1.0), None]);
- let f = Float32Array::from(vec![Some(0.0), None, Some(333.3), None, Some(5.25)]);
-
- let g_value = Int16Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10]);
-
- // Construct a buffer for value offsets, for the nested array:
- // [[1], [2, 3], [], [4, 5, 6], [7, 8, 9, 10]]
- let g_value_offsets =
- arrow::buffer::Buffer::from(&[0, 1, 3, 3, 6, 10].to_byte_slice());
-
- // Construct a list array from the above two
- let g_list_data = ArrayData::builder(struct_field_g.data_type().clone())
- .len(5)
- .add_buffer(g_value_offsets.clone())
- .add_child_data(g_value.data().clone())
- .build();
- let g = ListArray::from(g_list_data);
- // The difference between g and h is that h has a null bitmap
- let h_list_data = ArrayData::builder(struct_field_h.data_type().clone())
- .len(5)
- .add_buffer(g_value_offsets)
- .add_child_data(g_value.data().clone())
- .null_bit_buffer(Buffer::from(vec![0b00011011]))
- .build();
- let h = ListArray::from(h_list_data);
-
- let e = StructArray::from(vec![
- (struct_field_f, Arc::new(f) as ArrayRef),
- (struct_field_g, Arc::new(g) as ArrayRef),
- (struct_field_h, Arc::new(h) as ArrayRef),
- ]);
-
- let c = StructArray::from(vec![
- (struct_field_d, Arc::new(d) as ArrayRef),
- (struct_field_e, Arc::new(e) as ArrayRef),
- ]);
-
- // build a record batch
- let batch = RecordBatch::try_new(
- Arc::new(schema),
- vec![Arc::new(a), Arc::new(b), Arc::new(c)],
- )
- .unwrap();
-
- roundtrip(
- "test_arrow_writer_complex.parquet",
- batch.clone(),
- Some(SMALL_SIZE / 2),
- );
-
- roundtrip(
- "test_arrow_writer_complex_small_batch.parquet",
- batch,
- Some(SMALL_SIZE / 3),
- );
- }
-
- #[test]
- fn arrow_writer_complex_mixed() {
- // This test was added while investigating https://github.com/apache/arrow-rs/issues/244.
- // It was subsequently fixed while investigating https://github.com/apache/arrow-rs/issues/245.
-
- // define schema
- let offset_field = Field::new("offset", DataType::Int32, false);
- let partition_field = Field::new("partition", DataType::Int64, true);
- let topic_field = Field::new("topic", DataType::Utf8, true);
- let schema = Schema::new(vec![Field::new(
- "some_nested_object",
- DataType::Struct(vec![
- offset_field.clone(),
- partition_field.clone(),
- topic_field.clone(),
- ]),
- false,
- )]);
-
- // create some data
- let offset = Int32Array::from(vec![1, 2, 3, 4, 5]);
- let partition = Int64Array::from(vec![Some(1), None, None, Some(4), Some(5)]);
- let topic = StringArray::from(vec![Some("A"), None, Some("A"), Some(""), None]);
-
- let some_nested_object = StructArray::from(vec![
- (offset_field, Arc::new(offset) as ArrayRef),
- (partition_field, Arc::new(partition) as ArrayRef),
- (topic_field, Arc::new(topic) as ArrayRef),
- ]);
-
- // build a record batch
- let batch =
- RecordBatch::try_new(Arc::new(schema), vec![Arc::new(some_nested_object)])
- .unwrap();
-
- roundtrip(
- "test_arrow_writer_complex_mixed.parquet",
- batch,
- Some(SMALL_SIZE / 2),
- );
- }
-
- #[test]
- fn arrow_writer_2_level_struct() {
- // tests writing <struct<struct<primitive>>
- let field_c = Field::new("c", DataType::Int32, true);
- let field_b = Field::new("b", DataType::Struct(vec![field_c]), true);
- let field_a = Field::new("a", DataType::Struct(vec![field_b.clone()]), true);
- let schema = Schema::new(vec![field_a.clone()]);
-
- // create data
- let c = Int32Array::from(vec![Some(1), None, Some(3), None, None, Some(6)]);
- let b_data = ArrayDataBuilder::new(field_b.data_type().clone())
- .len(6)
- .null_bit_buffer(Buffer::from(vec![0b00100111]))
- .add_child_data(c.data().clone())
- .build();
- let b = StructArray::from(b_data);
- let a_data = ArrayDataBuilder::new(field_a.data_type().clone())
- .len(6)
- .null_bit_buffer(Buffer::from(vec![0b00101111]))
- .add_child_data(b.data().clone())
- .build();
- let a = StructArray::from(a_data);
-
- assert_eq!(a.null_count(), 1);
- assert_eq!(a.column(0).null_count(), 2);
-
- // build a racord batch
- let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a)]).unwrap();
-
- roundtrip(
- "test_arrow_writer_2_level_struct.parquet",
- batch,
- Some(SMALL_SIZE / 2),
- );
- }
-
- #[test]
- fn arrow_writer_2_level_struct_non_null() {
- // tests writing <struct<struct<primitive>>
- let field_c = Field::new("c", DataType::Int32, false);
- let field_b = Field::new("b", DataType::Struct(vec![field_c]), false);
- let field_a = Field::new("a", DataType::Struct(vec![field_b.clone()]), false);
- let schema = Schema::new(vec![field_a.clone()]);
-
- // create data
- let c = Int32Array::from(vec![1, 2, 3, 4, 5, 6]);
- let b_data = ArrayDataBuilder::new(field_b.data_type().clone())
- .len(6)
- .add_child_data(c.data().clone())
- .build();
- let b = StructArray::from(b_data);
- let a_data = ArrayDataBuilder::new(field_a.data_type().clone())
- .len(6)
- .add_child_data(b.data().clone())
- .build();
- let a = StructArray::from(a_data);
-
- assert_eq!(a.null_count(), 0);
- assert_eq!(a.column(0).null_count(), 0);
-
- // build a racord batch
- let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a)]).unwrap();
-
- roundtrip(
- "test_arrow_writer_2_level_struct_non_null.parquet",
- batch,
- Some(SMALL_SIZE / 2),
- );
- }
-
- #[test]
- fn arrow_writer_2_level_struct_mixed_null() {
- // tests writing <struct<struct<primitive>>
- let field_c = Field::new("c", DataType::Int32, false);
- let field_b = Field::new("b", DataType::Struct(vec![field_c]), true);
- let field_a = Field::new("a", DataType::Struct(vec![field_b.clone()]), false);
- let schema = Schema::new(vec![field_a.clone()]);
-
- // create data
- let c = Int32Array::from(vec![1, 2, 3, 4, 5, 6]);
- let b_data = ArrayDataBuilder::new(field_b.data_type().clone())
- .len(6)
- .null_bit_buffer(Buffer::from(vec![0b00100111]))
- .add_child_data(c.data().clone())
- .build();
- let b = StructArray::from(b_data);
- // a intentionally has no null buffer, to test that this is handled correctly
- let a_data = ArrayDataBuilder::new(field_a.data_type().clone())
- .len(6)
- .add_child_data(b.data().clone())
- .build();
- let a = StructArray::from(a_data);
-
- assert_eq!(a.null_count(), 0);
- assert_eq!(a.column(0).null_count(), 2);
-
- // build a racord batch
- let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a)]).unwrap();
-
- roundtrip(
- "test_arrow_writer_2_level_struct_mixed_null.parquet",
- batch,
- Some(SMALL_SIZE / 2),
- );
- }
-
- const SMALL_SIZE: usize = 7;
-
- fn roundtrip(
- filename: &str,
- expected_batch: RecordBatch,
- max_row_group_size: Option<usize>,
- ) -> File {
- let file = get_temp_file(filename, &[]);
-
- let mut writer = ArrowWriter::try_new(
- file.try_clone().unwrap(),
- expected_batch.schema(),
- max_row_group_size.map(|size| {
- WriterProperties::builder()
- .set_max_row_group_size(size)
- .build()
- }),
- )
- .expect("Unable to write file");
- writer.write(&expected_batch).unwrap();
- writer.close().unwrap();
-
- let reader = SerializedFileReader::new(file.try_clone().unwrap()).unwrap();
- let mut arrow_reader = ParquetFileArrowReader::new(Arc::new(reader));
- let mut record_batch_reader = arrow_reader.get_record_reader(1024).unwrap();
-
- let actual_batch = record_batch_reader
- .next()
- .expect("No batch found")
- .expect("Unable to get batch");
-
- assert_eq!(expected_batch.schema(), actual_batch.schema());
- assert_eq!(expected_batch.num_columns(), actual_batch.num_columns());
- assert_eq!(expected_batch.num_rows(), actual_batch.num_rows());
- for i in 0..expected_batch.num_columns() {
- let expected_data = expected_batch.column(i).data();
- let actual_data = actual_batch.column(i).data();
-
- assert_eq!(expected_data, actual_data);
- }
-
- file
- }
-
- fn one_column_roundtrip(
- filename: &str,
- values: ArrayRef,
- nullable: bool,
- max_row_group_size: Option<usize>,
- ) -> File {
- let schema = Schema::new(vec![Field::new(
- "col",
- values.data_type().clone(),
- nullable,
- )]);
- let expected_batch =
- RecordBatch::try_new(Arc::new(schema), vec![values]).unwrap();
-
- roundtrip(filename, expected_batch, max_row_group_size)
- }
-
- fn values_required<A, I>(iter: I, filename: &str)
- where
- A: From<Vec<I::Item>> + Array + 'static,
- I: IntoIterator,
- {
- let raw_values: Vec<_> = iter.into_iter().collect();
- let values = Arc::new(A::from(raw_values));
- one_column_roundtrip(filename, values, false, Some(SMALL_SIZE / 2));
- }
-
- fn values_optional<A, I>(iter: I, filename: &str)
- where
- A: From<Vec<Option<I::Item>>> + Array + 'static,
- I: IntoIterator,
- {
- let optional_raw_values: Vec<_> = iter
- .into_iter()
- .enumerate()
- .map(|(i, v)| if i % 2 == 0 { None } else { Some(v) })
- .collect();
- let optional_values = Arc::new(A::from(optional_raw_values));
- one_column_roundtrip(filename, optional_values, true, Some(SMALL_SIZE / 2));
- }
-
- fn required_and_optional<A, I>(iter: I, filename: &str)
- where
- A: From<Vec<I::Item>> + From<Vec<Option<I::Item>>> + Array + 'static,
- I: IntoIterator + Clone,
- {
- values_required::<A, I>(iter.clone(), filename);
- values_optional::<A, I>(iter, filename);
- }
-
- #[test]
- fn all_null_primitive_single_column() {
- let values = Arc::new(Int32Array::from(vec![None; SMALL_SIZE]));
- one_column_roundtrip(
- "all_null_primitive_single_column",
- values,
- true,
- Some(SMALL_SIZE / 2),
- );
- }
- #[test]
- fn null_single_column() {
- let values = Arc::new(NullArray::new(SMALL_SIZE));
- one_column_roundtrip("null_single_column", values, true, Some(SMALL_SIZE / 2));
- // null arrays are always nullable, a test with non-nullable nulls fails
- }
-
- #[test]
- fn bool_single_column() {
- required_and_optional::<BooleanArray, _>(
- [true, false].iter().cycle().copied().take(SMALL_SIZE),
- "bool_single_column",
- );
- }
-
- #[test]
- fn i8_single_column() {
- required_and_optional::<Int8Array, _>(0..SMALL_SIZE as i8, "i8_single_column");
- }
-
- #[test]
- fn i16_single_column() {
- required_and_optional::<Int16Array, _>(0..SMALL_SIZE as i16, "i16_single_column");
- }
-
- #[test]
- fn i32_single_column() {
- required_and_optional::<Int32Array, _>(0..SMALL_SIZE as i32, "i32_single_column");
- }
-
- #[test]
- fn i64_single_column() {
- required_and_optional::<Int64Array, _>(0..SMALL_SIZE as i64, "i64_single_column");
- }
-
- #[test]
- fn u8_single_column() {
- required_and_optional::<UInt8Array, _>(0..SMALL_SIZE as u8, "u8_single_column");
- }
-
- #[test]
- fn u16_single_column() {
- required_and_optional::<UInt16Array, _>(
- 0..SMALL_SIZE as u16,
- "u16_single_column",
- );
- }
-
- #[test]
- fn u32_single_column() {
- required_and_optional::<UInt32Array, _>(
- 0..SMALL_SIZE as u32,
- "u32_single_column",
- );
- }
-
- #[test]
- fn u64_single_column() {
- required_and_optional::<UInt64Array, _>(
- 0..SMALL_SIZE as u64,
- "u64_single_column",
- );
- }
-
- #[test]
- fn f32_single_column() {
- required_and_optional::<Float32Array, _>(
- (0..SMALL_SIZE).map(|i| i as f32),
- "f32_single_column",
- );
- }
-
- #[test]
- fn f64_single_column() {
- required_and_optional::<Float64Array, _>(
- (0..SMALL_SIZE).map(|i| i as f64),
- "f64_single_column",
- );
- }
-
- // The timestamp array types don't implement From<Vec<T>> because they need the timezone
- // argument, and they also doesn't support building from a Vec<Option<T>>, so call
- // one_column_roundtrip manually instead of calling required_and_optional for these tests.
-
- #[test]
- fn timestamp_second_single_column() {
- let raw_values: Vec<_> = (0..SMALL_SIZE as i64).collect();
- let values = Arc::new(TimestampSecondArray::from_vec(raw_values, None));
-
- one_column_roundtrip("timestamp_second_single_column", values, false, Some(3));
- }
-
- #[test]
- fn timestamp_millisecond_single_column() {
- let raw_values: Vec<_> = (0..SMALL_SIZE as i64).collect();
- let values = Arc::new(TimestampMillisecondArray::from_vec(raw_values, None));
-
- one_column_roundtrip(
- "timestamp_millisecond_single_column",
- values,
- false,
- Some(SMALL_SIZE / 2 + 1),
- );
- }
-
- #[test]
- fn timestamp_microsecond_single_column() {
- let raw_values: Vec<_> = (0..SMALL_SIZE as i64).collect();
- let values = Arc::new(TimestampMicrosecondArray::from_vec(raw_values, None));
-
- one_column_roundtrip(
- "timestamp_microsecond_single_column",
- values,
- false,
- Some(SMALL_SIZE / 2 + 2),
- );
- }
-
- #[test]
- fn timestamp_nanosecond_single_column() {
- let raw_values: Vec<_> = (0..SMALL_SIZE as i64).collect();
- let values = Arc::new(TimestampNanosecondArray::from_vec(raw_values, None));
-
- one_column_roundtrip(
- "timestamp_nanosecond_single_column",
- values,
- false,
- Some(SMALL_SIZE / 2),
- );
- }
-
- #[test]
- fn date32_single_column() {
- required_and_optional::<Date32Array, _>(
- 0..SMALL_SIZE as i32,
- "date32_single_column",
- );
- }
-
- #[test]
- fn date64_single_column() {
- // Date64 must be a multiple of 86400000, see ARROW-10925
- required_and_optional::<Date64Array, _>(
- (0..(SMALL_SIZE as i64 * 86400000)).step_by(86400000),
- "date64_single_column",
- );
- }
-
- #[test]
- fn time32_second_single_column() {
- required_and_optional::<Time32SecondArray, _>(
- 0..SMALL_SIZE as i32,
- "time32_second_single_column",
- );
- }
-
- #[test]
- fn time32_millisecond_single_column() {
- required_and_optional::<Time32MillisecondArray, _>(
- 0..SMALL_SIZE as i32,
- "time32_millisecond_single_column",
- );
- }
-
- #[test]
- fn time64_microsecond_single_column() {
- required_and_optional::<Time64MicrosecondArray, _>(
- 0..SMALL_SIZE as i64,
- "time64_microsecond_single_column",
- );
- }
-
- #[test]
- fn time64_nanosecond_single_column() {
- required_and_optional::<Time64NanosecondArray, _>(
- 0..SMALL_SIZE as i64,
- "time64_nanosecond_single_column",
- );
- }
-
- #[test]
- #[should_panic(expected = "Converting Duration to parquet not supported")]
- fn duration_second_single_column() {
- required_and_optional::<DurationSecondArray, _>(
- 0..SMALL_SIZE as i64,
- "duration_second_single_column",
- );
- }
-
- #[test]
- #[should_panic(expected = "Converting Duration to parquet not supported")]
- fn duration_millisecond_single_column() {
- required_and_optional::<DurationMillisecondArray, _>(
- 0..SMALL_SIZE as i64,
- "duration_millisecond_single_column",
- );
- }
-
- #[test]
- #[should_panic(expected = "Converting Duration to parquet not supported")]
- fn duration_microsecond_single_column() {
- required_and_optional::<DurationMicrosecondArray, _>(
- 0..SMALL_SIZE as i64,
- "duration_microsecond_single_column",
- );
- }
-
- #[test]
- #[should_panic(expected = "Converting Duration to parquet not supported")]
- fn duration_nanosecond_single_column() {
- required_and_optional::<DurationNanosecondArray, _>(
- 0..SMALL_SIZE as i64,
- "duration_nanosecond_single_column",
- );
- }
-
- #[test]
- fn interval_year_month_single_column() {
- required_and_optional::<IntervalYearMonthArray, _>(
- 0..SMALL_SIZE as i32,
- "interval_year_month_single_column",
- );
- }
-
- #[test]
- fn interval_day_time_single_column() {
- required_and_optional::<IntervalDayTimeArray, _>(
- 0..SMALL_SIZE as i64,
- "interval_day_time_single_column",
- );
- }
-
- #[test]
- fn binary_single_column() {
- let one_vec: Vec<u8> = (0..SMALL_SIZE as u8).collect();
- let many_vecs: Vec<_> = std::iter::repeat(one_vec).take(SMALL_SIZE).collect();
- let many_vecs_iter = many_vecs.iter().map(|v| v.as_slice());
-
- // BinaryArrays can't be built from Vec<Option<&str>>, so only call `values_required`
- values_required::<BinaryArray, _>(many_vecs_iter, "binary_single_column");
- }
-
- #[test]
- fn large_binary_single_column() {
- let one_vec: Vec<u8> = (0..SMALL_SIZE as u8).collect();
- let many_vecs: Vec<_> = std::iter::repeat(one_vec).take(SMALL_SIZE).collect();
- let many_vecs_iter = many_vecs.iter().map(|v| v.as_slice());
-
- // LargeBinaryArrays can't be built from Vec<Option<&str>>, so only call `values_required`
- values_required::<LargeBinaryArray, _>(
- many_vecs_iter,
- "large_binary_single_column",
- );
- }
-
- #[test]
- fn fixed_size_binary_single_column() {
- let mut builder = FixedSizeBinaryBuilder::new(16, 4);
- builder.append_value(b"0123").unwrap();
- builder.append_null().unwrap();
- builder.append_value(b"8910").unwrap();
- builder.append_value(b"1112").unwrap();
- let array = Arc::new(builder.finish());
-
- one_column_roundtrip(
- "fixed_size_binary_single_column",
- array,
- true,
- Some(SMALL_SIZE / 2),
- );
- }
-
- #[test]
- fn string_single_column() {
- let raw_values: Vec<_> = (0..SMALL_SIZE).map(|i| i.to_string()).collect();
- let raw_strs = raw_values.iter().map(|s| s.as_str());
-
- required_and_optional::<StringArray, _>(raw_strs, "string_single_column");
- }
-
- #[test]
- fn large_string_single_column() {
- let raw_values: Vec<_> = (0..SMALL_SIZE).map(|i| i.to_string()).collect();
- let raw_strs = raw_values.iter().map(|s| s.as_str());
-
- required_and_optional::<LargeStringArray, _>(
- raw_strs,
- "large_string_single_column",
- );
- }
-
- #[test]
- fn list_single_column() {
- let a_values = Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10]);
- let a_value_offsets =
- arrow::buffer::Buffer::from(&[0, 1, 3, 3, 6, 10].to_byte_slice());
- let a_list_data = ArrayData::builder(DataType::List(Box::new(Field::new(
- "item",
- DataType::Int32,
- false,
- ))))
- .len(5)
- .add_buffer(a_value_offsets)
- .null_bit_buffer(Buffer::from(vec![0b00011011]))
- .add_child_data(a_values.data().clone())
- .build();
-
- assert_eq!(a_list_data.null_count(), 1);
-
- let a = ListArray::from(a_list_data);
- let values = Arc::new(a);
-
- one_column_roundtrip("list_single_column", values, true, Some(SMALL_SIZE / 2));
- }
-
- #[test]
- fn large_list_single_column() {
- let a_values = Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10]);
- let a_value_offsets =
- arrow::buffer::Buffer::from(&[0i64, 1, 3, 3, 6, 10].to_byte_slice());
- let a_list_data = ArrayData::builder(DataType::LargeList(Box::new(Field::new(
- "large_item",
- DataType::Int32,
- true,
- ))))
- .len(5)
- .add_buffer(a_value_offsets)
- .add_child_data(a_values.data().clone())
- .null_bit_buffer(Buffer::from(vec![0b00011011]))
- .build();
-
- // I think this setup is incorrect because this should pass
- assert_eq!(a_list_data.null_count(), 1);
-
- let a = LargeListArray::from(a_list_data);
- let values = Arc::new(a);
-
- one_column_roundtrip(
- "large_list_single_column",
- values,
- true,
- Some(SMALL_SIZE / 2),
- );
- }
-
- #[test]
- fn struct_single_column() {
- let a_values = Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10]);
- let struct_field_a = Field::new("f", DataType::Int32, false);
- let s = StructArray::from(vec![(struct_field_a, Arc::new(a_values) as ArrayRef)]);
-
- let values = Arc::new(s);
- one_column_roundtrip("struct_single_column", values, false, Some(SMALL_SIZE / 2));
- }
-
- #[test]
- fn arrow_writer_string_dictionary() {
- // define schema
- let schema = Arc::new(Schema::new(vec![Field::new_dict(
- "dictionary",
- DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
- true,
- 42,
- true,
- )]));
-
- // create some data
- let d: Int32DictionaryArray = [Some("alpha"), None, Some("beta"), Some("alpha")]
- .iter()
- .copied()
- .collect();
-
- // build a record batch
- let expected_batch = RecordBatch::try_new(schema, vec![Arc::new(d)]).unwrap();
-
- roundtrip(
- "test_arrow_writer_string_dictionary.parquet",
- expected_batch,
- Some(SMALL_SIZE / 2),
- );
- }
-
- #[test]
- fn arrow_writer_primitive_dictionary() {
- // define schema
- let schema = Arc::new(Schema::new(vec![Field::new_dict(
- "dictionary",
- DataType::Dictionary(Box::new(DataType::UInt8), Box::new(DataType::UInt32)),
- true,
- 42,
- true,
- )]));
-
- // create some data
- let key_builder = PrimitiveBuilder::<UInt8Type>::new(3);
- let value_builder = PrimitiveBuilder::<UInt32Type>::new(2);
- let mut builder = PrimitiveDictionaryBuilder::new(key_builder, value_builder);
- builder.append(12345678).unwrap();
- builder.append_null().unwrap();
- builder.append(22345678).unwrap();
- builder.append(12345678).unwrap();
- let d = builder.finish();
-
- // build a record batch
- let expected_batch = RecordBatch::try_new(schema, vec![Arc::new(d)]).unwrap();
-
- roundtrip(
- "test_arrow_writer_primitive_dictionary.parquet",
- expected_batch,
- Some(SMALL_SIZE / 2),
- );
- }
-
- #[test]
- fn arrow_writer_string_dictionary_unsigned_index() {
- // define schema
- let schema = Arc::new(Schema::new(vec![Field::new_dict(
- "dictionary",
- DataType::Dictionary(Box::new(DataType::UInt8), Box::new(DataType::Utf8)),
- true,
- 42,
- true,
- )]));
-
- // create some data
- let d: UInt8DictionaryArray = [Some("alpha"), None, Some("beta"), Some("alpha")]
- .iter()
- .copied()
- .collect();
-
- // build a record batch
- let expected_batch = RecordBatch::try_new(schema, vec![Arc::new(d)]).unwrap();
-
- roundtrip(
- "test_arrow_writer_string_dictionary_unsigned_index.parquet",
- expected_batch,
- Some(SMALL_SIZE / 2),
- );
- }
-
- #[test]
- fn u32_min_max() {
- // check values roundtrip through parquet
- let values = Arc::new(UInt32Array::from_iter_values(vec![
- u32::MIN,
- u32::MIN + 1,
- (i32::MAX as u32) - 1,
- i32::MAX as u32,
- (i32::MAX as u32) + 1,
- u32::MAX - 1,
- u32::MAX,
- ]));
- let file = one_column_roundtrip("u32_min_max_single_column", values, false, None);
-
- // check statistics are valid
- let reader = SerializedFileReader::new(file).unwrap();
- let metadata = reader.metadata();
- assert_eq!(metadata.num_row_groups(), 1);
- let row_group = metadata.row_group(0);
- assert_eq!(row_group.num_columns(), 1);
- let column = row_group.column(0);
- let stats = column.statistics().unwrap();
- assert!(stats.has_min_max_set());
- if let Statistics::Int32(stats) = stats {
- assert_eq!(*stats.min() as u32, u32::MIN);
- assert_eq!(*stats.max() as u32, u32::MAX);
- } else {
- panic!("Statistics::Int32 missing")
- }
- }
-
- #[test]
- fn u64_min_max() {
- // check values roundtrip through parquet
- let values = Arc::new(UInt64Array::from_iter_values(vec![
- u64::MIN,
- u64::MIN + 1,
- (i64::MAX as u64) - 1,
- i64::MAX as u64,
- (i64::MAX as u64) + 1,
- u64::MAX - 1,
- u64::MAX,
- ]));
- let file = one_column_roundtrip("u64_min_max_single_column", values, false, None);
-
- // check statistics are valid
- let reader = SerializedFileReader::new(file).unwrap();
- let metadata = reader.metadata();
- assert_eq!(metadata.num_row_groups(), 1);
- let row_group = metadata.row_group(0);
- assert_eq!(row_group.num_columns(), 1);
- let column = row_group.column(0);
- let stats = column.statistics().unwrap();
- assert!(stats.has_min_max_set());
- if let Statistics::Int64(stats) = stats {
- assert_eq!(*stats.min() as u64, u64::MIN);
- assert_eq!(*stats.max() as u64, u64::MAX);
- } else {
- panic!("Statistics::Int64 missing")
- }
- }
-
- #[test]
- fn statistics_null_counts_only_nulls() {
- // check that null-count statistics for "only NULL"-columns are correct
- let values = Arc::new(UInt64Array::from(vec![None, None]));
- let file = one_column_roundtrip("null_counts", values, true, None);
-
- // check statistics are valid
- let reader = SerializedFileReader::new(file).unwrap();
- let metadata = reader.metadata();
- assert_eq!(metadata.num_row_groups(), 1);
- let row_group = metadata.row_group(0);
- assert_eq!(row_group.num_columns(), 1);
- let column = row_group.column(0);
- let stats = column.statistics().unwrap();
- assert_eq!(stats.null_count(), 2);
- }
-}
diff --git a/parquet/src/arrow/converter.rs b/parquet/src/arrow/converter.rs
deleted file mode 100644
index 1672be9..0000000
--- a/parquet/src/arrow/converter.rs
+++ /dev/null
@@ -1,454 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-use crate::data_type::{ByteArray, DataType, FixedLenByteArray, Int96};
-// TODO: clean up imports (best done when there are few moving parts)
-use arrow::array::{
- Array, ArrayRef, BinaryBuilder, DecimalBuilder, FixedSizeBinaryBuilder,
- IntervalDayTimeArray, IntervalDayTimeBuilder, IntervalYearMonthArray,
- IntervalYearMonthBuilder, LargeBinaryBuilder, LargeStringBuilder, PrimitiveBuilder,
- PrimitiveDictionaryBuilder, StringBuilder, StringDictionaryBuilder,
-};
-use arrow::compute::cast;
-use std::convert::{From, TryInto};
-use std::sync::Arc;
-
-use crate::errors::Result;
-use arrow::datatypes::{ArrowDictionaryKeyType, ArrowPrimitiveType};
-
-use arrow::array::{
- BinaryArray, DecimalArray, DictionaryArray, FixedSizeBinaryArray, LargeBinaryArray,
- LargeStringArray, PrimitiveArray, StringArray, TimestampNanosecondArray,
-};
-use std::marker::PhantomData;
-
-use crate::data_type::Int32Type as ParquetInt32Type;
-use arrow::datatypes::Int32Type;
-
-/// A converter is used to consume record reader's content and convert it to arrow
-/// primitive array.
-pub trait Converter<S, T> {
- /// This method converts record reader's buffered content into arrow array.
- /// It will consume record reader's data, but will not reset record reader's
- /// state.
- fn convert(&self, source: S) -> Result<T>;
-}
-
-pub struct FixedSizeArrayConverter {
- byte_width: i32,
-}
-
-impl FixedSizeArrayConverter {
- pub fn new(byte_width: i32) -> Self {
- Self { byte_width }
- }
-}
-
-impl Converter<Vec<Option<FixedLenByteArray>>, FixedSizeBinaryArray>
- for FixedSizeArrayConverter
-{
- fn convert(
- &self,
- source: Vec<Option<FixedLenByteArray>>,
- ) -> Result<FixedSizeBinaryArray> {
- let mut builder = FixedSizeBinaryBuilder::new(source.len(), self.byte_width);
- for v in source {
- match v {
- Some(array) => builder.append_value(array.data()),
- None => builder.append_null(),
- }?
- }
-
- Ok(builder.finish())
- }
-}
-
-pub struct DecimalArrayConverter {
- precision: i32,
- scale: i32,
-}
-
-impl DecimalArrayConverter {
- pub fn new(precision: i32, scale: i32) -> Self {
- Self { precision, scale }
- }
-
- fn from_bytes_to_i128(b: &[u8]) -> i128 {
- assert!(b.len() <= 16, "DecimalArray supports only up to size 16");
- let first_bit = b[0] & 128u8 == 128u8;
- let mut result = if first_bit { [255u8; 16] } else { [0u8; 16] };
- for (i, v) in b.iter().enumerate() {
- result[i + (16 - b.len())] = *v;
- }
- i128::from_be_bytes(result)
- }
-}
-
-impl Converter<Vec<Option<FixedLenByteArray>>, DecimalArray> for DecimalArrayConverter {
- fn convert(&self, source: Vec<Option<FixedLenByteArray>>) -> Result<DecimalArray> {
- let mut builder = DecimalBuilder::new(
- source.len(),
- self.precision as usize,
- self.scale as usize,
- );
- for v in source {
- match v {
- Some(array) => {
- builder.append_value(Self::from_bytes_to_i128(array.data()))
- }
- None => builder.append_null(),
- }?
- }
-
- Ok(builder.finish())
- }
-}
-/// An Arrow Interval converter, which reads the first 4 bytes of a Parquet interval,
-/// and interprets it as an i32 value representing the Arrow YearMonth value
-pub struct IntervalYearMonthArrayConverter {}
-
-impl Converter<Vec<Option<FixedLenByteArray>>, IntervalYearMonthArray>
- for IntervalYearMonthArrayConverter
-{
- fn convert(
- &self,
- source: Vec<Option<FixedLenByteArray>>,
- ) -> Result<IntervalYearMonthArray> {
- let mut builder = IntervalYearMonthBuilder::new(source.len());
- for v in source {
- match v {
- Some(array) => builder.append_value(i32::from_le_bytes(
- array.data()[0..4].try_into().unwrap(),
- )),
- None => builder.append_null(),
- }?
- }
-
- Ok(builder.finish())
- }
-}
-
-/// An Arrow Interval converter, which reads the last 8 bytes of a Parquet interval,
-/// and interprets it as an i32 value representing the Arrow DayTime value
-pub struct IntervalDayTimeArrayConverter {}
-
-impl Converter<Vec<Option<FixedLenByteArray>>, IntervalDayTimeArray>
- for IntervalDayTimeArrayConverter
-{
- fn convert(
- &self,
- source: Vec<Option<FixedLenByteArray>>,
- ) -> Result<IntervalDayTimeArray> {
- let mut builder = IntervalDayTimeBuilder::new(source.len());
- for v in source {
- match v {
- Some(array) => builder.append_value(i64::from_le_bytes(
- array.data()[4..12].try_into().unwrap(),
- )),
- None => builder.append_null(),
- }?
- }
-
- Ok(builder.finish())
- }
-}
-
-pub struct Int96ArrayConverter {
- pub timezone: Option<String>,
-}
-
-impl Converter<Vec<Option<Int96>>, TimestampNanosecondArray> for Int96ArrayConverter {
- fn convert(&self, source: Vec<Option<Int96>>) -> Result<TimestampNanosecondArray> {
- Ok(TimestampNanosecondArray::from_opt_vec(
- source
- .into_iter()
- .map(|int96| int96.map(|val| val.to_i64() * 1_000_000))
- .collect(),
- self.timezone.clone(),
- ))
- }
-}
-
-pub struct Utf8ArrayConverter {}
-
-impl Converter<Vec<Option<ByteArray>>, StringArray> for Utf8ArrayConverter {
- fn convert(&self, source: Vec<Option<ByteArray>>) -> Result<StringArray> {
- let data_size = source
- .iter()
- .map(|x| x.as_ref().map(|b| b.len()).unwrap_or(0))
- .sum();
-
- let mut builder = StringBuilder::with_capacity(source.len(), data_size);
- for v in source {
- match v {
- Some(array) => builder.append_value(array.as_utf8()?),
- None => builder.append_null(),
- }?
- }
-
- Ok(builder.finish())
- }
-}
-
-pub struct LargeUtf8ArrayConverter {}
-
-impl Converter<Vec<Option<ByteArray>>, LargeStringArray> for LargeUtf8ArrayConverter {
- fn convert(&self, source: Vec<Option<ByteArray>>) -> Result<LargeStringArray> {
- let data_size = source
- .iter()
- .map(|x| x.as_ref().map(|b| b.len()).unwrap_or(0))
- .sum();
-
- let mut builder = LargeStringBuilder::with_capacity(source.len(), data_size);
- for v in source {
- match v {
- Some(array) => builder.append_value(array.as_utf8()?),
- None => builder.append_null(),
- }?
- }
-
- Ok(builder.finish())
- }
-}
-
-pub struct BinaryArrayConverter {}
-
-impl Converter<Vec<Option<ByteArray>>, BinaryArray> for BinaryArrayConverter {
- fn convert(&self, source: Vec<Option<ByteArray>>) -> Result<BinaryArray> {
- let mut builder = BinaryBuilder::new(source.len());
- for v in source {
- match v {
- Some(array) => builder.append_value(array.data()),
- None => builder.append_null(),
- }?
- }
-
- Ok(builder.finish())
- }
-}
-
-pub struct LargeBinaryArrayConverter {}
-
-impl Converter<Vec<Option<ByteArray>>, LargeBinaryArray> for LargeBinaryArrayConverter {
- fn convert(&self, source: Vec<Option<ByteArray>>) -> Result<LargeBinaryArray> {
- let mut builder = LargeBinaryBuilder::new(source.len());
- for v in source {
- match v {
- Some(array) => builder.append_value(array.data()),
- None => builder.append_null(),
- }?
- }
-
- Ok(builder.finish())
- }
-}
-
-pub struct StringDictionaryArrayConverter {}
-
-impl<K: ArrowDictionaryKeyType> Converter<Vec<Option<ByteArray>>, DictionaryArray<K>>
- for StringDictionaryArrayConverter
-{
- fn convert(&self, source: Vec<Option<ByteArray>>) -> Result<DictionaryArray<K>> {
- let data_size = source
- .iter()
- .map(|x| x.as_ref().map(|b| b.len()).unwrap_or(0))
- .sum();
-
- let keys_builder = PrimitiveBuilder::<K>::new(source.len());
- let values_builder = StringBuilder::with_capacity(source.len(), data_size);
-
- let mut builder = StringDictionaryBuilder::new(keys_builder, values_builder);
- for v in source {
- match v {
- Some(array) => {
- let _ = builder.append(array.as_utf8()?)?;
- }
- None => builder.append_null()?,
- }
- }
-
- Ok(builder.finish())
- }
-}
-
-pub struct DictionaryArrayConverter<DictValueSourceType, DictValueTargetType, ParquetType>
-{
- _dict_value_source_marker: PhantomData<DictValueSourceType>,
- _dict_value_target_marker: PhantomData<DictValueTargetType>,
- _parquet_marker: PhantomData<ParquetType>,
-}
-
-impl<DictValueSourceType, DictValueTargetType, ParquetType>
- DictionaryArrayConverter<DictValueSourceType, DictValueTargetType, ParquetType>
-{
- pub fn new() -> Self {
- Self {
- _dict_value_source_marker: PhantomData,
- _dict_value_target_marker: PhantomData,
- _parquet_marker: PhantomData,
- }
- }
-}
-
-impl<K, DictValueSourceType, DictValueTargetType, ParquetType>
- Converter<Vec<Option<<ParquetType as DataType>::T>>, DictionaryArray<K>>
- for DictionaryArrayConverter<DictValueSourceType, DictValueTargetType, ParquetType>
-where
- K: ArrowPrimitiveType,
- DictValueSourceType: ArrowPrimitiveType,
- DictValueTargetType: ArrowPrimitiveType,
- ParquetType: DataType,
- PrimitiveArray<DictValueSourceType>: From<Vec<Option<<ParquetType as DataType>::T>>>,
-{
- fn convert(
- &self,
- source: Vec<Option<<ParquetType as DataType>::T>>,
- ) -> Result<DictionaryArray<K>> {
- let keys_builder = PrimitiveBuilder::<K>::new(source.len());
- let values_builder = PrimitiveBuilder::<DictValueTargetType>::new(source.len());
-
- let mut builder = PrimitiveDictionaryBuilder::new(keys_builder, values_builder);
-
- let source_array: Arc<dyn Array> =
- Arc::new(PrimitiveArray::<DictValueSourceType>::from(source));
- let target_array = cast(&source_array, &DictValueTargetType::DATA_TYPE)?;
- let target = target_array
- .as_any()
- .downcast_ref::<PrimitiveArray<DictValueTargetType>>()
- .unwrap();
-
- for i in 0..target.len() {
- if target.is_null(i) {
- builder.append_null()?;
- } else {
- let _ = builder.append(target.value(i))?;
- }
- }
-
- Ok(builder.finish())
- }
-}
-
-pub type Utf8Converter =
- ArrayRefConverter<Vec<Option<ByteArray>>, StringArray, Utf8ArrayConverter>;
-pub type LargeUtf8Converter =
- ArrayRefConverter<Vec<Option<ByteArray>>, LargeStringArray, LargeUtf8ArrayConverter>;
-pub type BinaryConverter =
- ArrayRefConverter<Vec<Option<ByteArray>>, BinaryArray, BinaryArrayConverter>;
-pub type LargeBinaryConverter = ArrayRefConverter<
- Vec<Option<ByteArray>>,
- LargeBinaryArray,
- LargeBinaryArrayConverter,
->;
-pub type StringDictionaryConverter<T> = ArrayRefConverter<
- Vec<Option<ByteArray>>,
- DictionaryArray<T>,
- StringDictionaryArrayConverter,
->;
-pub type DictionaryConverter<K, SV, TV, P> = ArrayRefConverter<
- Vec<Option<<P as DataType>::T>>,
- DictionaryArray<K>,
- DictionaryArrayConverter<SV, TV, P>,
->;
-pub type PrimitiveDictionaryConverter<K, V> = ArrayRefConverter<
- Vec<Option<<ParquetInt32Type as DataType>::T>>,
- DictionaryArray<K>,
- DictionaryArrayConverter<Int32Type, V, ParquetInt32Type>,
->;
-
-pub type Int96Converter =
- ArrayRefConverter<Vec<Option<Int96>>, TimestampNanosecondArray, Int96ArrayConverter>;
-
-pub type FixedLenBinaryConverter = ArrayRefConverter<
- Vec<Option<FixedLenByteArray>>,
- FixedSizeBinaryArray,
- FixedSizeArrayConverter,
->;
-pub type IntervalYearMonthConverter = ArrayRefConverter<
- Vec<Option<FixedLenByteArray>>,
- IntervalYearMonthArray,
- IntervalYearMonthArrayConverter,
->;
-pub type IntervalDayTimeConverter = ArrayRefConverter<
- Vec<Option<FixedLenByteArray>>,
- IntervalDayTimeArray,
- IntervalDayTimeArrayConverter,
->;
-
-pub type DecimalConverter = ArrayRefConverter<
- Vec<Option<FixedLenByteArray>>,
- DecimalArray,
- DecimalArrayConverter,
->;
-
-pub struct FromConverter<S, T> {
- _source: PhantomData<S>,
- _dest: PhantomData<T>,
-}
-
-impl<S, T> FromConverter<S, T>
-where
- T: From<S>,
-{
- pub fn new() -> Self {
- Self {
- _source: PhantomData,
- _dest: PhantomData,
- }
- }
-}
-
-impl<S, T> Converter<S, T> for FromConverter<S, T>
-where
- T: From<S>,
-{
- fn convert(&self, source: S) -> Result<T> {
- Ok(T::from(source))
- }
-}
-
-pub struct ArrayRefConverter<S, A, C> {
- _source: PhantomData<S>,
- _array: PhantomData<A>,
- converter: C,
-}
-
-impl<S, A, C> ArrayRefConverter<S, A, C>
-where
- A: Array + 'static,
- C: Converter<S, A> + 'static,
-{
- pub fn new(converter: C) -> Self {
- Self {
- _source: PhantomData,
- _array: PhantomData,
- converter,
- }
- }
-}
-
-impl<S, A, C> Converter<S, ArrayRef> for ArrayRefConverter<S, A, C>
-where
- A: Array + 'static,
- C: Converter<S, A> + 'static,
-{
- fn convert(&self, source: S) -> Result<ArrayRef> {
- self.converter
- .convert(source)
- .map(|array| Arc::new(array) as ArrayRef)
- }
-}
diff --git a/parquet/src/arrow/levels.rs b/parquet/src/arrow/levels.rs
deleted file mode 100644
index 2e95039..0000000
--- a/parquet/src/arrow/levels.rs
+++ /dev/null
@@ -1,1548 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-//! Parquet definition and repetition levels
-//!
-//! Contains the algorithm for computing definition and repetition levels.
-//! The algorithm works by tracking the slots of an array that should
-//! ultimately be populated when writing to Parquet.
-//! Parquet achieves nesting through definition levels and repetition levels \[1\].
-//! Definition levels specify how many optional fields in the part for the column
-//! are defined.
-//! Repetition levels specify at what repeated field (list) in the path a column
-//! is defined.
-//!
-//! In a nested data structure such as `a.b.c`, one can see levels as defining
-//! whether a record is defined at `a`, `a.b`, or `a.b.c`.
-//! Optional fields are nullable fields, thus if all 3 fields
-//! are nullable, the maximum definition could be = 3 if there are no lists.
-//!
-//! The algorithm in this module computes the necessary information to enable
-//! the writer to keep track of which columns are at which levels, and to extract
-//! the correct values at the correct slots from Arrow arrays.
-//!
-//! It works by walking a record batch's arrays, keeping track of what values
-//! are non-null, their positions and computing what their levels are.
-//!
-//! \[1\] [parquet-format#nested-encoding](https://github.com/apache/parquet-format#nested-encoding)
-
-use arrow::array::{make_array, ArrayRef, StructArray};
-use arrow::datatypes::{DataType, Field};
-
-/// Keeps track of the level information per array that is needed to write an Arrow array to Parquet.
-///
-/// When a nested schema is traversed, intermediate [LevelInfo] structs are created to track
-/// the state of parent arrays. When a primitive Arrow array is encountered, a final [LevelInfo]
-/// is created, and this is what is used to index into the array when writing data to Parquet.
-#[derive(Debug, Eq, PartialEq, Clone)]
-pub(crate) struct LevelInfo {
- /// Array's definition levels
- pub definition: Vec<i16>,
- /// Array's optional repetition levels
- pub repetition: Option<Vec<i16>>,
- /// Array's offsets, 64-bit is used to accommodate large offset arrays
- pub array_offsets: Vec<i64>,
- // TODO: Convert to an Arrow Buffer after ARROW-10766 is merged.
- /// Array's logical validity mask, whcih gets unpacked for list children.
- /// If the parent of an array is null, all children are logically treated as
- /// null. This mask keeps track of that.
- ///
- pub array_mask: Vec<bool>,
- /// The maximum definition at this level, 0 at the record batch
- pub max_definition: i16,
- /// The type of array represented by this level info
- pub level_type: LevelType,
- /// The offset of the current level's array
- pub offset: usize,
- /// The length of the current level's array
- pub length: usize,
-}
-
-/// LevelType defines the type of level, and whether it is nullable or not
-#[derive(Debug, Eq, PartialEq, Clone, Copy)]
-pub(crate) enum LevelType {
- Root,
- List(bool),
- Struct(bool),
- Primitive(bool),
-}
-
-impl LevelType {
- #[inline]
- const fn level_increment(&self) -> i16 {
- match self {
- LevelType::Root => 0,
- LevelType::List(is_nullable)
- | LevelType::Struct(is_nullable)
- | LevelType::Primitive(is_nullable) => *is_nullable as i16,
- }
- }
-}
-
-impl LevelInfo {
- /// Create a new [LevelInfo] by filling `length` slots, and setting an initial offset.
- ///
- /// This is a convenience function to populate the starting point of the traversal.
- pub(crate) fn new(offset: usize, length: usize) -> Self {
- Self {
- // a batch has no definition level yet
- definition: vec![0; length],
- // a batch has no repetition as it is not a list
- repetition: None,
- // a batch has sequential offsets, should be num_rows + 1
- array_offsets: (0..=(length as i64)).collect(),
- // all values at a batch-level are non-null
- array_mask: vec![true; length],
- max_definition: 0,
- level_type: LevelType::Root,
- offset,
- length,
- }
- }
-
- /// Compute nested levels of the Arrow array, recursing into lists and structs.
- ///
- /// Returns a list of `LevelInfo`, where each level is for nested primitive arrays.
- ///
- /// The parent struct's nullness is tracked, as it determines whether the child
- /// max_definition should be incremented.
- /// The 'is_parent_struct' variable asks "is this field's parent a struct?".
- /// * If we are starting at a [RecordBatch], this is `false`.
- /// * If we are calculating a list's child, this is `false`.
- /// * If we are calculating a struct (i.e. `field.data_type90 == Struct`),
- /// this depends on whether the struct is a child of a struct.
- /// * If we are calculating a field inside a [StructArray], this is 'true'.
- pub(crate) fn calculate_array_levels(
- &self,
- array: &ArrayRef,
- field: &Field,
- ) -> Vec<Self> {
- let (array_offsets, array_mask) =
- Self::get_array_offsets_and_masks(array, self.offset, self.length);
- match array.data_type() {
- DataType::Null => vec![Self {
- definition: self.definition.clone(),
- repetition: self.repetition.clone(),
- array_offsets,
- array_mask,
- max_definition: self.max_definition.max(1),
- // Null type is always nullable
- level_type: LevelType::Primitive(true),
- offset: self.offset,
- length: self.length,
- }],
- DataType::Boolean
- | DataType::Int8
- | DataType::Int16
- | DataType::Int32
- | DataType::Int64
- | DataType::UInt8
- | DataType::UInt16
- | DataType::UInt32
- | DataType::UInt64
- | DataType::Float16
- | DataType::Float32
- | DataType::Float64
- | DataType::Utf8
- | DataType::LargeUtf8
- | DataType::Timestamp(_, _)
- | DataType::Date32
- | DataType::Date64
- | DataType::Time32(_)
- | DataType::Time64(_)
- | DataType::Duration(_)
- | DataType::Interval(_)
- | DataType::Binary
- | DataType::LargeBinary
- | DataType::Decimal(_, _)
- | DataType::FixedSizeBinary(_) => {
- // we return a vector of 1 value to represent the primitive
- vec![self.calculate_child_levels(
- array_offsets,
- array_mask,
- LevelType::Primitive(field.is_nullable()),
- )]
- }
- DataType::List(list_field) | DataType::LargeList(list_field) => {
- let child_offset = array_offsets[0] as usize;
- let child_len = *array_offsets.last().unwrap() as usize;
- // Calculate the list level
- let list_level = self.calculate_child_levels(
- array_offsets,
- array_mask,
- LevelType::List(field.is_nullable()),
- );
-
- // Construct the child array of the list, and get its offset + mask
- let array_data = array.data();
- let child_data = array_data.child_data().get(0).unwrap();
- let child_array = make_array(child_data.clone());
- let (child_offsets, child_mask) = Self::get_array_offsets_and_masks(
- &child_array,
- child_offset,
- child_len - child_offset,
- );
-
- match child_array.data_type() {
- // TODO: The behaviour of a <list<null>> is untested
- DataType::Null => vec![list_level],
- DataType::Boolean
- | DataType::Int8
- | DataType::Int16
- | DataType::Int32
- | DataType::Int64
- | DataType::UInt8
- | DataType::UInt16
- | DataType::UInt32
- | DataType::UInt64
- | DataType::Float16
- | DataType::Float32
- | DataType::Float64
- | DataType::Timestamp(_, _)
- | DataType::Date32
- | DataType::Date64
- | DataType::Time32(_)
- | DataType::Time64(_)
- | DataType::Duration(_)
- | DataType::Interval(_)
- | DataType::Binary
- | DataType::LargeBinary
- | DataType::Utf8
- | DataType::LargeUtf8
- | DataType::Dictionary(_, _)
- | DataType::Decimal(_, _)
- | DataType::FixedSizeBinary(_) => {
- vec![list_level.calculate_child_levels(
- child_offsets,
- child_mask,
- LevelType::Primitive(list_field.is_nullable()),
- )]
- }
- DataType::List(_) | DataType::LargeList(_) | DataType::Struct(_) => {
- list_level.calculate_array_levels(&child_array, list_field)
- }
- DataType::FixedSizeList(_, _) => unimplemented!(),
- DataType::Union(_) => unimplemented!(),
- }
- }
- DataType::FixedSizeList(_, _) => unimplemented!(),
- DataType::Struct(struct_fields) => {
- let struct_array: &StructArray = array
- .as_any()
- .downcast_ref::<StructArray>()
- .expect("Unable to get struct array");
- let struct_level = self.calculate_child_levels(
- array_offsets,
- array_mask,
- LevelType::Struct(field.is_nullable()),
- );
- let mut struct_levels = vec![];
- struct_array
- .columns()
- .into_iter()
- .zip(struct_fields)
- .for_each(|(child_array, child_field)| {
- let mut levels =
- struct_level.calculate_array_levels(child_array, child_field);
- struct_levels.append(&mut levels);
- });
- struct_levels
- }
- DataType::Union(_) => unimplemented!(),
- DataType::Dictionary(_, _) => {
- // Need to check for these cases not implemented in C++:
- // - "Writing DictionaryArray with nested dictionary type not yet supported"
- // - "Writing DictionaryArray with null encoded in dictionary type not yet supported"
- // vec![self.get_primitive_def_levels(array, field, array_mask)]
- vec![self.calculate_child_levels(
- array_offsets,
- array_mask,
- LevelType::Primitive(field.is_nullable()),
- )]
- }
- }
- }
-
- /// Calculate child/leaf array levels.
- ///
- /// The algorithm works by incrementing definitions of array values based on whether:
- /// - a value is optional or required (is_nullable)
- /// - a list value is repeated + optional or required (is_list)
- ///
- /// A record batch always starts at a populated definition = level 0.
- /// When a batch only has a primitive, i.e. `<batch<primitive[a]>>, column `a`
- /// can only have a maximum level of 1 if it is not null.
- /// If it is not null, we increment by 1, such that the null slots will = level 1.
- /// The above applies to types that have no repetition (anything not a list or map).
- ///
- /// If a batch has lists, then we increment by up to 2 levels:
- /// - 1 level for the list (repeated)
- /// - 1 level if the list itself is nullable (optional)
- ///
- /// A list's child then gets incremented using the above rules.
- ///
- /// *Exceptions*
- ///
- /// There are 2 exceptions from the above rules:
- ///
- /// 1. When at the root of the schema: We always increment the
- /// level regardless of whether the child is nullable or not. If we do not do
- /// this, we could have a non-nullable array having a definition of 0.
- ///
- /// 2. List parent, non-list child: We always increment the level in this case,
- /// regardless of whether the child is nullable or not.
- ///
- /// *Examples*
- ///
- /// A batch with only a primitive that's non-nullable. `<primitive[required]>`:
- /// * We don't increment the definition level as the array is not optional.
- /// * This would leave us with a definition of 0, so the first exception applies.
- /// * The definition level becomes 1.
- ///
- /// A batch with only a primitive that's nullable. `<primitive[optional]>`:
- /// * The definition level becomes 1, as we increment it once.
- ///
- /// A batch with a single non-nullable list (both list and child not null):
- /// * We calculate the level twice, for the list, and for the child.
- /// * At the list, the level becomes 1, where 0 indicates that the list is
- /// empty, and 1 says it's not (determined through offsets).
- /// * At the primitive level, the second exception applies. The level becomes 2.
- fn calculate_child_levels(
- &self,
- // we use 64-bit offsets to also accommodate large arrays
- array_offsets: Vec<i64>,
- array_mask: Vec<bool>,
- level_type: LevelType,
- ) -> Self {
- let min_len = *(array_offsets.last().unwrap()) as usize;
- let mut definition = Vec::with_capacity(min_len);
- let mut repetition = Vec::with_capacity(min_len);
- let mut merged_array_mask = Vec::with_capacity(min_len);
-
- let max_definition = match (self.level_type, level_type) {
- (LevelType::Root, LevelType::Struct(is_nullable)) => {
- // If the struct is non-nullable, its def level doesn't increment
- is_nullable as i16
- }
- (LevelType::Root, _) => 1,
- (_, LevelType::Root) => {
- unreachable!("Cannot have a root as a child")
- }
- (LevelType::List(_), _) => {
- self.max_definition + 1 + level_type.level_increment()
- }
- (LevelType::Struct(_), _) => {
- self.max_definition + level_type.level_increment()
- }
- (_, LevelType::List(is_nullable)) => {
- // if the child is a list, even if its parent is a root
- self.max_definition + 1 + is_nullable as i16
- }
- (LevelType::Primitive(_), _) => {
- unreachable!("Cannot have a primitive parent for any type")
- }
- };
-
- match (self.level_type, level_type) {
- (LevelType::List(_), LevelType::List(is_nullable)) => {
- // parent is a list or descendant of a list, and child is a list
- let reps = self.repetition.clone().unwrap();
- // Calculate the 2 list hierarchy definitions in advance
- // List is not empty, but null
- let l2 = max_definition - is_nullable as i16;
- // List is not empty, and not null
- let l3 = max_definition;
-
- let mut nulls_seen = 0;
-
- self.array_offsets.windows(2).for_each(|w| {
- let start = w[0] as usize;
- let end = w[1] as usize;
- let parent_len = end - start;
-
- if parent_len == 0 {
- // If the parent length is 0, there won't be a slot for the child
- let index = start + nulls_seen - self.offset;
- definition.push(self.definition[index]);
- repetition.push(0);
- merged_array_mask.push(self.array_mask[index]);
- nulls_seen += 1;
- } else {
- (start..end).for_each(|parent_index| {
- let index = parent_index + nulls_seen - self.offset;
- let parent_index = parent_index - self.offset;
-
- // parent is either defined at this level, or earlier
- let parent_def = self.definition[index];
- let parent_rep = reps[index];
- let parent_mask = self.array_mask[index];
-
- // valid parent, index into children
- let child_start = array_offsets[parent_index] as usize;
- let child_end = array_offsets[parent_index + 1] as usize;
- let child_len = child_end - child_start;
- let child_mask = array_mask[parent_index];
- let merged_mask = parent_mask && child_mask;
-
- if child_len == 0 {
- definition.push(parent_def);
- repetition.push(parent_rep);
- merged_array_mask.push(merged_mask);
- } else {
- (child_start..child_end).for_each(|child_index| {
- let rep = match (
- parent_index == start,
- child_index == child_start,
- ) {
- (true, true) => parent_rep,
- (true, false) => parent_rep + 2,
- (false, true) => parent_rep,
- (false, false) => parent_rep + 1,
- };
-
- definition.push(if !parent_mask {
- parent_def
- } else if child_mask {
- l3
- } else {
- l2
- });
- repetition.push(rep);
- merged_array_mask.push(merged_mask);
- });
- }
- });
- }
- });
-
- debug_assert_eq!(definition.len(), merged_array_mask.len());
-
- let offset = *array_offsets.first().unwrap() as usize;
- let length = *array_offsets.last().unwrap() as usize - offset;
-
- Self {
- definition,
- repetition: Some(repetition),
- array_offsets,
- array_mask: merged_array_mask,
- max_definition,
- level_type,
- offset: offset + self.offset,
- length,
- }
- }
- (LevelType::List(_), _) => {
- // List and primitive (or struct).
- // The list can have more values than the primitive, indicating that there
- // are slots where the list is empty. We use a counter to track this behaviour.
- let mut nulls_seen = 0;
-
- // let child_max_definition = list_max_definition + is_nullable as i16;
- // child values are a function of parent list offsets
- let reps = self.repetition.as_deref().unwrap();
- self.array_offsets.windows(2).for_each(|w| {
- let start = w[0] as usize;
- let end = w[1] as usize;
- let parent_len = end - start;
-
- if parent_len == 0 {
- let index = start + nulls_seen - self.offset;
- definition.push(self.definition[index]);
- repetition.push(reps[index]);
- merged_array_mask.push(self.array_mask[index]);
- nulls_seen += 1;
- } else {
- // iterate through the array, adjusting child definitions for nulls
- (start..end).for_each(|child_index| {
- let index = child_index + nulls_seen - self.offset;
- let child_mask = array_mask[child_index - self.offset];
- let parent_mask = self.array_mask[index];
- let parent_def = self.definition[index];
-
- if !parent_mask || parent_def < self.max_definition {
- definition.push(parent_def);
- repetition.push(reps[index]);
- merged_array_mask.push(parent_mask);
- } else {
- definition.push(max_definition - !child_mask as i16);
- repetition.push(reps[index]);
- merged_array_mask.push(child_mask);
- }
- });
- }
- });
-
- debug_assert_eq!(definition.len(), merged_array_mask.len());
-
- let offset = *array_offsets.first().unwrap() as usize;
- let length = *array_offsets.last().unwrap() as usize - offset;
-
- Self {
- definition,
- repetition: Some(repetition),
- array_offsets: self.array_offsets.clone(),
- array_mask: merged_array_mask,
- max_definition,
- level_type,
- offset: offset + self.offset,
- length,
- }
- }
- (_, LevelType::List(is_nullable)) => {
- // Encountering a list for the first time.
- // Calculate the 2 list hierarchy definitions in advance
-
- // List is not empty, but null (if nullable)
- let l2 = max_definition - is_nullable as i16;
- // List is not empty, and not null
- let l3 = max_definition;
-
- self.definition
- .iter()
- .enumerate()
- .for_each(|(parent_index, def)| {
- let child_from = array_offsets[parent_index];
- let child_to = array_offsets[parent_index + 1];
- let child_len = child_to - child_from;
- let child_mask = array_mask[parent_index];
- let parent_mask = self.array_mask[parent_index];
-
- match (parent_mask, child_len) {
- (true, 0) => {
- // empty slot that is valid, i.e. {"parent": {"child": [] } }
- definition.push(if child_mask { l3 } else { l2 });
- repetition.push(0);
- merged_array_mask.push(child_mask);
- }
- (false, 0) => {
- definition.push(*def);
- repetition.push(0);
- merged_array_mask.push(child_mask);
- }
- (true, _) => {
- (child_from..child_to).for_each(|child_index| {
- definition.push(if child_mask { l3 } else { l2 });
- // mark the first child slot as 0, and the next as 1
- repetition.push(if child_index == child_from {
- 0
- } else {
- 1
- });
- merged_array_mask.push(child_mask);
- });
- }
- (false, _) => {
- (child_from..child_to).for_each(|child_index| {
- definition.push(*def);
- // mark the first child slot as 0, and the next as 1
- repetition.push(if child_index == child_from {
- 0
- } else {
- 1
- });
- merged_array_mask.push(false);
- });
- }
- }
- });
-
- debug_assert_eq!(definition.len(), merged_array_mask.len());
-
- let offset = *array_offsets.first().unwrap() as usize;
- let length = *array_offsets.last().unwrap() as usize - offset;
-
- Self {
- definition,
- repetition: Some(repetition),
- array_offsets,
- array_mask: merged_array_mask,
- max_definition,
- level_type,
- offset,
- length,
- }
- }
- (_, _) => {
- self.definition
- .iter()
- .zip(array_mask.into_iter().zip(&self.array_mask))
- .for_each(|(current_def, (child_mask, parent_mask))| {
- merged_array_mask.push(*parent_mask && child_mask);
- match (parent_mask, child_mask) {
- (true, true) => {
- definition.push(max_definition);
- }
- (true, false) => {
- // The child is only legally null if its array is nullable.
- // Thus parent's max_definition is lower
- definition.push(if *current_def <= self.max_definition {
- *current_def
- } else {
- self.max_definition
- });
- }
- // if the parent was false, retain its definitions
- (false, _) => {
- definition.push(*current_def);
- }
- }
- });
-
- debug_assert_eq!(definition.len(), merged_array_mask.len());
-
- Self {
- definition,
- repetition: self.repetition.clone(), // it's None
- array_offsets,
- array_mask: merged_array_mask,
- max_definition,
- level_type,
- // Inherit parent offset and length
- offset: self.offset,
- length: self.length,
- }
- }
- }
- }
-
- /// Get the offsets of an array as 64-bit values, and validity masks as booleans
- /// - Primitive, binary and struct arrays' offsets will be a sequence, masks obtained
- /// from validity bitmap
- /// - List array offsets will be the value offsets, masks are computed from offsets
- fn get_array_offsets_and_masks(
- array: &ArrayRef,
- offset: usize,
- len: usize,
- ) -> (Vec<i64>, Vec<bool>) {
- match array.data_type() {
- DataType::Null
- | DataType::Boolean
- | DataType::Int8
- | DataType::Int16
- | DataType::Int32
- | DataType::Int64
- | DataType::UInt8
- | DataType::UInt16
- | DataType::UInt32
- | DataType::UInt64
- | DataType::Float16
- | DataType::Float32
- | DataType::Float64
- | DataType::Timestamp(_, _)
- | DataType::Date32
- | DataType::Date64
- | DataType::Time32(_)
- | DataType::Time64(_)
- | DataType::Duration(_)
- | DataType::Interval(_)
- | DataType::Binary
- | DataType::LargeBinary
- | DataType::Utf8
- | DataType::LargeUtf8
- | DataType::Struct(_)
- | DataType::Dictionary(_, _)
- | DataType::Decimal(_, _) => {
- let array_mask = match array.data().null_buffer() {
- Some(buf) => get_bool_array_slice(buf, array.offset() + offset, len),
- None => vec![true; len],
- };
- ((0..=(len as i64)).collect(), array_mask)
- }
- DataType::List(_) => {
- let data = array.data();
- let offsets = unsafe { data.buffers()[0].typed_data::<i32>() };
- let offsets = offsets
- .to_vec()
- .into_iter()
- .skip(offset)
- .take(len + 1)
- .map(|v| v as i64)
- .collect::<Vec<i64>>();
- let array_mask = match array.data().null_buffer() {
- Some(buf) => get_bool_array_slice(buf, array.offset() + offset, len),
- None => vec![true; len],
- };
- (offsets, array_mask)
- }
- DataType::LargeList(_) => {
- let offsets = unsafe { array.data().buffers()[0].typed_data::<i64>() }
- .iter()
- .skip(offset)
- .take(len + 1)
- .copied()
- .collect();
- let array_mask = match array.data().null_buffer() {
- Some(buf) => get_bool_array_slice(buf, array.offset() + offset, len),
- None => vec![true; len],
- };
- (offsets, array_mask)
- }
- DataType::FixedSizeBinary(value_len) => {
- let array_mask = match array.data().null_buffer() {
- Some(buf) => get_bool_array_slice(buf, array.offset() + offset, len),
- None => vec![true; len],
- };
- let value_len = *value_len as i64;
- (
- (0..=(len as i64)).map(|v| v * value_len).collect(),
- array_mask,
- )
- }
- DataType::FixedSizeList(_, _) | DataType::Union(_) => {
- unimplemented!("Getting offsets not yet implemented")
- }
- }
- }
-
- /// Given a level's information, calculate the offsets required to index an array correctly.
- pub(crate) fn filter_array_indices(&self) -> Vec<usize> {
- // happy path if not dealing with lists
- let is_nullable = match self.level_type {
- LevelType::Primitive(is_nullable) => is_nullable,
- _ => panic!(
- "Cannot filter indices on a non-primitive array, found {:?}",
- self.level_type
- ),
- };
- if self.repetition.is_none() {
- return self
- .definition
- .iter()
- .enumerate()
- .filter_map(|(i, def)| {
- if *def == self.max_definition {
- Some(i)
- } else {
- None
- }
- })
- .collect();
- }
- let mut filtered = vec![];
- // remove slots that are false from definition_mask
- let mut index = 0;
- self.definition.iter().for_each(|def| {
- if *def == self.max_definition {
- filtered.push(index);
- }
- if *def >= self.max_definition - is_nullable as i16 {
- index += 1;
- }
- });
- filtered
- }
-}
-
-/// Convert an Arrow buffer to a boolean array slice
-/// TODO: this was created for buffers, so might not work for bool array, might be slow too
-#[inline]
-fn get_bool_array_slice(
- buffer: &arrow::buffer::Buffer,
- offset: usize,
- len: usize,
-) -> Vec<bool> {
- let data = buffer.as_slice();
- (offset..(len + offset))
- .map(|i| arrow::util::bit_util::get_bit(data, i))
- .collect()
-}
-
-#[cfg(test)]
-mod tests {
- use super::*;
-
- use std::sync::Arc;
-
- use arrow::array::*;
- use arrow::buffer::Buffer;
- use arrow::datatypes::{Schema, ToByteSlice};
- use arrow::record_batch::RecordBatch;
-
- #[test]
- fn test_calculate_array_levels_twitter_example() {
- // based on the example at https://blog.twitter.com/engineering/en_us/a/2013/dremel-made-simple-with-parquet.html
- // [[a, b, c], [d, e, f, g]], [[h], [i,j]]
- let parent_levels = LevelInfo {
- definition: vec![0, 0],
- repetition: None,
- array_offsets: vec![0, 1, 2], // 2 records, root offsets always sequential
- array_mask: vec![true, true], // both lists defined
- max_definition: 0,
- level_type: LevelType::Root,
- offset: 0,
- length: 2,
- };
- // offset into array, each level1 has 2 values
- let array_offsets = vec![0, 2, 4];
- let array_mask = vec![true, true];
-
- // calculate level1 levels
- let levels = parent_levels.calculate_child_levels(
- array_offsets.clone(),
- array_mask,
- LevelType::List(false),
- );
- //
- let expected_levels = LevelInfo {
- definition: vec![1, 1, 1, 1],
- repetition: Some(vec![0, 1, 0, 1]),
- array_offsets,
- array_mask: vec![true, true, true, true],
- max_definition: 1,
- level_type: LevelType::List(false),
- offset: 0,
- length: 4,
- };
- // the separate asserts make it easier to see what's failing
- assert_eq!(&levels.definition, &expected_levels.definition);
- assert_eq!(&levels.repetition, &expected_levels.repetition);
- assert_eq!(&levels.array_mask, &expected_levels.array_mask);
- assert_eq!(&levels.array_offsets, &expected_levels.array_offsets);
- assert_eq!(&levels.max_definition, &expected_levels.max_definition);
- assert_eq!(&levels.level_type, &expected_levels.level_type);
- // this assert is to help if there are more variables added to the struct
- assert_eq!(&levels, &expected_levels);
-
- // level2
- let parent_levels = levels;
- let array_offsets = vec![0, 3, 7, 8, 10];
- let array_mask = vec![true, true, true, true];
- let levels = parent_levels.calculate_child_levels(
- array_offsets.clone(),
- array_mask,
- LevelType::List(false),
- );
- let expected_levels = LevelInfo {
- definition: vec![2, 2, 2, 2, 2, 2, 2, 2, 2, 2],
- repetition: Some(vec![0, 2, 2, 1, 2, 2, 2, 0, 1, 2]),
- array_offsets,
- array_mask: vec![true; 10],
- max_definition: 2,
- level_type: LevelType::List(false),
- offset: 0,
- length: 10,
- };
- assert_eq!(&levels.definition, &expected_levels.definition);
- assert_eq!(&levels.repetition, &expected_levels.repetition);
- assert_eq!(&levels.array_mask, &expected_levels.array_mask);
- assert_eq!(&levels.max_definition, &expected_levels.max_definition);
- assert_eq!(&levels.array_offsets, &expected_levels.array_offsets);
- assert_eq!(&levels.level_type, &expected_levels.level_type);
- assert_eq!(&levels, &expected_levels);
- }
-
- #[test]
- fn test_calculate_one_level_1() {
- // This test calculates the levels for a non-null primitive array
- let parent_levels = LevelInfo {
- definition: vec![0; 10],
- repetition: None,
- array_offsets: (0..=10).collect(),
- array_mask: vec![true; 10],
- max_definition: 0,
- level_type: LevelType::Root,
- offset: 0,
- length: 10,
- };
- let array_offsets: Vec<i64> = (0..=10).collect();
- let array_mask = vec![true; 10];
-
- let levels = parent_levels.calculate_child_levels(
- array_offsets.clone(),
- array_mask.clone(),
- LevelType::Primitive(false),
- );
- let expected_levels = LevelInfo {
- definition: vec![1; 10],
- repetition: None,
- array_offsets,
- array_mask,
- max_definition: 1,
- level_type: LevelType::Primitive(false),
- offset: 0,
- length: 10,
- };
- assert_eq!(&levels, &expected_levels);
- }
-
- #[test]
- fn test_calculate_one_level_2() {
- // This test calculates the levels for a non-null primitive array
- let parent_levels = LevelInfo {
- definition: vec![0; 5],
- repetition: None,
- array_offsets: (0..=5).collect(),
- array_mask: vec![true, true, true, true, true],
- max_definition: 0,
- level_type: LevelType::Root,
- offset: 0,
- length: 5,
- };
- let array_offsets: Vec<i64> = (0..=5).collect();
- let array_mask = vec![true, false, true, true, false];
-
- let levels = parent_levels.calculate_child_levels(
- array_offsets.clone(),
- array_mask.clone(),
- LevelType::Primitive(true),
- );
- let expected_levels = LevelInfo {
- definition: vec![1, 0, 1, 1, 0],
- repetition: None,
- array_offsets,
- array_mask,
- max_definition: 1,
- level_type: LevelType::Primitive(true),
- offset: 0,
- length: 5,
- };
- assert_eq!(&levels, &expected_levels);
- }
-
- #[test]
- fn test_calculate_array_levels_1() {
- // if all array values are defined (e.g. batch<list<_>>)
- // [[0], [1], [2], [3], [4]]
- let parent_levels = LevelInfo {
- definition: vec![0; 5],
- repetition: None,
- array_offsets: vec![0, 1, 2, 3, 4, 5],
- array_mask: vec![true, true, true, true, true],
- max_definition: 0,
- level_type: LevelType::Root,
- offset: 0,
- length: 5,
- };
- let array_offsets = vec![0, 2, 2, 4, 8, 11];
- let array_mask = vec![true, false, true, true, true];
-
- let levels = parent_levels.calculate_child_levels(
- array_offsets.clone(),
- array_mask,
- LevelType::List(true),
- );
- // array: [[0, 0], _1_, [2, 2], [3, 3, 3, 3], [4, 4, 4]]
- // all values are defined as we do not have nulls on the root (batch)
- // repetition:
- // 0: 0, 1
- // 1:
- // 2: 0, 1
- // 3: 0, 1, 1, 1
- // 4: 0, 1, 1
- let expected_levels = LevelInfo {
- // The levels are normally 2 because we:
- // - Calculate the level at the list
- // - Calculate the level at the list's child
- // We do not do this in these tests, thus the levels are 1 less.
- definition: vec![1, 1, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1],
- repetition: Some(vec![0, 1, 0, 0, 1, 0, 1, 1, 1, 0, 1, 1]),
- array_offsets,
- array_mask: vec![
- true, true, false, true, true, true, true, true, true, true, true, true,
- ],
- max_definition: 1,
- level_type: LevelType::List(true),
- offset: 0,
- length: 11, // the child has 11 slots
- };
- assert_eq!(&levels.definition, &expected_levels.definition);
- assert_eq!(&levels.repetition, &expected_levels.repetition);
- assert_eq!(&levels.array_offsets, &expected_levels.array_offsets);
- assert_eq!(&levels.max_definition, &expected_levels.max_definition);
- assert_eq!(&levels.level_type, &expected_levels.level_type);
- assert_eq!(&levels, &expected_levels);
- }
-
- #[test]
- fn test_calculate_array_levels_2() {
- // If some values are null
- //
- // This emulates an array in the form: <struct<list<?>>
- // with values:
- // - 0: [0, 1], but is null because of the struct
- // - 1: []
- // - 2: [2, 3], but is null because of the struct
- // - 3: [4, 5, 6, 7]
- // - 4: [8, 9, 10]
- //
- // If the first values of a list are null due to a parent, we have to still account for them
- // while indexing, because they would affect the way the child is indexed
- // i.e. in the above example, we have to know that [0, 1] has to be skipped
- let parent_levels = LevelInfo {
- definition: vec![0, 1, 0, 1, 1],
- repetition: None,
- array_offsets: vec![0, 1, 2, 3, 4, 5],
- array_mask: vec![false, true, false, true, true],
- max_definition: 1,
- level_type: LevelType::Struct(true),
- offset: 0,
- length: 5,
- };
- let array_offsets = vec![0, 2, 2, 4, 8, 11];
- let array_mask = vec![true, false, true, true, true];
-
- let levels = parent_levels.calculate_child_levels(
- array_offsets.clone(),
- array_mask,
- LevelType::List(true),
- );
- let expected_levels = LevelInfo {
- // 0 1 [2] are 0 (not defined at level 1)
- // [2] is 1, but has 0 slots so is not populated (defined at level 1 only)
- // 2 3 [4] are 0
- // 4 5 6 7 [8] are 1 (defined at level 1 only)
- // 8 9 10 [11] are 2 (defined at both levels)
- definition: vec![0, 0, 1, 0, 0, 2, 2, 2, 2, 2, 2, 2],
- repetition: Some(vec![0, 1, 0, 0, 1, 0, 1, 1, 1, 0, 1, 1]),
- array_offsets,
- array_mask: vec![
- false, false, false, false, false, true, true, true, true, true, true,
- true,
- ],
- max_definition: 2,
- level_type: LevelType::List(true),
- offset: 0,
- length: 11,
- };
- assert_eq!(&levels.definition, &expected_levels.definition);
- assert_eq!(&levels.repetition, &expected_levels.repetition);
- assert_eq!(&levels.array_offsets, &expected_levels.array_offsets);
- assert_eq!(&levels.max_definition, &expected_levels.max_definition);
- assert_eq!(&levels.level_type, &expected_levels.level_type);
- assert_eq!(&levels, &expected_levels);
-
- // nested lists (using previous test)
- let nested_parent_levels = levels;
- let array_offsets = vec![0, 2, 4, 6, 8, 10, 12, 14, 16, 18, 20, 22];
- let array_mask = vec![
- true, true, true, true, true, true, true, true, true, true, true,
- ];
- let levels = nested_parent_levels.calculate_child_levels(
- array_offsets.clone(),
- array_mask,
- LevelType::List(true),
- );
- let expected_levels = LevelInfo {
- // (def: 0) 0 1 [2] are 0 (take parent)
- // (def: 0) 2 3 [4] are 0 (take parent)
- // (def: 0) 4 5 [6] are 0 (take parent)
- // (def: 0) 6 7 [8] are 0 (take parent)
- // (def: 1) 8 9 [10] are 1 (take parent)
- // (def: 1) 10 11 [12] are 1 (take parent)
- // (def: 1) 12 23 [14] are 1 (take parent)
- // (def: 1) 14 15 [16] are 1 (take parent)
- // (def: 2) 16 17 [18] are 2 (defined at all levels)
- // (def: 2) 18 19 [20] are 2 (defined at all levels)
- // (def: 2) 20 21 [22] are 2 (defined at all levels)
- //
- // 0 1 [2] are 0 (not defined at level 1)
- // [2] is 1, but has 0 slots so is not populated (defined at level 1 only)
- // 2 3 [4] are 0
- // 4 5 6 7 [8] are 1 (defined at level 1 only)
- // 8 9 10 [11] are 2 (defined at both levels)
- //
- // 0: [[100, 101], [102, 103]]
- // 1: []
- // 2: [[104, 105], [106, 107]]
- // 3: [[108, 109], [110, 111], [112, 113], [114, 115]]
- // 4: [[116, 117], [118, 119], [120, 121]]
- definition: vec![
- 0, 0, 0, 0, 1, 0, 0, 0, 0, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4,
- ],
- repetition: Some(vec![
- 0, 2, 1, 2, 0, 0, 2, 1, 2, 0, 2, 1, 2, 1, 2, 1, 2, 0, 2, 1, 2, 1, 2,
- ]),
- array_offsets,
- array_mask: vec![
- false, false, false, false, false, false, false, false, false, true,
- true, true, true, true, true, true, true, true, true, true, true, true,
- true,
- ],
- max_definition: 4,
- level_type: LevelType::List(true),
- offset: 0,
- length: 22,
- };
- assert_eq!(&levels.definition, &expected_levels.definition);
- assert_eq!(&levels.repetition, &expected_levels.repetition);
- assert_eq!(&levels.array_offsets, &expected_levels.array_offsets);
- assert_eq!(&levels.array_mask, &expected_levels.array_mask);
- assert_eq!(&levels.max_definition, &expected_levels.max_definition);
- assert_eq!(&levels.level_type, &expected_levels.level_type);
- assert_eq!(&levels, &expected_levels);
- }
-
- #[test]
- fn test_calculate_array_levels_nested_list() {
- // if all array values are defined (e.g. batch<list<_>>)
- // The array at this level looks like:
- // 0: [a]
- // 1: [a]
- // 2: [a]
- // 3: [a]
- let parent_levels = LevelInfo {
- definition: vec![1, 1, 1, 1],
- repetition: None,
- array_offsets: vec![0, 1, 2, 3, 4],
- array_mask: vec![true, true, true, true],
- max_definition: 1,
- level_type: LevelType::Struct(true),
- offset: 0,
- length: 4,
- };
- // 0: null ([], but mask is false, so it's not just an empty list)
- // 1: [1, 2, 3]
- // 2: [4, 5]
- // 3: [6, 7]
- let array_offsets = vec![0, 1, 4, 6, 8];
- let array_mask = vec![false, true, true, true];
-
- let levels = parent_levels.calculate_child_levels(
- array_offsets.clone(),
- array_mask,
- LevelType::List(true),
- );
- // 0: [null], level 1 is defined, but not 2
- // 1: [1, 2, 3]
- // 2: [4, 5]
- // 3: [6, 7]
- let expected_levels = LevelInfo {
- definition: vec![1, 2, 2, 2, 2, 2, 2, 2],
- repetition: Some(vec![0, 0, 1, 1, 0, 1, 0, 1]),
- array_offsets,
- array_mask: vec![false, true, true, true, true, true, true, true],
- max_definition: 2,
- level_type: LevelType::List(true),
- offset: 0,
- length: 8,
- };
- assert_eq!(&levels.definition, &expected_levels.definition);
- assert_eq!(&levels.repetition, &expected_levels.repetition);
- assert_eq!(&levels.array_offsets, &expected_levels.array_offsets);
- assert_eq!(&levels.max_definition, &expected_levels.max_definition);
- assert_eq!(&levels.level_type, &expected_levels.level_type);
- assert_eq!(&levels, &expected_levels);
-
- // nested lists (using previous test)
- let nested_parent_levels = levels;
- // 0: [null] (was a populated null slot at the parent)
- // 1: [201]
- // 2: [202, 203]
- // 3: null ([])
- // 4: [204, 205, 206]
- // 5: [207, 208, 209, 210]
- // 6: [] (tests a non-null empty list slot)
- // 7: [211, 212, 213, 214, 215]
- let array_offsets = vec![0, 1, 2, 4, 4, 7, 11, 11, 16];
- // logically, the fist slot of the mask is false
- let array_mask = vec![true, true, true, false, true, true, true, true];
- let levels = nested_parent_levels.calculate_child_levels(
- array_offsets.clone(),
- array_mask,
- LevelType::List(true),
- );
- // We have 7 array values, and at least 15 primitives (from array_offsets)
- // 0: (-)[null], parent was null, no value populated here
- // 1: (0)[201], (1)[202, 203], (2)[[null]]
- // 2: (3)[204, 205, 206], (4)[207, 208, 209, 210]
- // 3: (5)[[]], (6)[211, 212, 213, 214, 215]
- //
- // In a JSON syntax with the schema: <struct<list<list<primitive>>>>, this translates into:
- // 0: {"struct": [ null ]}
- // 1: {"struct": [ [201], [202, 203], [] ]}
- // 2: {"struct": [ [204, 205, 206], [207, 208, 209, 210] ]}
- // 3: {"struct": [ [], [211, 212, 213, 214, 215] ]}
- let expected_levels = LevelInfo {
- definition: vec![1, 4, 4, 4, 2, 4, 4, 4, 4, 4, 4, 4, 2, 4, 4, 4, 4, 4],
- repetition: Some(vec![0, 0, 1, 2, 1, 0, 2, 2, 1, 2, 2, 2, 0, 1, 2, 2, 2, 2]),
- array_mask: vec![
- false, true, true, true, false, true, true, true, true, true, true, true,
- true, true, true, true, true, true,
- ],
- array_offsets,
- max_definition: 4,
- level_type: LevelType::List(true),
- offset: 0,
- length: 16,
- };
- assert_eq!(&levels.definition, &expected_levels.definition);
- assert_eq!(&levels.repetition, &expected_levels.repetition);
- assert_eq!(&levels.array_offsets, &expected_levels.array_offsets);
- assert_eq!(&levels.array_mask, &expected_levels.array_mask);
- assert_eq!(&levels.max_definition, &expected_levels.max_definition);
- assert_eq!(&levels.level_type, &expected_levels.level_type);
- assert_eq!(&levels, &expected_levels);
- }
-
- #[test]
- fn test_calculate_nested_struct_levels() {
- // tests a <struct[a]<struct[b]<int[c]>>
- // array:
- // - {a: {b: {c: 1}}}
- // - {a: {b: {c: null}}}
- // - {a: {b: {c: 3}}}
- // - {a: {b: null}}
- // - {a: null}}
- // - {a: {b: {c: 6}}}
- let a_levels = LevelInfo {
- definition: vec![1, 1, 1, 1, 0, 1],
- repetition: None,
- array_offsets: (0..=6).collect(),
- array_mask: vec![true, true, true, true, false, true],
- max_definition: 1,
- level_type: LevelType::Struct(true),
- offset: 0,
- length: 6,
- };
- // b's offset and mask
- let b_offsets: Vec<i64> = (0..=6).collect();
- let b_mask = vec![true, true, true, false, false, true];
- // b's expected levels
- let b_expected_levels = LevelInfo {
- definition: vec![2, 2, 2, 1, 0, 2],
- repetition: None,
- array_offsets: (0..=6).collect(),
- array_mask: vec![true, true, true, false, false, true],
- max_definition: 2,
- level_type: LevelType::Struct(true),
- offset: 0,
- length: 6,
- };
- let b_levels = a_levels.calculate_child_levels(
- b_offsets.clone(),
- b_mask,
- LevelType::Struct(true),
- );
- assert_eq!(&b_expected_levels, &b_levels);
-
- // c's offset and mask
- let c_offsets = b_offsets;
- let c_mask = vec![true, false, true, false, false, true];
- // c's expected levels
- let c_expected_levels = LevelInfo {
- definition: vec![3, 2, 3, 1, 0, 3],
- repetition: None,
- array_offsets: c_offsets.clone(),
- array_mask: vec![true, false, true, false, false, true],
- max_definition: 3,
- level_type: LevelType::Struct(true),
- offset: 0,
- length: 6,
- };
- let c_levels =
- b_levels.calculate_child_levels(c_offsets, c_mask, LevelType::Struct(true));
- assert_eq!(&c_expected_levels, &c_levels);
- }
-
- #[test]
- fn list_single_column() {
- // this tests the level generation from the arrow_writer equivalent test
-
- let a_values = Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10]);
- let a_value_offsets =
- arrow::buffer::Buffer::from(&[0, 1, 3, 3, 6, 10].to_byte_slice());
- let a_list_type =
- DataType::List(Box::new(Field::new("item", DataType::Int32, true)));
- let a_list_data = ArrayData::builder(a_list_type.clone())
- .len(5)
- .add_buffer(a_value_offsets)
- .null_bit_buffer(Buffer::from(vec![0b00011011]))
- .add_child_data(a_values.data().clone())
- .build();
-
- assert_eq!(a_list_data.null_count(), 1);
-
- let a = ListArray::from(a_list_data);
- let values = Arc::new(a);
-
- let schema = Schema::new(vec![Field::new("item", a_list_type, true)]);
-
- let batch = RecordBatch::try_new(Arc::new(schema), vec![values]).unwrap();
-
- let expected_batch_level = LevelInfo {
- definition: vec![0; 2],
- repetition: None,
- array_offsets: (0..=2).collect(),
- array_mask: vec![true, true],
- max_definition: 0,
- level_type: LevelType::Root,
- offset: 2,
- length: 2,
- };
-
- let batch_level = LevelInfo::new(2, 2);
- assert_eq!(&batch_level, &expected_batch_level);
-
- // calculate the list's level
- let mut levels = vec![];
- batch
- .columns()
- .iter()
- .zip(batch.schema().fields())
- .for_each(|(array, field)| {
- let mut array_levels = batch_level.calculate_array_levels(array, field);
- levels.append(&mut array_levels);
- });
- assert_eq!(levels.len(), 1);
-
- let list_level = levels.get(0).unwrap();
-
- let expected_level = LevelInfo {
- definition: vec![0, 3, 3, 3],
- repetition: Some(vec![0, 0, 1, 1]),
- array_offsets: vec![3, 3, 6],
- array_mask: vec![false, true, true, true],
- max_definition: 3,
- level_type: LevelType::Primitive(true),
- offset: 3,
- length: 3,
- };
- assert_eq!(&list_level.definition, &expected_level.definition);
- assert_eq!(&list_level.repetition, &expected_level.repetition);
- assert_eq!(&list_level.array_offsets, &expected_level.array_offsets);
- assert_eq!(&list_level.array_mask, &expected_level.array_mask);
- assert_eq!(&list_level.max_definition, &expected_level.max_definition);
- assert_eq!(&list_level.level_type, &expected_level.level_type);
- assert_eq!(list_level, &expected_level);
- }
-
- #[test]
- fn mixed_struct_list() {
- // this tests the level generation from the equivalent arrow_writer_complex test
-
- // define schema
- let struct_field_d = Field::new("d", DataType::Float64, true);
- let struct_field_f = Field::new("f", DataType::Float32, true);
- let struct_field_g = Field::new(
- "g",
- DataType::List(Box::new(Field::new("items", DataType::Int16, false))),
- false,
- );
- let struct_field_e = Field::new(
- "e",
- DataType::Struct(vec![struct_field_f.clone(), struct_field_g.clone()]),
- true,
- );
- let schema = Schema::new(vec![
- Field::new("a", DataType::Int32, false),
- Field::new("b", DataType::Int32, true),
- Field::new(
- "c",
- DataType::Struct(vec![struct_field_d.clone(), struct_field_e.clone()]),
- true, // https://github.com/apache/arrow-rs/issues/245
- ),
- ]);
-
- // create some data
- let a = Int32Array::from(vec![1, 2, 3, 4, 5]);
- let b = Int32Array::from(vec![Some(1), None, None, Some(4), Some(5)]);
- let d = Float64Array::from(vec![None, None, None, Some(1.0), None]);
- let f = Float32Array::from(vec![Some(0.0), None, Some(333.3), None, Some(5.25)]);
-
- let g_value = Int16Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10]);
-
- // Construct a buffer for value offsets, for the nested array:
- // [[1], [2, 3], null, [4, 5, 6], [7, 8, 9, 10]]
- let g_value_offsets =
- arrow::buffer::Buffer::from(&[0, 1, 3, 3, 6, 10].to_byte_slice());
-
- // Construct a list array from the above two
- let g_list_data = ArrayData::builder(struct_field_g.data_type().clone())
- .len(5)
- .add_buffer(g_value_offsets)
- .add_child_data(g_value.data().clone())
- .build();
- let g = ListArray::from(g_list_data);
-
- let e = StructArray::from(vec![
- (struct_field_f, Arc::new(f) as ArrayRef),
- (struct_field_g, Arc::new(g) as ArrayRef),
- ]);
-
- let c = StructArray::from(vec![
- (struct_field_d, Arc::new(d) as ArrayRef),
- (struct_field_e, Arc::new(e) as ArrayRef),
- ]);
-
- // build a record batch
- let batch = RecordBatch::try_new(
- Arc::new(schema),
- vec![Arc::new(a), Arc::new(b), Arc::new(c)],
- )
- .unwrap();
-
- //////////////////////////////////////////////
- let expected_batch_level = LevelInfo {
- definition: vec![0; 5],
- repetition: None,
- array_offsets: (0..=5).collect(),
- array_mask: vec![true, true, true, true, true],
- max_definition: 0,
- level_type: LevelType::Root,
- offset: 0,
- length: 5,
- };
-
- let batch_level = LevelInfo::new(0, 5);
- assert_eq!(&batch_level, &expected_batch_level);
-
- // calculate the list's level
- let mut levels = vec![];
- batch
- .columns()
- .iter()
- .zip(batch.schema().fields())
- .for_each(|(array, field)| {
- let mut array_levels = batch_level.calculate_array_levels(array, field);
- levels.append(&mut array_levels);
- });
- assert_eq!(levels.len(), 5);
-
- // test "a" levels
- let list_level = levels.get(0).unwrap();
-
- let expected_level = LevelInfo {
- definition: vec![1, 1, 1, 1, 1],
- repetition: None,
- array_offsets: vec![0, 1, 2, 3, 4, 5],
- array_mask: vec![true, true, true, true, true],
- max_definition: 1,
- level_type: LevelType::Primitive(false),
- offset: 0,
- length: 5,
- };
- assert_eq!(list_level, &expected_level);
-
- // test "b" levels
- let list_level = levels.get(1).unwrap();
-
- let expected_level = LevelInfo {
- definition: vec![1, 0, 0, 1, 1],
- repetition: None,
- array_offsets: vec![0, 1, 2, 3, 4, 5],
- array_mask: vec![true, false, false, true, true],
- max_definition: 1,
- level_type: LevelType::Primitive(true),
- offset: 0,
- length: 5,
- };
- assert_eq!(list_level, &expected_level);
-
- // test "d" levels
- let list_level = levels.get(2).unwrap();
-
- let expected_level = LevelInfo {
- definition: vec![1, 1, 1, 2, 1],
- repetition: None,
- array_offsets: vec![0, 1, 2, 3, 4, 5],
- array_mask: vec![false, false, false, true, false],
- max_definition: 2,
- level_type: LevelType::Primitive(true),
- offset: 0,
- length: 5,
- };
- assert_eq!(list_level, &expected_level);
-
- // test "f" levels
- let list_level = levels.get(3).unwrap();
-
- let expected_level = LevelInfo {
- definition: vec![3, 2, 3, 2, 3],
- repetition: None,
- array_offsets: vec![0, 1, 2, 3, 4, 5],
- array_mask: vec![true, false, true, false, true],
- max_definition: 3,
- level_type: LevelType::Primitive(true),
- offset: 0,
- length: 5,
- };
- assert_eq!(list_level, &expected_level);
- }
-
- #[test]
- fn test_filter_array_indices() {
- let level = LevelInfo {
- definition: vec![3, 3, 3, 1, 3, 3, 3],
- repetition: Some(vec![0, 1, 1, 0, 0, 1, 1]),
- array_offsets: vec![0, 3, 3, 6],
- array_mask: vec![true, true, true, false, true, true, true],
- max_definition: 3,
- level_type: LevelType::Primitive(true),
- offset: 0,
- length: 6,
- };
-
- let expected = vec![0, 1, 2, 3, 4, 5];
- let filter = level.filter_array_indices();
- assert_eq!(expected, filter);
- }
-
- #[test]
- fn test_null_vs_nonnull_struct() {
- // define schema
- let offset_field = Field::new("offset", DataType::Int32, true);
- let schema = Schema::new(vec![Field::new(
- "some_nested_object",
- DataType::Struct(vec![offset_field.clone()]),
- false,
- )]);
-
- // create some data
- let offset = Int32Array::from(vec![1, 2, 3, 4, 5]);
-
- let some_nested_object =
- StructArray::from(vec![(offset_field, Arc::new(offset) as ArrayRef)]);
-
- // build a record batch
- let batch =
- RecordBatch::try_new(Arc::new(schema), vec![Arc::new(some_nested_object)])
- .unwrap();
-
- let batch_level = LevelInfo::new(0, batch.num_rows());
- let struct_null_level =
- batch_level.calculate_array_levels(batch.column(0), batch.schema().field(0));
-
- // create second batch
- // define schema
- let offset_field = Field::new("offset", DataType::Int32, true);
- let schema = Schema::new(vec![Field::new(
- "some_nested_object",
- DataType::Struct(vec![offset_field.clone()]),
- true,
- )]);
-
- // create some data
- let offset = Int32Array::from(vec![1, 2, 3, 4, 5]);
-
- let some_nested_object =
- StructArray::from(vec![(offset_field, Arc::new(offset) as ArrayRef)]);
-
- // build a record batch
- let batch =
- RecordBatch::try_new(Arc::new(schema), vec![Arc::new(some_nested_object)])
- .unwrap();
-
- let batch_level = LevelInfo::new(0, batch.num_rows());
- let struct_non_null_level =
- batch_level.calculate_array_levels(batch.column(0), batch.schema().field(0));
-
- // The 2 levels should not be the same
- if struct_non_null_level == struct_null_level {
- panic!("Levels should not be equal, to reflect the difference in struct nullness");
- }
- }
-}
diff --git a/parquet/src/arrow/mod.rs b/parquet/src/arrow/mod.rs
deleted file mode 100644
index afb1fdc..0000000
--- a/parquet/src/arrow/mod.rs
+++ /dev/null
@@ -1,70 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-//! [Apache Arrow](http://arrow.apache.org/) is a cross-language development platform for
-//! in-memory data.
-//!
-//! This mod provides API for converting between arrow and parquet.
-//!
-//! # Example of reading parquet file into arrow record batch
-//!
-//! ```rust, no_run
-//! use arrow::record_batch::RecordBatchReader;
-//! use parquet::file::reader::SerializedFileReader;
-//! use parquet::arrow::{ParquetFileArrowReader, ArrowReader};
-//! use std::sync::Arc;
-//! use std::fs::File;
-//!
-//! let file = File::open("parquet.file").unwrap();
-//! let file_reader = SerializedFileReader::new(file).unwrap();
-//! let mut arrow_reader = ParquetFileArrowReader::new(Arc::new(file_reader));
-//!
-//! println!("Converted arrow schema is: {}", arrow_reader.get_schema().unwrap());
-//! println!("Arrow schema after projection is: {}",
-//! arrow_reader.get_schema_by_columns(vec![2, 4, 6], true).unwrap());
-//!
-//! let mut record_batch_reader = arrow_reader.get_record_reader(2048).unwrap();
-//!
-//! for maybe_record_batch in record_batch_reader {
-//! let record_batch = maybe_record_batch.unwrap();
-//! if record_batch.num_rows() > 0 {
-//! println!("Read {} records.", record_batch.num_rows());
-//! } else {
-//! println!("End of file!");
-//! }
-//!}
-//! ```
-
-pub mod array_reader;
-pub mod arrow_array_reader;
-pub mod arrow_reader;
-pub mod arrow_writer;
-pub mod converter;
-pub(in crate::arrow) mod levels;
-pub(in crate::arrow) mod record_reader;
-pub mod schema;
-
-pub use self::arrow_reader::ArrowReader;
-pub use self::arrow_reader::ParquetFileArrowReader;
-pub use self::arrow_writer::ArrowWriter;
-pub use self::schema::{
- arrow_to_parquet_schema, parquet_to_arrow_schema, parquet_to_arrow_schema_by_columns,
- parquet_to_arrow_schema_by_root_columns,
-};
-
-/// Schema metadata key used to store serialized Arrow IPC schema
-pub const ARROW_SCHEMA_META_KEY: &str = "ARROW:schema";
diff --git a/parquet/src/arrow/record_reader.rs b/parquet/src/arrow/record_reader.rs
deleted file mode 100644
index 4dd7da9..0000000
--- a/parquet/src/arrow/record_reader.rs
+++ /dev/null
@@ -1,802 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-use std::cmp::{max, min};
-use std::mem::{replace, size_of};
-
-use crate::column::{page::PageReader, reader::ColumnReaderImpl};
-use crate::data_type::DataType;
-use crate::errors::{ParquetError, Result};
-use crate::schema::types::ColumnDescPtr;
-use arrow::array::BooleanBufferBuilder;
-use arrow::bitmap::Bitmap;
-use arrow::buffer::{Buffer, MutableBuffer};
-
-const MIN_BATCH_SIZE: usize = 1024;
-
-/// A `RecordReader` is a stateful column reader that delimits semantic records.
-pub struct RecordReader<T: DataType> {
- column_desc: ColumnDescPtr,
-
- records: MutableBuffer,
- def_levels: Option<MutableBuffer>,
- rep_levels: Option<MutableBuffer>,
- null_bitmap: Option<BooleanBufferBuilder>,
- column_reader: Option<ColumnReaderImpl<T>>,
-
- /// Number of records accumulated in records
- num_records: usize,
- /// Number of values `num_records` contains.
- num_values: usize,
-
- values_seen: usize,
- /// Starts from 1, number of values have been written to buffer
- values_written: usize,
- in_middle_of_record: bool,
-}
-
-impl<T: DataType> RecordReader<T> {
- pub fn new(column_schema: ColumnDescPtr) -> Self {
- let (def_levels, null_map) = if column_schema.max_def_level() > 0 {
- (
- Some(MutableBuffer::new(MIN_BATCH_SIZE)),
- Some(BooleanBufferBuilder::new(0)),
- )
- } else {
- (None, None)
- };
-
- let rep_levels = if column_schema.max_rep_level() > 0 {
- Some(MutableBuffer::new(MIN_BATCH_SIZE))
- } else {
- None
- };
-
- Self {
- records: MutableBuffer::new(MIN_BATCH_SIZE),
- def_levels,
- rep_levels,
- null_bitmap: null_map,
- column_reader: None,
- column_desc: column_schema,
- num_records: 0,
- num_values: 0,
- values_seen: 0,
- values_written: 0,
- in_middle_of_record: false,
- }
- }
-
- /// Set the current page reader.
- pub fn set_page_reader(&mut self, page_reader: Box<dyn PageReader>) -> Result<()> {
- self.column_reader =
- Some(ColumnReaderImpl::new(self.column_desc.clone(), page_reader));
- Ok(())
- }
-
- /// Try to read `num_records` of column data into internal buffer.
- ///
- /// # Returns
- ///
- /// Number of actual records read.
- pub fn read_records(&mut self, num_records: usize) -> Result<usize> {
- if self.column_reader.is_none() {
- return Ok(0);
- }
-
- let mut records_read = 0;
-
- // Used to mark whether we have reached the end of current
- // column chunk
- let mut end_of_column = false;
-
- loop {
- // Try to find some records from buffers that has been read into memory
- // but not counted as seen records.
- records_read += self.split_records(num_records - records_read)?;
-
- // Since page reader contains complete records, so if we reached end of a
- // page reader, we should reach the end of a record
- if end_of_column
- && self.values_seen >= self.values_written
- && self.in_middle_of_record
- {
- self.num_records += 1;
- self.num_values = self.values_seen;
- self.in_middle_of_record = false;
- records_read += 1;
- }
-
- if (records_read >= num_records) || end_of_column {
- break;
- }
-
- let batch_size = max(num_records - records_read, MIN_BATCH_SIZE);
-
- // Try to more value from parquet pages
- let values_read = self.read_one_batch(batch_size)?;
- if values_read < batch_size {
- end_of_column = true;
- }
- }
-
- Ok(records_read)
- }
-
- /// Returns number of records stored in buffer.
- pub fn num_records(&self) -> usize {
- self.num_records
- }
-
- /// Return number of values stored in buffer.
- /// If the parquet column is not repeated, it should be equals to `num_records`,
- /// otherwise it should be larger than or equal to `num_records`.
- pub fn num_values(&self) -> usize {
- self.num_values
- }
-
- /// Returns definition level data.
- /// The implementation has side effects. It will create a new buffer to hold those
- /// definition level values that have already been read into memory but not counted
- /// as record values, e.g. those from `self.num_values` to `self.values_written`.
- pub fn consume_def_levels(&mut self) -> Result<Option<Buffer>> {
- let new_buffer = if let Some(ref mut def_levels_buf) = &mut self.def_levels {
- let num_left_values = self.values_written - self.num_values;
- // create an empty buffer, as it will be resized below
- let mut new_buffer = MutableBuffer::new(0);
- let num_bytes = num_left_values * size_of::<i16>();
- let new_len = self.num_values * size_of::<i16>();
-
- new_buffer.resize(num_bytes, 0);
-
- let new_def_levels = new_buffer.as_slice_mut();
- let left_def_levels = &def_levels_buf.as_slice_mut()[new_len..];
-
- new_def_levels[0..num_bytes].copy_from_slice(&left_def_levels[0..num_bytes]);
-
- def_levels_buf.resize(new_len, 0);
- Some(new_buffer)
- } else {
- None
- };
-
- Ok(replace(&mut self.def_levels, new_buffer).map(|x| x.into()))
- }
-
- /// Return repetition level data.
- /// The side effect is similar to `consume_def_levels`.
- pub fn consume_rep_levels(&mut self) -> Result<Option<Buffer>> {
- // TODO: Optimize to reduce the copy
- let new_buffer = if let Some(ref mut rep_levels_buf) = &mut self.rep_levels {
- let num_left_values = self.values_written - self.num_values;
- // create an empty buffer, as it will be resized below
- let mut new_buffer = MutableBuffer::new(0);
- let num_bytes = num_left_values * size_of::<i16>();
- let new_len = self.num_values * size_of::<i16>();
-
- new_buffer.resize(num_bytes, 0);
-
- let new_rep_levels = new_buffer.as_slice_mut();
- let left_rep_levels = &rep_levels_buf.as_slice_mut()[new_len..];
-
- new_rep_levels[0..num_bytes].copy_from_slice(&left_rep_levels[0..num_bytes]);
-
- rep_levels_buf.resize(new_len, 0);
-
- Some(new_buffer)
- } else {
- None
- };
-
- Ok(replace(&mut self.rep_levels, new_buffer).map(|x| x.into()))
- }
-
- /// Returns currently stored buffer data.
- /// The side effect is similar to `consume_def_levels`.
- pub fn consume_record_data(&mut self) -> Result<Buffer> {
- // TODO: Optimize to reduce the copy
- let num_left_values = self.values_written - self.num_values;
- // create an empty buffer, as it will be resized below
- let mut new_buffer = MutableBuffer::new(0);
- let num_bytes = num_left_values * T::get_type_size();
- let new_len = self.num_values * T::get_type_size();
-
- new_buffer.resize(num_bytes, 0);
-
- let new_records = new_buffer.as_slice_mut();
- let left_records = &mut self.records.as_slice_mut()[new_len..];
-
- new_records[0..num_bytes].copy_from_slice(&left_records[0..num_bytes]);
-
- self.records.resize(new_len, 0);
-
- Ok(replace(&mut self.records, new_buffer).into())
- }
-
- /// Returns currently stored null bitmap data.
- /// The side effect is similar to `consume_def_levels`.
- pub fn consume_bitmap_buffer(&mut self) -> Result<Option<Buffer>> {
- // TODO: Optimize to reduce the copy
- if self.column_desc.max_def_level() > 0 {
- assert!(self.null_bitmap.is_some());
- let num_left_values = self.values_written - self.num_values;
- let new_bitmap_builder = Some(BooleanBufferBuilder::new(max(
- MIN_BATCH_SIZE,
- num_left_values,
- )));
-
- let old_bitmap = replace(&mut self.null_bitmap, new_bitmap_builder)
- .map(|mut builder| builder.finish())
- .unwrap();
-
- let old_bitmap = Bitmap::from(old_bitmap);
-
- for i in self.num_values..self.values_written {
- self.null_bitmap
- .as_mut()
- .unwrap()
- .append(old_bitmap.is_set(i));
- }
-
- Ok(Some(old_bitmap.into_buffer()))
- } else {
- Ok(None)
- }
- }
-
- /// Reset state of record reader.
- /// Should be called after consuming data, e.g. `consume_rep_levels`,
- /// `consume_rep_levels`, `consume_record_data` and `consume_bitmap_buffer`.
- pub fn reset(&mut self) {
- self.values_written -= self.num_values;
- self.num_records = 0;
- self.num_values = 0;
- self.values_seen = 0;
- self.in_middle_of_record = false;
- }
-
- /// Returns bitmap data.
- pub fn consume_bitmap(&mut self) -> Result<Option<Bitmap>> {
- self.consume_bitmap_buffer()
- .map(|buffer| buffer.map(Bitmap::from))
- }
-
- /// Try to read one batch of data.
- fn read_one_batch(&mut self, batch_size: usize) -> Result<usize> {
- // Reserve spaces
- self.records
- .resize(self.records.len() + batch_size * T::get_type_size(), 0);
- if let Some(ref mut buf) = self.rep_levels {
- buf.resize(buf.len() + batch_size * size_of::<i16>(), 0);
- }
- if let Some(ref mut buf) = self.def_levels {
- buf.resize(buf.len() + batch_size * size_of::<i16>(), 0);
- }
-
- let values_written = self.values_written;
-
- // Convert mutable buffer spaces to mutable slices
- let (prefix, values, suffix) =
- unsafe { self.records.as_slice_mut().align_to_mut::<T::T>() };
- assert!(prefix.is_empty() && suffix.is_empty());
- let values = &mut values[values_written..];
-
- let def_levels = self.def_levels.as_mut().map(|buf| {
- let (prefix, def_levels, suffix) =
- unsafe { buf.as_slice_mut().align_to_mut::<i16>() };
- assert!(prefix.is_empty() && suffix.is_empty());
- &mut def_levels[values_written..]
- });
-
- let rep_levels = self.rep_levels.as_mut().map(|buf| {
- let (prefix, rep_levels, suffix) =
- unsafe { buf.as_slice_mut().align_to_mut::<i16>() };
- assert!(prefix.is_empty() && suffix.is_empty());
- &mut rep_levels[values_written..]
- });
-
- let (values_read, levels_read) = self
- .column_reader
- .as_mut()
- .unwrap()
- .read_batch(batch_size, def_levels, rep_levels, values)?;
-
- // get new references for the def levels.
- let def_levels = self.def_levels.as_ref().map(|buf| {
- let (prefix, def_levels, suffix) =
- unsafe { buf.as_slice().align_to::<i16>() };
- assert!(prefix.is_empty() && suffix.is_empty());
- &def_levels[values_written..]
- });
-
- let max_def_level = self.column_desc.max_def_level();
-
- if values_read < levels_read {
- let def_levels = def_levels.ok_or_else(|| {
- general_err!(
- "Definition levels should exist when data is less than levels!"
- )
- })?;
-
- // Fill spaces in column data with default values
- let mut values_pos = values_read;
- let mut level_pos = levels_read;
-
- while level_pos > values_pos {
- if def_levels[level_pos - 1] == max_def_level {
- // This values is not empty
- // We use swap rather than assign here because T::T doesn't
- // implement Copy
- values.swap(level_pos - 1, values_pos - 1);
- values_pos -= 1;
- } else {
- values[level_pos - 1] = T::T::default();
- }
-
- level_pos -= 1;
- }
- }
-
- // Fill in bitmap data
- if let Some(null_buffer) = self.null_bitmap.as_mut() {
- let def_levels = def_levels.ok_or_else(|| {
- general_err!(
- "Definition levels should exist when data is less than levels!"
- )
- })?;
- (0..levels_read)
- .for_each(|idx| null_buffer.append(def_levels[idx] == max_def_level));
- }
-
- let values_read = max(values_read, levels_read);
- self.set_values_written(self.values_written + values_read)?;
- Ok(values_read)
- }
-
- /// Split values into records according repetition definition and returns number of
- /// records read.
- #[allow(clippy::unnecessary_wraps)]
- fn split_records(&mut self, records_to_read: usize) -> Result<usize> {
- let rep_levels = self.rep_levels.as_ref().map(|buf| {
- let (prefix, rep_levels, suffix) =
- unsafe { buf.as_slice().align_to::<i16>() };
- assert!(prefix.is_empty() && suffix.is_empty());
- rep_levels
- });
-
- match rep_levels {
- Some(buf) => {
- let mut records_read = 0;
-
- while (self.values_seen < self.values_written)
- && (records_read < records_to_read)
- {
- if buf[self.values_seen] == 0 {
- if self.in_middle_of_record {
- records_read += 1;
- self.num_records += 1;
- self.num_values = self.values_seen;
- }
- self.in_middle_of_record = true;
- }
- self.values_seen += 1;
- }
-
- Ok(records_read)
- }
- None => {
- let records_read =
- min(records_to_read, self.values_written - self.values_seen);
- self.num_records += records_read;
- self.num_values += records_read;
- self.values_seen += records_read;
- self.in_middle_of_record = false;
-
- Ok(records_read)
- }
- }
- }
-
- #[allow(clippy::unnecessary_wraps)]
- fn set_values_written(&mut self, new_values_written: usize) -> Result<()> {
- self.values_written = new_values_written;
- self.records
- .resize(self.values_written * T::get_type_size(), 0);
-
- let new_levels_len = self.values_written * size_of::<i16>();
-
- if let Some(ref mut buf) = self.rep_levels {
- buf.resize(new_levels_len, 0)
- };
-
- if let Some(ref mut buf) = self.def_levels {
- buf.resize(new_levels_len, 0)
- };
-
- Ok(())
- }
-}
-
-#[cfg(test)]
-mod tests {
- use super::RecordReader;
- use crate::basic::Encoding;
- use crate::column::page::Page;
- use crate::column::page::PageReader;
- use crate::data_type::Int32Type;
- use crate::errors::Result;
- use crate::schema::parser::parse_message_type;
- use crate::schema::types::SchemaDescriptor;
- use crate::util::test_common::page_util::{DataPageBuilder, DataPageBuilderImpl};
- use arrow::array::{BooleanBufferBuilder, Int16BufferBuilder, Int32BufferBuilder};
- use arrow::bitmap::Bitmap;
- use std::sync::Arc;
-
- struct TestPageReader {
- pages: Box<dyn Iterator<Item = Page>>,
- }
-
- impl TestPageReader {
- pub fn new(pages: Vec<Page>) -> Self {
- Self {
- pages: Box::new(pages.into_iter()),
- }
- }
- }
-
- impl PageReader for TestPageReader {
- fn get_next_page(&mut self) -> Result<Option<Page>> {
- Ok(self.pages.next())
- }
- }
-
- impl Iterator for TestPageReader {
- type Item = Result<Page>;
-
- fn next(&mut self) -> Option<Self::Item> {
- self.get_next_page().transpose()
- }
- }
-
- #[test]
- fn test_read_required_records() {
- // Construct column schema
- let message_type = "
- message test_schema {
- REQUIRED INT32 leaf;
- }
- ";
- let desc = parse_message_type(message_type)
- .map(|t| SchemaDescriptor::new(Arc::new(t)))
- .map(|s| s.column(0))
- .unwrap();
-
- // Construct record reader
- let mut record_reader = RecordReader::<Int32Type>::new(desc.clone());
-
- // First page
-
- // Records data:
- // test_schema
- // leaf: 4
- // test_schema
- // leaf: 7
- // test_schema
- // leaf: 6
- // test_schema
- // left: 3
- // test_schema
- // left: 2
- {
- let values = [4, 7, 6, 3, 2];
- let mut pb = DataPageBuilderImpl::new(desc.clone(), 5, true);
- pb.add_values::<Int32Type>(Encoding::PLAIN, &values);
- let page = pb.consume();
-
- let page_reader = Box::new(TestPageReader::new(vec![page]));
- record_reader.set_page_reader(page_reader).unwrap();
- assert_eq!(2, record_reader.read_records(2).unwrap());
- assert_eq!(2, record_reader.num_records());
- assert_eq!(2, record_reader.num_values());
- assert_eq!(3, record_reader.read_records(3).unwrap());
- assert_eq!(5, record_reader.num_records());
- assert_eq!(5, record_reader.num_values());
- }
-
- // Second page
-
- // Records data:
- // test_schema
- // leaf: 8
- // test_schema
- // leaf: 9
- {
- let values = [8, 9];
- let mut pb = DataPageBuilderImpl::new(desc, 2, true);
- pb.add_values::<Int32Type>(Encoding::PLAIN, &values);
- let page = pb.consume();
-
- let page_reader = Box::new(TestPageReader::new(vec![page]));
- record_reader.set_page_reader(page_reader).unwrap();
- assert_eq!(2, record_reader.read_records(10).unwrap());
- assert_eq!(7, record_reader.num_records());
- assert_eq!(7, record_reader.num_values());
- }
-
- let mut bb = Int32BufferBuilder::new(7);
- bb.append_slice(&[4, 7, 6, 3, 2, 8, 9]);
- let expected_buffer = bb.finish();
- assert_eq!(
- expected_buffer,
- record_reader.consume_record_data().unwrap()
- );
- assert_eq!(None, record_reader.consume_def_levels().unwrap());
- assert_eq!(None, record_reader.consume_bitmap().unwrap());
- }
-
- #[test]
- fn test_read_optional_records() {
- // Construct column schema
- let message_type = "
- message test_schema {
- OPTIONAL Group test_struct {
- OPTIONAL INT32 leaf;
- }
- }
- ";
-
- let desc = parse_message_type(message_type)
- .map(|t| SchemaDescriptor::new(Arc::new(t)))
- .map(|s| s.column(0))
- .unwrap();
-
- // Construct record reader
- let mut record_reader = RecordReader::<Int32Type>::new(desc.clone());
-
- // First page
-
- // Records data:
- // test_schema
- // test_struct
- // test_schema
- // test_struct
- // left: 7
- // test_schema
- // test_schema
- // test_struct
- // leaf: 6
- // test_schema
- // test_struct
- // leaf: 6
- {
- let values = [7, 6, 3];
- //empty, non-empty, empty, non-empty, non-empty
- let def_levels = [1i16, 2i16, 0i16, 2i16, 2i16];
- let mut pb = DataPageBuilderImpl::new(desc.clone(), 5, true);
- pb.add_def_levels(2, &def_levels);
- pb.add_values::<Int32Type>(Encoding::PLAIN, &values);
- let page = pb.consume();
-
- let page_reader = Box::new(TestPageReader::new(vec![page]));
- record_reader.set_page_reader(page_reader).unwrap();
- assert_eq!(2, record_reader.read_records(2).unwrap());
- assert_eq!(2, record_reader.num_records());
- assert_eq!(2, record_reader.num_values());
- assert_eq!(3, record_reader.read_records(3).unwrap());
- assert_eq!(5, record_reader.num_records());
- assert_eq!(5, record_reader.num_values());
- }
-
- // Second page
-
- // Records data:
- // test_schema
- // test_schema
- // test_struct
- // left: 8
- {
- let values = [8];
- //empty, non-empty
- let def_levels = [0i16, 2i16];
- let mut pb = DataPageBuilderImpl::new(desc, 2, true);
- pb.add_def_levels(2, &def_levels);
- pb.add_values::<Int32Type>(Encoding::PLAIN, &values);
- let page = pb.consume();
-
- let page_reader = Box::new(TestPageReader::new(vec![page]));
- record_reader.set_page_reader(page_reader).unwrap();
- assert_eq!(2, record_reader.read_records(10).unwrap());
- assert_eq!(7, record_reader.num_records());
- assert_eq!(7, record_reader.num_values());
- }
-
- // Verify result record data
- let mut bb = Int32BufferBuilder::new(7);
- bb.append_slice(&[0, 7, 0, 6, 3, 0, 8]);
- let expected_buffer = bb.finish();
- assert_eq!(
- expected_buffer,
- record_reader.consume_record_data().unwrap()
- );
-
- // Verify result def levels
- let mut bb = Int16BufferBuilder::new(7);
- bb.append_slice(&[1i16, 2i16, 0i16, 2i16, 2i16, 0i16, 2i16]);
- let expected_def_levels = bb.finish();
- assert_eq!(
- Some(expected_def_levels),
- record_reader.consume_def_levels().unwrap()
- );
-
- // Verify bitmap
- let mut bb = BooleanBufferBuilder::new(7);
- bb.append_slice(&[false, true, false, true, true, false, true]);
- let expected_bitmap = Bitmap::from(bb.finish());
- assert_eq!(
- Some(expected_bitmap),
- record_reader.consume_bitmap().unwrap()
- );
- }
-
- #[test]
- fn test_read_repeated_records() {
- // Construct column schema
- let message_type = "
- message test_schema {
- REPEATED Group test_struct {
- REPEATED INT32 leaf;
- }
- }
- ";
-
- let desc = parse_message_type(message_type)
- .map(|t| SchemaDescriptor::new(Arc::new(t)))
- .map(|s| s.column(0))
- .unwrap();
-
- // Construct record reader
- let mut record_reader = RecordReader::<Int32Type>::new(desc.clone());
-
- // First page
-
- // Records data:
- // test_schema
- // test_struct
- // leaf: 4
- // test_schema
- // test_schema
- // test_struct
- // test_struct
- // leaf: 7
- // leaf: 6
- // leaf: 3
- // test_struct
- // leaf: 2
- {
- let values = [4, 7, 6, 3, 2];
- let def_levels = [2i16, 0i16, 1i16, 2i16, 2i16, 2i16, 2i16];
- let rep_levels = [0i16, 0i16, 0i16, 1i16, 2i16, 2i16, 1i16];
- let mut pb = DataPageBuilderImpl::new(desc.clone(), 7, true);
- pb.add_rep_levels(2, &rep_levels);
- pb.add_def_levels(2, &def_levels);
- pb.add_values::<Int32Type>(Encoding::PLAIN, &values);
- let page = pb.consume();
-
- let page_reader = Box::new(TestPageReader::new(vec![page]));
- record_reader.set_page_reader(page_reader).unwrap();
-
- assert_eq!(1, record_reader.read_records(1).unwrap());
- assert_eq!(1, record_reader.num_records());
- assert_eq!(1, record_reader.num_values());
- assert_eq!(2, record_reader.read_records(3).unwrap());
- assert_eq!(3, record_reader.num_records());
- assert_eq!(7, record_reader.num_values());
- }
-
- // Second page
-
- // Records data:
- // test_schema
- // test_struct
- // leaf: 8
- // leaf: 9
- {
- let values = [8, 9];
- let def_levels = [2i16, 2i16];
- let rep_levels = [0i16, 2i16];
- let mut pb = DataPageBuilderImpl::new(desc, 2, true);
- pb.add_rep_levels(2, &rep_levels);
- pb.add_def_levels(2, &def_levels);
- pb.add_values::<Int32Type>(Encoding::PLAIN, &values);
- let page = pb.consume();
-
- let page_reader = Box::new(TestPageReader::new(vec![page]));
- record_reader.set_page_reader(page_reader).unwrap();
-
- assert_eq!(1, record_reader.read_records(10).unwrap());
- assert_eq!(4, record_reader.num_records());
- assert_eq!(9, record_reader.num_values());
- }
-
- // Verify result record data
- let mut bb = Int32BufferBuilder::new(9);
- bb.append_slice(&[4, 0, 0, 7, 6, 3, 2, 8, 9]);
- let expected_buffer = bb.finish();
- assert_eq!(
- expected_buffer,
- record_reader.consume_record_data().unwrap()
- );
-
- // Verify result def levels
- let mut bb = Int16BufferBuilder::new(9);
- bb.append_slice(&[2i16, 0i16, 1i16, 2i16, 2i16, 2i16, 2i16, 2i16, 2i16]);
- let expected_def_levels = bb.finish();
- assert_eq!(
- Some(expected_def_levels),
- record_reader.consume_def_levels().unwrap()
- );
-
- // Verify bitmap
- let mut bb = BooleanBufferBuilder::new(9);
- bb.append_slice(&[true, false, false, true, true, true, true, true, true]);
- let expected_bitmap = Bitmap::from(bb.finish());
- assert_eq!(
- Some(expected_bitmap),
- record_reader.consume_bitmap().unwrap()
- );
- }
-
- #[test]
- fn test_read_more_than_one_batch() {
- // Construct column schema
- let message_type = "
- message test_schema {
- REPEATED INT32 leaf;
- }
- ";
-
- let desc = parse_message_type(message_type)
- .map(|t| SchemaDescriptor::new(Arc::new(t)))
- .map(|s| s.column(0))
- .unwrap();
-
- // Construct record reader
- let mut record_reader = RecordReader::<Int32Type>::new(desc.clone());
-
- {
- let values = [100; 5000];
- let def_levels = [1i16; 5000];
- let mut rep_levels = [1i16; 5000];
- for idx in 0..1000 {
- rep_levels[idx * 5] = 0i16;
- }
-
- let mut pb = DataPageBuilderImpl::new(desc, 5000, true);
- pb.add_rep_levels(1, &rep_levels);
- pb.add_def_levels(1, &def_levels);
- pb.add_values::<Int32Type>(Encoding::PLAIN, &values);
- let page = pb.consume();
-
- let page_reader = Box::new(TestPageReader::new(vec![page]));
- record_reader.set_page_reader(page_reader).unwrap();
-
- assert_eq!(1000, record_reader.read_records(1000).unwrap());
- assert_eq!(1000, record_reader.num_records());
- assert_eq!(5000, record_reader.num_values());
- }
- }
-}
diff --git a/parquet/src/arrow/schema.rs b/parquet/src/arrow/schema.rs
deleted file mode 100644
index 6c04b70..0000000
--- a/parquet/src/arrow/schema.rs
+++ /dev/null
@@ -1,1943 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-//! Provides API for converting parquet schema to arrow schema and vice versa.
-//!
-//! The main interfaces for converting parquet schema to arrow schema are
-//! `parquet_to_arrow_schema`, `parquet_to_arrow_schema_by_columns` and
-//! `parquet_to_arrow_field`.
-//!
-//! The interfaces for converting arrow schema to parquet schema is coming.
-
-use std::collections::{HashMap, HashSet};
-use std::sync::Arc;
-
-use arrow::datatypes::{DataType, Field, IntervalUnit, Schema, TimeUnit};
-use arrow::ipc::writer;
-
-use crate::errors::{ParquetError::ArrowError, Result};
-use crate::file::{metadata::KeyValue, properties::WriterProperties};
-use crate::schema::types::{ColumnDescriptor, SchemaDescriptor, Type, TypePtr};
-use crate::{
- basic::{
- ConvertedType, DecimalType, IntType, LogicalType, Repetition, TimeType,
- TimeUnit as ParquetTimeUnit, TimestampType, Type as PhysicalType,
- },
- errors::ParquetError,
-};
-
-/// Convert Parquet schema to Arrow schema including optional metadata.
-/// Attempts to decode any existing Arrow schema metadata, falling back
-/// to converting the Parquet schema column-wise
-pub fn parquet_to_arrow_schema(
- parquet_schema: &SchemaDescriptor,
- key_value_metadata: &Option<Vec<KeyValue>>,
-) -> Result<Schema> {
- let mut metadata = parse_key_value_metadata(key_value_metadata).unwrap_or_default();
- metadata
- .remove(super::ARROW_SCHEMA_META_KEY)
- .map(|encoded| get_arrow_schema_from_metadata(&encoded))
- .unwrap_or(parquet_to_arrow_schema_by_columns(
- parquet_schema,
- 0..parquet_schema.columns().len(),
- key_value_metadata,
- ))
-}
-
-/// Convert parquet schema to arrow schema including optional metadata,
-/// only preserving some root columns.
-/// This is useful if we have columns `a.b`, `a.c.e` and `a.d`,
-/// and want `a` with all its child fields
-pub fn parquet_to_arrow_schema_by_root_columns<T>(
- parquet_schema: &SchemaDescriptor,
- column_indices: T,
- key_value_metadata: &Option<Vec<KeyValue>>,
-) -> Result<Schema>
-where
- T: IntoIterator<Item = usize>,
-{
- // Reconstruct the index ranges of the parent columns
- // An Arrow struct gets represented by 1+ columns based on how many child fields the
- // struct has. This means that getting fields 1 and 2 might return the struct twice,
- // if field 1 is the struct having say 3 fields, and field 2 is a primitive.
- //
- // The below gets the parent columns, and counts the number of child fields in each parent,
- // such that we would end up with:
- // - field 1 - columns: [0, 1, 2]
- // - field 2 - columns: [3]
- let mut parent_columns = vec![];
- let mut curr_name = "";
- let mut prev_name = "";
- let mut indices = vec![];
- (0..(parquet_schema.num_columns())).for_each(|i| {
- let p_type = parquet_schema.get_column_root(i);
- curr_name = p_type.get_basic_info().name();
- if prev_name.is_empty() {
- // first index
- indices.push(i);
- prev_name = curr_name;
- } else if curr_name != prev_name {
- prev_name = curr_name;
- parent_columns.push((curr_name.to_string(), indices.clone()));
- indices = vec![i];
- } else {
- indices.push(i);
- }
- });
- // push the last column if indices has values
- if !indices.is_empty() {
- parent_columns.push((curr_name.to_string(), indices));
- }
-
- // gather the required leaf columns
- let leaf_columns = column_indices
- .into_iter()
- .flat_map(|i| parent_columns[i].1.clone());
-
- parquet_to_arrow_schema_by_columns(parquet_schema, leaf_columns, key_value_metadata)
-}
-
-/// Convert parquet schema to arrow schema including optional metadata,
-/// only preserving some leaf columns.
-pub fn parquet_to_arrow_schema_by_columns<T>(
- parquet_schema: &SchemaDescriptor,
- column_indices: T,
- key_value_metadata: &Option<Vec<KeyValue>>,
-) -> Result<Schema>
-where
- T: IntoIterator<Item = usize>,
-{
- let mut metadata = parse_key_value_metadata(key_value_metadata).unwrap_or_default();
- let arrow_schema_metadata = metadata
- .remove(super::ARROW_SCHEMA_META_KEY)
- .map(|encoded| get_arrow_schema_from_metadata(&encoded))
- .map_or(Ok(None), |v| v.map(Some))?;
-
- // add the Arrow metadata to the Parquet metadata
- if let Some(arrow_schema) = &arrow_schema_metadata {
- arrow_schema.metadata().iter().for_each(|(k, v)| {
- metadata.insert(k.clone(), v.clone());
- });
- }
-
- let mut base_nodes = Vec::new();
- let mut base_nodes_set = HashSet::new();
- let mut leaves = HashSet::new();
-
- enum FieldType<'a> {
- Parquet(&'a Type),
- Arrow(Field),
- }
-
- for c in column_indices {
- let column = parquet_schema.column(c);
- let name = column.name();
-
- if let Some(field) = arrow_schema_metadata
- .as_ref()
- .and_then(|schema| schema.field_with_name(name).ok().cloned())
- {
- base_nodes.push(FieldType::Arrow(field));
- } else {
- let column = column.self_type() as *const Type;
- let root = parquet_schema.get_column_root(c);
- let root_raw_ptr = root as *const Type;
-
- leaves.insert(column);
- if !base_nodes_set.contains(&root_raw_ptr) {
- base_nodes.push(FieldType::Parquet(root));
- base_nodes_set.insert(root_raw_ptr);
- }
- }
- }
-
- base_nodes
- .into_iter()
- .map(|t| match t {
- FieldType::Parquet(t) => ParquetTypeConverter::new(t, &leaves).to_field(),
- FieldType::Arrow(f) => Ok(Some(f)),
- })
- .collect::<Result<Vec<Option<Field>>>>()
- .map(|result| result.into_iter().flatten().collect::<Vec<Field>>())
- .map(|fields| Schema::new_with_metadata(fields, metadata))
-}
-
-/// Try to convert Arrow schema metadata into a schema
-fn get_arrow_schema_from_metadata(encoded_meta: &str) -> Result<Schema> {
- let decoded = base64::decode(encoded_meta);
- match decoded {
- Ok(bytes) => {
- let slice = if bytes[0..4] == [255u8; 4] {
- &bytes[8..]
- } else {
- bytes.as_slice()
- };
- match arrow::ipc::root_as_message(slice) {
- Ok(message) => message
- .header_as_schema()
- .map(arrow::ipc::convert::fb_to_schema)
- .ok_or(ArrowError("the message is not Arrow Schema".to_string())),
- Err(err) => {
- // The flatbuffers implementation returns an error on verification error.
- Err(ArrowError(format!(
- "Unable to get root as message stored in {}: {:?}",
- super::ARROW_SCHEMA_META_KEY,
- err
- )))
- }
- }
- }
- Err(err) => {
- // The C++ implementation returns an error if the schema can't be parsed.
- Err(ArrowError(format!(
- "Unable to decode the encoded schema stored in {}, {:?}",
- super::ARROW_SCHEMA_META_KEY,
- err
- )))
- }
- }
-}
-
-/// Encodes the Arrow schema into the IPC format, and base64 encodes it
-fn encode_arrow_schema(schema: &Schema) -> String {
- let options = writer::IpcWriteOptions::default();
- let data_gen = arrow::ipc::writer::IpcDataGenerator::default();
- let mut serialized_schema = data_gen.schema_to_bytes(&schema, &options);
-
- // manually prepending the length to the schema as arrow uses the legacy IPC format
- // TODO: change after addressing ARROW-9777
- let schema_len = serialized_schema.ipc_message.len();
- let mut len_prefix_schema = Vec::with_capacity(schema_len + 8);
- len_prefix_schema.append(&mut vec![255u8, 255, 255, 255]);
- len_prefix_schema.append((schema_len as u32).to_le_bytes().to_vec().as_mut());
- len_prefix_schema.append(&mut serialized_schema.ipc_message);
-
- base64::encode(&len_prefix_schema)
-}
-
-/// Mutates writer metadata by storing the encoded Arrow schema.
-/// If there is an existing Arrow schema metadata, it is replaced.
-pub(crate) fn add_encoded_arrow_schema_to_metadata(
- schema: &Schema,
- props: &mut WriterProperties,
-) {
- let encoded = encode_arrow_schema(schema);
-
- let schema_kv = KeyValue {
- key: super::ARROW_SCHEMA_META_KEY.to_string(),
- value: Some(encoded),
- };
-
- let mut meta = props.key_value_metadata.clone().unwrap_or_default();
- // check if ARROW:schema exists, and overwrite it
- let schema_meta = meta
- .iter()
- .enumerate()
- .find(|(_, kv)| kv.key.as_str() == super::ARROW_SCHEMA_META_KEY);
- match schema_meta {
- Some((i, _)) => {
- meta.remove(i);
- meta.push(schema_kv);
- }
- None => {
- meta.push(schema_kv);
- }
- }
- props.key_value_metadata = Some(meta);
-}
-
-/// Convert arrow schema to parquet schema
-pub fn arrow_to_parquet_schema(schema: &Schema) -> Result<SchemaDescriptor> {
- let fields: Result<Vec<TypePtr>> = schema
- .fields()
- .iter()
- .map(|field| arrow_to_parquet_type(field).map(Arc::new))
- .collect();
- let group = Type::group_type_builder("arrow_schema")
- .with_fields(&mut fields?)
- .build()?;
- Ok(SchemaDescriptor::new(Arc::new(group)))
-}
-
-fn parse_key_value_metadata(
- key_value_metadata: &Option<Vec<KeyValue>>,
-) -> Option<HashMap<String, String>> {
- match key_value_metadata {
- Some(key_values) => {
- let map: HashMap<String, String> = key_values
- .iter()
- .filter_map(|kv| {
- kv.value
- .as_ref()
- .map(|value| (kv.key.clone(), value.clone()))
- })
- .collect();
-
- if map.is_empty() {
- None
- } else {
- Some(map)
- }
- }
- None => None,
- }
-}
-
-/// Convert parquet column schema to arrow field.
-pub fn parquet_to_arrow_field(parquet_column: &ColumnDescriptor) -> Result<Field> {
- let schema = parquet_column.self_type();
-
- let mut leaves = HashSet::new();
- leaves.insert(parquet_column.self_type() as *const Type);
-
- ParquetTypeConverter::new(schema, &leaves)
- .to_field()
- .map(|opt| opt.unwrap())
-}
-
-pub fn decimal_length_from_precision(precision: usize) -> usize {
- (10.0_f64.powi(precision as i32).log2() / 8.0).ceil() as usize
-}
-
-/// Convert an arrow field to a parquet `Type`
-fn arrow_to_parquet_type(field: &Field) -> Result<Type> {
- let name = field.name().as_str();
- let repetition = if field.is_nullable() {
- Repetition::OPTIONAL
- } else {
- Repetition::REQUIRED
- };
- // create type from field
- match field.data_type() {
- DataType::Null => Type::primitive_type_builder(name, PhysicalType::INT32)
- .with_logical_type(Some(LogicalType::UNKNOWN(Default::default())))
- .with_repetition(repetition)
- .build(),
- DataType::Boolean => Type::primitive_type_builder(name, PhysicalType::BOOLEAN)
- .with_repetition(repetition)
- .build(),
- DataType::Int8 => Type::primitive_type_builder(name, PhysicalType::INT32)
- .with_logical_type(Some(LogicalType::INTEGER(IntType {
- bit_width: 8,
- is_signed: true,
- })))
- .with_repetition(repetition)
- .build(),
- DataType::Int16 => Type::primitive_type_builder(name, PhysicalType::INT32)
- .with_logical_type(Some(LogicalType::INTEGER(IntType {
- bit_width: 16,
- is_signed: true,
- })))
- .with_repetition(repetition)
- .build(),
- DataType::Int32 => Type::primitive_type_builder(name, PhysicalType::INT32)
- .with_repetition(repetition)
- .build(),
- DataType::Int64 => Type::primitive_type_builder(name, PhysicalType::INT64)
- .with_repetition(repetition)
- .build(),
- DataType::UInt8 => Type::primitive_type_builder(name, PhysicalType::INT32)
- .with_logical_type(Some(LogicalType::INTEGER(IntType {
- bit_width: 8,
- is_signed: false,
- })))
- .with_repetition(repetition)
- .build(),
- DataType::UInt16 => Type::primitive_type_builder(name, PhysicalType::INT32)
- .with_logical_type(Some(LogicalType::INTEGER(IntType {
- bit_width: 16,
- is_signed: false,
- })))
- .with_repetition(repetition)
- .build(),
- DataType::UInt32 => Type::primitive_type_builder(name, PhysicalType::INT32)
- .with_logical_type(Some(LogicalType::INTEGER(IntType {
- bit_width: 32,
- is_signed: false,
- })))
- .with_repetition(repetition)
- .build(),
- DataType::UInt64 => Type::primitive_type_builder(name, PhysicalType::INT64)
- .with_logical_type(Some(LogicalType::INTEGER(IntType {
- bit_width: 64,
- is_signed: false,
- })))
- .with_repetition(repetition)
- .build(),
- DataType::Float16 => Err(ArrowError("Float16 arrays not supported".to_string())),
- DataType::Float32 => Type::primitive_type_builder(name, PhysicalType::FLOAT)
- .with_repetition(repetition)
- .build(),
- DataType::Float64 => Type::primitive_type_builder(name, PhysicalType::DOUBLE)
- .with_repetition(repetition)
- .build(),
- DataType::Timestamp(time_unit, zone) => Type::primitive_type_builder(
- name,
- PhysicalType::INT64,
- )
- .with_logical_type(Some(LogicalType::TIMESTAMP(TimestampType {
- is_adjusted_to_u_t_c: matches!(zone, Some(z) if !z.as_str().is_empty()),
- unit: match time_unit {
- TimeUnit::Second => ParquetTimeUnit::MILLIS(Default::default()),
- TimeUnit::Millisecond => ParquetTimeUnit::MILLIS(Default::default()),
- TimeUnit::Microsecond => ParquetTimeUnit::MICROS(Default::default()),
- TimeUnit::Nanosecond => ParquetTimeUnit::NANOS(Default::default()),
- },
- })))
- .with_repetition(repetition)
- .build(),
- DataType::Date32 => Type::primitive_type_builder(name, PhysicalType::INT32)
- .with_logical_type(Some(LogicalType::DATE(Default::default())))
- .with_repetition(repetition)
- .build(),
- // date64 is cast to date32
- DataType::Date64 => Type::primitive_type_builder(name, PhysicalType::INT32)
- .with_logical_type(Some(LogicalType::DATE(Default::default())))
- .with_repetition(repetition)
- .build(),
- DataType::Time32(_) => Type::primitive_type_builder(name, PhysicalType::INT32)
- .with_logical_type(Some(LogicalType::TIME(TimeType {
- is_adjusted_to_u_t_c: false,
- unit: ParquetTimeUnit::MILLIS(Default::default()),
- })))
- .with_repetition(repetition)
- .build(),
- DataType::Time64(unit) => Type::primitive_type_builder(name, PhysicalType::INT64)
- .with_logical_type(Some(LogicalType::TIME(TimeType {
- is_adjusted_to_u_t_c: false,
- unit: match unit {
- TimeUnit::Microsecond => ParquetTimeUnit::MICROS(Default::default()),
- TimeUnit::Nanosecond => ParquetTimeUnit::NANOS(Default::default()),
- u => unreachable!("Invalid unit for Time64: {:?}", u),
- },
- })))
- .with_repetition(repetition)
- .build(),
- DataType::Duration(_) => Err(ArrowError(
- "Converting Duration to parquet not supported".to_string(),
- )),
- DataType::Interval(_) => {
- Type::primitive_type_builder(name, PhysicalType::FIXED_LEN_BYTE_ARRAY)
- .with_converted_type(ConvertedType::INTERVAL)
- .with_repetition(repetition)
- .with_length(12)
- .build()
- }
- DataType::Binary | DataType::LargeBinary => {
- Type::primitive_type_builder(name, PhysicalType::BYTE_ARRAY)
- .with_repetition(repetition)
- .build()
- }
- DataType::FixedSizeBinary(length) => {
- Type::primitive_type_builder(name, PhysicalType::FIXED_LEN_BYTE_ARRAY)
- .with_repetition(repetition)
- .with_length(*length)
- .build()
- }
- DataType::Decimal(precision, scale) => {
- // Decimal precision determines the Parquet physical type to use.
- // TODO(ARROW-12018): Enable the below after ARROW-10818 Decimal support
- //
- // let (physical_type, length) = if *precision > 1 && *precision <= 9 {
- // (PhysicalType::INT32, -1)
- // } else if *precision <= 18 {
- // (PhysicalType::INT64, -1)
- // } else {
- // (
- // PhysicalType::FIXED_LEN_BYTE_ARRAY,
- // decimal_length_from_precision(*precision) as i32,
- // )
- // };
- Type::primitive_type_builder(name, PhysicalType::FIXED_LEN_BYTE_ARRAY)
- .with_repetition(repetition)
- .with_length(decimal_length_from_precision(*precision) as i32)
- .with_logical_type(Some(LogicalType::DECIMAL(DecimalType {
- scale: *scale as i32,
- precision: *precision as i32,
- })))
- .with_precision(*precision as i32)
- .with_scale(*scale as i32)
- .build()
- }
- DataType::Utf8 | DataType::LargeUtf8 => {
- Type::primitive_type_builder(name, PhysicalType::BYTE_ARRAY)
- .with_logical_type(Some(LogicalType::STRING(Default::default())))
- .with_repetition(repetition)
- .build()
- }
- DataType::List(f) | DataType::FixedSizeList(f, _) | DataType::LargeList(f) => {
- Type::group_type_builder(name)
- .with_fields(&mut vec![Arc::new(
- Type::group_type_builder("list")
- .with_fields(&mut vec![Arc::new(arrow_to_parquet_type(f)?)])
- .with_repetition(Repetition::REPEATED)
- .build()?,
- )])
- .with_logical_type(Some(LogicalType::LIST(Default::default())))
- .with_repetition(repetition)
- .build()
- }
- DataType::Struct(fields) => {
- if fields.is_empty() {
- return Err(ArrowError(
- "Parquet does not support writing empty structs".to_string(),
- ));
- }
- // recursively convert children to types/nodes
- let fields: Result<Vec<TypePtr>> = fields
- .iter()
- .map(|f| arrow_to_parquet_type(f).map(Arc::new))
- .collect();
- Type::group_type_builder(name)
- .with_fields(&mut fields?)
- .with_repetition(repetition)
- .build()
- }
- DataType::Union(_) => unimplemented!("See ARROW-8817."),
- DataType::Dictionary(_, ref value) => {
- // Dictionary encoding not handled at the schema level
- let dict_field = Field::new(name, *value.clone(), field.is_nullable());
- arrow_to_parquet_type(&dict_field)
- }
- }
-}
-/// This struct is used to group methods and data structures used to convert parquet
-/// schema together.
-struct ParquetTypeConverter<'a> {
- schema: &'a Type,
- /// This is the columns that need to be converted to arrow schema.
- columns_to_convert: &'a HashSet<*const Type>,
-}
-
-impl<'a> ParquetTypeConverter<'a> {
- fn new(schema: &'a Type, columns_to_convert: &'a HashSet<*const Type>) -> Self {
- Self {
- schema,
- columns_to_convert,
- }
- }
-
- fn clone_with_schema(&self, other: &'a Type) -> Self {
- Self {
- schema: other,
- columns_to_convert: self.columns_to_convert,
- }
- }
-}
-
-impl ParquetTypeConverter<'_> {
- // Public interfaces.
-
- /// Converts parquet schema to arrow data type.
- ///
- /// This function discards schema name.
- ///
- /// If this schema is a primitive type and not included in the leaves, the result is
- /// Ok(None).
- ///
- /// If this schema is a group type and none of its children is reserved in the
- /// conversion, the result is Ok(None).
- fn to_data_type(&self) -> Result<Option<DataType>> {
- match self.schema {
- Type::PrimitiveType { .. } => self.to_primitive_type(),
- Type::GroupType { .. } => self.to_group_type(),
- }
- }
-
- /// Converts parquet schema to arrow field.
- ///
- /// This method is roughly the same as
- /// [`to_data_type`](`ParquetTypeConverter::to_data_type`), except it reserves schema
- /// name.
- fn to_field(&self) -> Result<Option<Field>> {
- self.to_data_type().map(|opt| {
- opt.map(|dt| Field::new(self.schema.name(), dt, self.is_nullable()))
- })
- }
-
- // Utility functions.
-
- /// Checks whether this schema is nullable.
- fn is_nullable(&self) -> bool {
- let basic_info = self.schema.get_basic_info();
- if basic_info.has_repetition() {
- match basic_info.repetition() {
- Repetition::OPTIONAL => true,
- Repetition::REPEATED => true,
- Repetition::REQUIRED => false,
- }
- } else {
- false
- }
- }
-
- fn is_repeated(&self) -> bool {
- let basic_info = self.schema.get_basic_info();
-
- basic_info.has_repetition() && basic_info.repetition() == Repetition::REPEATED
- }
-
- fn is_self_included(&self) -> bool {
- self.columns_to_convert
- .contains(&(self.schema as *const Type))
- }
-
- // Functions for primitive types.
-
- /// Entry point for converting parquet primitive type to arrow type.
- ///
- /// This function takes care of repetition.
- fn to_primitive_type(&self) -> Result<Option<DataType>> {
- if self.is_self_included() {
- self.to_primitive_type_inner().map(|dt| {
- if self.is_repeated() {
- Some(DataType::List(Box::new(Field::new(
- self.schema.name(),
- dt,
- self.is_nullable(),
- ))))
- } else {
- Some(dt)
- }
- })
- } else {
- Ok(None)
- }
- }
-
- /// Converting parquet primitive type to arrow data type.
- fn to_primitive_type_inner(&self) -> Result<DataType> {
- match self.schema.get_physical_type() {
- PhysicalType::BOOLEAN => Ok(DataType::Boolean),
- PhysicalType::INT32 => self.from_int32(),
- PhysicalType::INT64 => self.from_int64(),
- PhysicalType::INT96 => Ok(DataType::Timestamp(TimeUnit::Nanosecond, None)),
- PhysicalType::FLOAT => Ok(DataType::Float32),
- PhysicalType::DOUBLE => Ok(DataType::Float64),
- PhysicalType::BYTE_ARRAY => self.from_byte_array(),
- PhysicalType::FIXED_LEN_BYTE_ARRAY => self.from_fixed_len_byte_array(),
- }
- }
-
- fn from_int32(&self) -> Result<DataType> {
- match (
- self.schema.get_basic_info().logical_type(),
- self.schema.get_basic_info().converted_type(),
- ) {
- (None, ConvertedType::NONE) => Ok(DataType::Int32),
- (Some(LogicalType::INTEGER(t)), _) => match (t.bit_width, t.is_signed) {
- (8, true) => Ok(DataType::Int8),
- (16, true) => Ok(DataType::Int16),
- (32, true) => Ok(DataType::Int32),
- (8, false) => Ok(DataType::UInt8),
- (16, false) => Ok(DataType::UInt16),
- (32, false) => Ok(DataType::UInt32),
- _ => Err(ArrowError(format!(
- "Cannot create INT32 physical type from {:?}",
- t
- ))),
- },
- (Some(LogicalType::DECIMAL(_)), _) => Ok(self.to_decimal()),
- (Some(LogicalType::DATE(_)), _) => Ok(DataType::Date32),
- (Some(LogicalType::TIME(t)), _) => match t.unit {
- ParquetTimeUnit::MILLIS(_) => Ok(DataType::Time32(TimeUnit::Millisecond)),
- _ => Err(ArrowError(format!(
- "Cannot create INT32 physical type from {:?}",
- t.unit
- ))),
- },
- (None, ConvertedType::UINT_8) => Ok(DataType::UInt8),
- (None, ConvertedType::UINT_16) => Ok(DataType::UInt16),
- (None, ConvertedType::UINT_32) => Ok(DataType::UInt32),
- (None, ConvertedType::INT_8) => Ok(DataType::Int8),
- (None, ConvertedType::INT_16) => Ok(DataType::Int16),
- (None, ConvertedType::INT_32) => Ok(DataType::Int32),
- (None, ConvertedType::DATE) => Ok(DataType::Date32),
- (None, ConvertedType::TIME_MILLIS) => {
- Ok(DataType::Time32(TimeUnit::Millisecond))
- }
- (None, ConvertedType::DECIMAL) => Ok(self.to_decimal()),
- (logical, converted) => Err(ArrowError(format!(
- "Unable to convert parquet INT32 logical type {:?} or converted type {}",
- logical, converted
- ))),
- }
- }
-
- fn from_int64(&self) -> Result<DataType> {
- match (
- self.schema.get_basic_info().logical_type(),
- self.schema.get_basic_info().converted_type(),
- ) {
- (None, ConvertedType::NONE) => Ok(DataType::Int64),
- (Some(LogicalType::INTEGER(t)), _) if t.bit_width == 64 => {
- match t.is_signed {
- true => Ok(DataType::Int64),
- false => Ok(DataType::UInt64),
- }
- }
- (Some(LogicalType::TIME(t)), _) => match t.unit {
- ParquetTimeUnit::MILLIS(_) => Err(ArrowError(
- "Cannot create INT64 from MILLIS time unit".to_string(),
- )),
- ParquetTimeUnit::MICROS(_) => Ok(DataType::Time64(TimeUnit::Microsecond)),
- ParquetTimeUnit::NANOS(_) => Ok(DataType::Time64(TimeUnit::Nanosecond)),
- },
- (Some(LogicalType::TIMESTAMP(t)), _) => Ok(DataType::Timestamp(
- match t.unit {
- ParquetTimeUnit::MILLIS(_) => TimeUnit::Millisecond,
- ParquetTimeUnit::MICROS(_) => TimeUnit::Microsecond,
- ParquetTimeUnit::NANOS(_) => TimeUnit::Nanosecond,
- },
- if t.is_adjusted_to_u_t_c {
- Some("UTC".to_string())
- } else {
- None
- },
- )),
- (None, ConvertedType::INT_64) => Ok(DataType::Int64),
- (None, ConvertedType::UINT_64) => Ok(DataType::UInt64),
- (None, ConvertedType::TIME_MICROS) => {
- Ok(DataType::Time64(TimeUnit::Microsecond))
- }
- (None, ConvertedType::TIMESTAMP_MILLIS) => {
- Ok(DataType::Timestamp(TimeUnit::Millisecond, None))
- }
- (None, ConvertedType::TIMESTAMP_MICROS) => {
- Ok(DataType::Timestamp(TimeUnit::Microsecond, None))
- }
- (Some(LogicalType::DECIMAL(_)), _) => Ok(self.to_decimal()),
- (None, ConvertedType::DECIMAL) => Ok(self.to_decimal()),
- (logical, converted) => Err(ArrowError(format!(
- "Unable to convert parquet INT64 logical type {:?} or converted type {}",
- logical, converted
- ))),
- }
- }
-
- fn from_fixed_len_byte_array(&self) -> Result<DataType> {
- match (
- self.schema.get_basic_info().logical_type(),
- self.schema.get_basic_info().converted_type(),
- ) {
- (Some(LogicalType::DECIMAL(_)), _) => Ok(self.to_decimal()),
- (None, ConvertedType::DECIMAL) => Ok(self.to_decimal()),
- (None, ConvertedType::INTERVAL) => {
- // There is currently no reliable way of determining which IntervalUnit
- // to return. Thus without the original Arrow schema, the results
- // would be incorrect if all 12 bytes of the interval are populated
- Ok(DataType::Interval(IntervalUnit::DayTime))
- }
- _ => {
- let byte_width = match self.schema {
- Type::PrimitiveType {
- ref type_length, ..
- } => *type_length,
- _ => {
- return Err(ArrowError(
- "Expected a physical type, not a group type".to_string(),
- ))
- }
- };
-
- Ok(DataType::FixedSizeBinary(byte_width))
- }
- }
- }
-
- fn to_decimal(&self) -> DataType {
- assert!(self.schema.is_primitive());
- DataType::Decimal(
- self.schema.get_precision() as usize,
- self.schema.get_scale() as usize,
- )
- }
-
- fn from_byte_array(&self) -> Result<DataType> {
- match (self.schema.get_basic_info().logical_type(), self.schema.get_basic_info().converted_type()) {
- (Some(LogicalType::STRING(_)), _) => Ok(DataType::Utf8),
- (Some(LogicalType::JSON(_)), _) => Ok(DataType::Binary),
- (Some(LogicalType::BSON(_)), _) => Ok(DataType::Binary),
- (Some(LogicalType::ENUM(_)), _) => Ok(DataType::Binary),
- (None, ConvertedType::NONE) => Ok(DataType::Binary),
- (None, ConvertedType::JSON) => Ok(DataType::Binary),
- (None, ConvertedType::BSON) => Ok(DataType::Binary),
- (None, ConvertedType::ENUM) => Ok(DataType::Binary),
- (None, ConvertedType::UTF8) => Ok(DataType::Utf8),
- (logical, converted) => Err(ArrowError(format!(
- "Unable to convert parquet BYTE_ARRAY logical type {:?} or converted type {}",
- logical, converted
- ))),
- }
- }
-
- // Functions for group types.
-
- /// Entry point for converting parquet group type.
- ///
- /// This function takes care of logical type and repetition.
- fn to_group_type(&self) -> Result<Option<DataType>> {
- if self.is_repeated() {
- self.to_struct().map(|opt| {
- opt.map(|dt| {
- DataType::List(Box::new(Field::new(
- self.schema.name(),
- dt,
- self.is_nullable(),
- )))
- })
- })
- } else {
- match (
- self.schema.get_basic_info().logical_type(),
- self.schema.get_basic_info().converted_type(),
- ) {
- (Some(LogicalType::LIST(_)), _) => self.to_list(),
- (None, ConvertedType::LIST) => self.to_list(),
- _ => self.to_struct(),
- }
- }
- }
-
- /// Converts a parquet group type to arrow struct.
- fn to_struct(&self) -> Result<Option<DataType>> {
- match self.schema {
- Type::PrimitiveType { .. } => Err(ParquetError::General(format!(
- "{:?} is a struct type, and can't be processed as primitive.",
- self.schema
- ))),
- Type::GroupType {
- basic_info: _,
- fields,
- } => fields
- .iter()
- .map(|field_ptr| self.clone_with_schema(field_ptr).to_field())
- .collect::<Result<Vec<Option<Field>>>>()
- .map(|result| result.into_iter().flatten().collect::<Vec<Field>>())
- .map(|fields| {
- if fields.is_empty() {
- None
- } else {
- Some(DataType::Struct(fields))
- }
- }),
- }
- }
-
- /// Converts a parquet list to arrow list.
- ///
- /// To fully understand this algorithm, please refer to
- /// [parquet doc](https://github.com/apache/parquet-format/blob/master/LogicalTypes.md).
- fn to_list(&self) -> Result<Option<DataType>> {
- match self.schema {
- Type::PrimitiveType { .. } => Err(ParquetError::General(format!(
- "{:?} is a list type and can't be processed as primitive.",
- self.schema
- ))),
- Type::GroupType {
- basic_info: _,
- fields,
- } if fields.len() == 1 => {
- let list_item = fields.first().unwrap();
- let item_converter = self.clone_with_schema(list_item);
-
- let item_type = match list_item.as_ref() {
- Type::PrimitiveType { .. } => {
- if item_converter.is_repeated() {
- item_converter.to_primitive_type_inner().map(Some)
- } else {
- Err(ArrowError(
- "Primitive element type of list must be repeated."
- .to_string(),
- ))
- }
- }
- Type::GroupType {
- basic_info: _,
- fields,
- } => {
- if fields.len() > 1 {
- item_converter.to_struct()
- } else if fields.len() == 1
- && list_item.name() != "array"
- && list_item.name() != format!("{}_tuple", self.schema.name())
- {
- let nested_item = fields.first().unwrap();
- let nested_item_converter =
- self.clone_with_schema(nested_item);
-
- nested_item_converter.to_data_type()
- } else {
- item_converter.to_struct()
- }
- }
- };
-
- // Check that the name of the list child is "list", in which case we
- // get the child nullability and name (normally "element") from the nested
- // group type.
- // Without this step, the child incorrectly inherits the parent's optionality
- let (list_item_name, item_is_optional) = match &item_converter.schema {
- Type::GroupType { basic_info, fields }
- if basic_info.name() == "list" && fields.len() == 1 =>
- {
- let field = fields.first().unwrap();
- (field.name(), field.is_optional())
- }
- _ => (list_item.name(), list_item.is_optional()),
- };
-
- item_type.map(|opt| {
- opt.map(|dt| {
- DataType::List(Box::new(Field::new(
- list_item_name,
- dt,
- item_is_optional,
- )))
- })
- })
- }
- _ => Err(ArrowError(
- "Group element type of list can only contain one field.".to_string(),
- )),
- }
- }
-}
-
-#[cfg(test)]
-mod tests {
- use super::*;
-
- use std::{collections::HashMap, convert::TryFrom, sync::Arc};
-
- use arrow::datatypes::{DataType, Field, IntervalUnit, TimeUnit};
-
- use crate::file::{metadata::KeyValue, reader::SerializedFileReader};
- use crate::{
- arrow::{ArrowReader, ArrowWriter, ParquetFileArrowReader},
- schema::{parser::parse_message_type, types::SchemaDescriptor},
- util::test_common::get_temp_file,
- };
-
- #[test]
- fn test_flat_primitives() {
- let message_type = "
- message test_schema {
- REQUIRED BOOLEAN boolean;
- REQUIRED INT32 int8 (INT_8);
- REQUIRED INT32 int16 (INT_16);
- REQUIRED INT32 uint8 (INTEGER(8,false));
- REQUIRED INT32 uint16 (INTEGER(16,false));
- REQUIRED INT32 int32;
- REQUIRED INT64 int64 ;
- OPTIONAL DOUBLE double;
- OPTIONAL FLOAT float;
- OPTIONAL BINARY string (UTF8);
- OPTIONAL BINARY string_2 (STRING);
- }
- ";
- let parquet_group_type = parse_message_type(message_type).unwrap();
-
- let parquet_schema = SchemaDescriptor::new(Arc::new(parquet_group_type));
- let converted_arrow_schema =
- parquet_to_arrow_schema(&parquet_schema, &None).unwrap();
-
- let arrow_fields = vec![
- Field::new("boolean", DataType::Boolean, false),
- Field::new("int8", DataType::Int8, false),
- Field::new("int16", DataType::Int16, false),
- Field::new("uint8", DataType::UInt8, false),
- Field::new("uint16", DataType::UInt16, false),
- Field::new("int32", DataType::Int32, false),
- Field::new("int64", DataType::Int64, false),
- Field::new("double", DataType::Float64, true),
- Field::new("float", DataType::Float32, true),
- Field::new("string", DataType::Utf8, true),
- Field::new("string_2", DataType::Utf8, true),
- ];
-
- assert_eq!(&arrow_fields, converted_arrow_schema.fields());
- }
-
- #[test]
- fn test_byte_array_fields() {
- let message_type = "
- message test_schema {
- REQUIRED BYTE_ARRAY binary;
- REQUIRED FIXED_LEN_BYTE_ARRAY (20) fixed_binary;
- }
- ";
-
- let parquet_group_type = parse_message_type(message_type).unwrap();
-
- let parquet_schema = SchemaDescriptor::new(Arc::new(parquet_group_type));
- let converted_arrow_schema =
- parquet_to_arrow_schema(&parquet_schema, &None).unwrap();
-
- let arrow_fields = vec![
- Field::new("binary", DataType::Binary, false),
- Field::new("fixed_binary", DataType::FixedSizeBinary(20), false),
- ];
- assert_eq!(&arrow_fields, converted_arrow_schema.fields());
- }
-
- #[test]
- fn test_duplicate_fields() {
- let message_type = "
- message test_schema {
- REQUIRED BOOLEAN boolean;
- REQUIRED INT32 int8 (INT_8);
- }
- ";
-
- let parquet_group_type = parse_message_type(message_type).unwrap();
-
- let parquet_schema = SchemaDescriptor::new(Arc::new(parquet_group_type));
- let converted_arrow_schema =
- parquet_to_arrow_schema(&parquet_schema, &None).unwrap();
-
- let arrow_fields = vec![
- Field::new("boolean", DataType::Boolean, false),
- Field::new("int8", DataType::Int8, false),
- ];
- assert_eq!(&arrow_fields, converted_arrow_schema.fields());
-
- let converted_arrow_schema = parquet_to_arrow_schema_by_columns(
- &parquet_schema,
- vec![0usize, 1usize],
- &None,
- )
- .unwrap();
- assert_eq!(&arrow_fields, converted_arrow_schema.fields());
- }
-
- #[test]
- fn test_parquet_lists() {
- let mut arrow_fields = Vec::new();
-
- // LIST encoding example taken from parquet-format/LogicalTypes.md
- let message_type = "
- message test_schema {
- REQUIRED GROUP my_list (LIST) {
- REPEATED GROUP list {
- OPTIONAL BINARY element (UTF8);
- }
- }
- OPTIONAL GROUP my_list (LIST) {
- REPEATED GROUP list {
- REQUIRED BINARY element (UTF8);
- }
- }
- OPTIONAL GROUP array_of_arrays (LIST) {
- REPEATED GROUP list {
- REQUIRED GROUP element (LIST) {
- REPEATED GROUP list {
- REQUIRED INT32 element;
- }
- }
- }
- }
- OPTIONAL GROUP my_list (LIST) {
- REPEATED GROUP element {
- REQUIRED BINARY str (UTF8);
- }
- }
- OPTIONAL GROUP my_list (LIST) {
- REPEATED INT32 element;
- }
- OPTIONAL GROUP my_list (LIST) {
- REPEATED GROUP element {
- REQUIRED BINARY str (UTF8);
- REQUIRED INT32 num;
- }
- }
- OPTIONAL GROUP my_list (LIST) {
- REPEATED GROUP array {
- REQUIRED BINARY str (UTF8);
- }
-
- }
- OPTIONAL GROUP my_list (LIST) {
- REPEATED GROUP my_list_tuple {
- REQUIRED BINARY str (UTF8);
- }
- }
- REPEATED INT32 name;
- }
- ";
-
- // // List<String> (list non-null, elements nullable)
- // required group my_list (LIST) {
- // repeated group list {
- // optional binary element (UTF8);
- // }
- // }
- {
- arrow_fields.push(Field::new(
- "my_list",
- DataType::List(Box::new(Field::new("element", DataType::Utf8, true))),
- false,
- ));
- }
-
- // // List<String> (list nullable, elements non-null)
- // optional group my_list (LIST) {
- // repeated group list {
- // required binary element (UTF8);
- // }
- // }
- {
- arrow_fields.push(Field::new(
- "my_list",
- DataType::List(Box::new(Field::new("element", DataType::Utf8, false))),
- true,
- ));
- }
-
- // Element types can be nested structures. For example, a list of lists:
- //
- // // List<List<Integer>>
- // optional group array_of_arrays (LIST) {
- // repeated group list {
- // required group element (LIST) {
- // repeated group list {
- // required int32 element;
- // }
- // }
- // }
- // }
- {
- let arrow_inner_list =
- DataType::List(Box::new(Field::new("element", DataType::Int32, false)));
- arrow_fields.push(Field::new(
- "array_of_arrays",
- DataType::List(Box::new(Field::new("element", arrow_inner_list, false))),
- true,
- ));
- }
-
- // // List<String> (list nullable, elements non-null)
- // optional group my_list (LIST) {
- // repeated group element {
- // required binary str (UTF8);
- // };
- // }
- {
- arrow_fields.push(Field::new(
- "my_list",
- DataType::List(Box::new(Field::new("element", DataType::Utf8, true))),
- true,
- ));
- }
-
- // // List<Integer> (nullable list, non-null elements)
- // optional group my_list (LIST) {
- // repeated int32 element;
- // }
- {
- arrow_fields.push(Field::new(
- "my_list",
- DataType::List(Box::new(Field::new("element", DataType::Int32, true))),
- true,
- ));
- }
-
- // // List<Tuple<String, Integer>> (nullable list, non-null elements)
- // optional group my_list (LIST) {
- // repeated group element {
- // required binary str (UTF8);
- // required int32 num;
- // };
- // }
- {
- let arrow_struct = DataType::Struct(vec![
- Field::new("str", DataType::Utf8, false),
- Field::new("num", DataType::Int32, false),
- ]);
- arrow_fields.push(Field::new(
- "my_list",
- DataType::List(Box::new(Field::new("element", arrow_struct, true))),
- true,
- ));
- }
-
- // // List<OneTuple<String>> (nullable list, non-null elements)
- // optional group my_list (LIST) {
- // repeated group array {
- // required binary str (UTF8);
- // };
- // }
- // Special case: group is named array
- {
- let arrow_struct =
- DataType::Struct(vec![Field::new("str", DataType::Utf8, false)]);
- arrow_fields.push(Field::new(
- "my_list",
- DataType::List(Box::new(Field::new("array", arrow_struct, true))),
- true,
- ));
- }
-
- // // List<OneTuple<String>> (nullable list, non-null elements)
- // optional group my_list (LIST) {
- // repeated group my_list_tuple {
- // required binary str (UTF8);
- // };
- // }
- // Special case: group named ends in _tuple
- {
- let arrow_struct =
- DataType::Struct(vec![Field::new("str", DataType::Utf8, false)]);
- arrow_fields.push(Field::new(
- "my_list",
- DataType::List(Box::new(Field::new("my_list_tuple", arrow_struct, true))),
- true,
- ));
- }
-
- // One-level encoding: Only allows required lists with required cells
- // repeated value_type name
- {
- arrow_fields.push(Field::new(
- "name",
- DataType::List(Box::new(Field::new("name", DataType::Int32, true))),
- true,
- ));
- }
-
- let parquet_group_type = parse_message_type(message_type).unwrap();
-
- let parquet_schema = SchemaDescriptor::new(Arc::new(parquet_group_type));
- let converted_arrow_schema =
- parquet_to_arrow_schema(&parquet_schema, &None).unwrap();
- let converted_fields = converted_arrow_schema.fields();
-
- assert_eq!(arrow_fields.len(), converted_fields.len());
- for i in 0..arrow_fields.len() {
- assert_eq!(arrow_fields[i], converted_fields[i]);
- }
- }
-
- #[test]
- fn test_parquet_list_nullable() {
- let mut arrow_fields = Vec::new();
-
- let message_type = "
- message test_schema {
- REQUIRED GROUP my_list1 (LIST) {
- REPEATED GROUP list {
- OPTIONAL BINARY element (UTF8);
- }
- }
- OPTIONAL GROUP my_list2 (LIST) {
- REPEATED GROUP list {
- REQUIRED BINARY element (UTF8);
- }
- }
- REQUIRED GROUP my_list3 (LIST) {
- REPEATED GROUP list {
- REQUIRED BINARY element (UTF8);
- }
- }
- }
- ";
-
- // // List<String> (list non-null, elements nullable)
- // required group my_list1 (LIST) {
- // repeated group list {
- // optional binary element (UTF8);
- // }
- // }
- {
- arrow_fields.push(Field::new(
- "my_list1",
- DataType::List(Box::new(Field::new("element", DataType::Utf8, true))),
- false,
- ));
- }
-
- // // List<String> (list nullable, elements non-null)
- // optional group my_list2 (LIST) {
- // repeated group list {
- // required binary element (UTF8);
- // }
- // }
- {
- arrow_fields.push(Field::new(
- "my_list2",
- DataType::List(Box::new(Field::new("element", DataType::Utf8, false))),
- true,
- ));
- }
-
- // // List<String> (list non-null, elements non-null)
- // repeated group my_list3 (LIST) {
- // repeated group list {
- // required binary element (UTF8);
- // }
- // }
- {
- arrow_fields.push(Field::new(
- "my_list3",
- DataType::List(Box::new(Field::new("element", DataType::Utf8, false))),
- false,
- ));
- }
-
- let parquet_group_type = parse_message_type(message_type).unwrap();
-
- let parquet_schema = SchemaDescriptor::new(Arc::new(parquet_group_type));
- let converted_arrow_schema =
- parquet_to_arrow_schema(&parquet_schema, &None).unwrap();
- let converted_fields = converted_arrow_schema.fields();
-
- assert_eq!(arrow_fields.len(), converted_fields.len());
- for i in 0..arrow_fields.len() {
- assert_eq!(arrow_fields[i], converted_fields[i]);
- }
- }
-
- #[test]
- fn test_nested_schema() {
- let mut arrow_fields = Vec::new();
- {
- let group1_fields = vec![
- Field::new("leaf1", DataType::Boolean, false),
- Field::new("leaf2", DataType::Int32, false),
- ];
- let group1_struct =
- Field::new("group1", DataType::Struct(group1_fields), false);
- arrow_fields.push(group1_struct);
-
- let leaf3_field = Field::new("leaf3", DataType::Int64, false);
- arrow_fields.push(leaf3_field);
- }
-
- let message_type = "
- message test_schema {
- REQUIRED GROUP group1 {
- REQUIRED BOOLEAN leaf1;
- REQUIRED INT32 leaf2;
- }
- REQUIRED INT64 leaf3;
- }
- ";
- let parquet_group_type = parse_message_type(message_type).unwrap();
-
- let parquet_schema = SchemaDescriptor::new(Arc::new(parquet_group_type));
- let converted_arrow_schema =
- parquet_to_arrow_schema(&parquet_schema, &None).unwrap();
- let converted_fields = converted_arrow_schema.fields();
-
- assert_eq!(arrow_fields.len(), converted_fields.len());
- for i in 0..arrow_fields.len() {
- assert_eq!(arrow_fields[i], converted_fields[i]);
- }
- }
-
- #[test]
- fn test_nested_schema_partial() {
- let mut arrow_fields = Vec::new();
- {
- let group1_fields = vec![Field::new("leaf1", DataType::Int64, false)];
- let group1 = Field::new("group1", DataType::Struct(group1_fields), false);
- arrow_fields.push(group1);
-
- let group2_fields = vec![Field::new("leaf4", DataType::Int64, false)];
- let group2 = Field::new("group2", DataType::Struct(group2_fields), false);
- arrow_fields.push(group2);
-
- arrow_fields.push(Field::new("leaf5", DataType::Int64, false));
- }
-
- let message_type = "
- message test_schema {
- REQUIRED GROUP group1 {
- REQUIRED INT64 leaf1;
- REQUIRED INT64 leaf2;
- }
- REQUIRED GROUP group2 {
- REQUIRED INT64 leaf3;
- REQUIRED INT64 leaf4;
- }
- REQUIRED INT64 leaf5;
- }
- ";
- let parquet_group_type = parse_message_type(message_type).unwrap();
-
- // Expected partial arrow schema (columns 0, 3, 4):
- // required group group1 {
- // required int64 leaf1;
- // }
- // required group group2 {
- // required int64 leaf4;
- // }
- // required int64 leaf5;
-
- let parquet_schema = SchemaDescriptor::new(Arc::new(parquet_group_type));
- let converted_arrow_schema =
- parquet_to_arrow_schema_by_columns(&parquet_schema, vec![0, 3, 4], &None)
- .unwrap();
- let converted_fields = converted_arrow_schema.fields();
-
- assert_eq!(arrow_fields.len(), converted_fields.len());
- for i in 0..arrow_fields.len() {
- assert_eq!(arrow_fields[i], converted_fields[i]);
- }
- }
-
- #[test]
- fn test_nested_schema_partial_ordering() {
- let mut arrow_fields = Vec::new();
- {
- let group2_fields = vec![Field::new("leaf4", DataType::Int64, false)];
- let group2 = Field::new("group2", DataType::Struct(group2_fields), false);
- arrow_fields.push(group2);
-
- arrow_fields.push(Field::new("leaf5", DataType::Int64, false));
-
- let group1_fields = vec![Field::new("leaf1", DataType::Int64, false)];
- let group1 = Field::new("group1", DataType::Struct(group1_fields), false);
- arrow_fields.push(group1);
- }
-
- let message_type = "
- message test_schema {
- REQUIRED GROUP group1 {
- REQUIRED INT64 leaf1;
- REQUIRED INT64 leaf2;
- }
- REQUIRED GROUP group2 {
- REQUIRED INT64 leaf3;
- REQUIRED INT64 leaf4;
- }
- REQUIRED INT64 leaf5;
- }
- ";
- let parquet_group_type = parse_message_type(message_type).unwrap();
-
- // Expected partial arrow schema (columns 3, 4, 0):
- // required group group1 {
- // required int64 leaf1;
- // }
- // required group group2 {
- // required int64 leaf4;
- // }
- // required int64 leaf5;
-
- let parquet_schema = SchemaDescriptor::new(Arc::new(parquet_group_type));
- let converted_arrow_schema =
- parquet_to_arrow_schema_by_columns(&parquet_schema, vec![3, 4, 0], &None)
- .unwrap();
- let converted_fields = converted_arrow_schema.fields();
-
- assert_eq!(arrow_fields.len(), converted_fields.len());
- for i in 0..arrow_fields.len() {
- assert_eq!(arrow_fields[i], converted_fields[i]);
- }
- }
-
- #[test]
- fn test_repeated_nested_schema() {
- let mut arrow_fields = Vec::new();
- {
- arrow_fields.push(Field::new("leaf1", DataType::Int32, true));
-
- let inner_group_list = Field::new(
- "innerGroup",
- DataType::List(Box::new(Field::new(
- "innerGroup",
- DataType::Struct(vec![Field::new("leaf3", DataType::Int32, true)]),
- true,
- ))),
- true,
- );
-
- let outer_group_list = Field::new(
- "outerGroup",
- DataType::List(Box::new(Field::new(
- "outerGroup",
- DataType::Struct(vec![
- Field::new("leaf2", DataType::Int32, true),
- inner_group_list,
- ]),
- true,
- ))),
- true,
- );
- arrow_fields.push(outer_group_list);
- }
-
- let message_type = "
- message test_schema {
- OPTIONAL INT32 leaf1;
- REPEATED GROUP outerGroup {
- OPTIONAL INT32 leaf2;
- REPEATED GROUP innerGroup {
- OPTIONAL INT32 leaf3;
- }
- }
- }
- ";
- let parquet_group_type = parse_message_type(message_type).unwrap();
-
- let parquet_schema = SchemaDescriptor::new(Arc::new(parquet_group_type));
- let converted_arrow_schema =
- parquet_to_arrow_schema(&parquet_schema, &None).unwrap();
- let converted_fields = converted_arrow_schema.fields();
-
- assert_eq!(arrow_fields.len(), converted_fields.len());
- for i in 0..arrow_fields.len() {
- assert_eq!(arrow_fields[i], converted_fields[i]);
- }
- }
-
- #[test]
- fn test_column_desc_to_field() {
- let message_type = "
- message test_schema {
- REQUIRED BOOLEAN boolean;
- REQUIRED INT32 int8 (INT_8);
- REQUIRED INT32 uint8 (INTEGER(8,false));
- REQUIRED INT32 int16 (INT_16);
- REQUIRED INT32 uint16 (INTEGER(16,false));
- REQUIRED INT32 int32;
- REQUIRED INT64 int64;
- OPTIONAL DOUBLE double;
- OPTIONAL FLOAT float;
- OPTIONAL BINARY string (UTF8);
- REPEATED BOOLEAN bools;
- OPTIONAL INT32 date (DATE);
- OPTIONAL INT32 time_milli (TIME_MILLIS);
- OPTIONAL INT64 time_micro (TIME_MICROS);
- OPTIONAL INT64 time_nano (TIME(NANOS,false));
- OPTIONAL INT64 ts_milli (TIMESTAMP_MILLIS);
- REQUIRED INT64 ts_micro (TIMESTAMP_MICROS);
- REQUIRED INT64 ts_nano (TIMESTAMP(NANOS,true));
- }
- ";
- let parquet_group_type = parse_message_type(message_type).unwrap();
-
- let parquet_schema = SchemaDescriptor::new(Arc::new(parquet_group_type));
- let converted_arrow_fields = parquet_schema
- .columns()
- .iter()
- .map(|c| parquet_to_arrow_field(c).unwrap())
- .collect::<Vec<Field>>();
-
- let arrow_fields = vec![
- Field::new("boolean", DataType::Boolean, false),
- Field::new("int8", DataType::Int8, false),
- Field::new("uint8", DataType::UInt8, false),
- Field::new("int16", DataType::Int16, false),
- Field::new("uint16", DataType::UInt16, false),
- Field::new("int32", DataType::Int32, false),
- Field::new("int64", DataType::Int64, false),
- Field::new("double", DataType::Float64, true),
- Field::new("float", DataType::Float32, true),
- Field::new("string", DataType::Utf8, true),
- Field::new(
- "bools",
- DataType::List(Box::new(Field::new("bools", DataType::Boolean, true))),
- true,
- ),
- Field::new("date", DataType::Date32, true),
- Field::new("time_milli", DataType::Time32(TimeUnit::Millisecond), true),
- Field::new("time_micro", DataType::Time64(TimeUnit::Microsecond), true),
- Field::new("time_nano", DataType::Time64(TimeUnit::Nanosecond), true),
- Field::new(
- "ts_milli",
- DataType::Timestamp(TimeUnit::Millisecond, None),
- true,
- ),
- Field::new(
- "ts_micro",
- DataType::Timestamp(TimeUnit::Microsecond, None),
- false,
- ),
- Field::new(
- "ts_nano",
- DataType::Timestamp(TimeUnit::Nanosecond, Some("UTC".to_string())),
- false,
- ),
- ];
-
- assert_eq!(arrow_fields, converted_arrow_fields);
- }
-
- #[test]
- fn test_field_to_column_desc() {
- let message_type = "
- message arrow_schema {
- REQUIRED BOOLEAN boolean;
- REQUIRED INT32 int8 (INT_8);
- REQUIRED INT32 int16 (INTEGER(16,true));
- REQUIRED INT32 int32;
- REQUIRED INT64 int64;
- OPTIONAL DOUBLE double;
- OPTIONAL FLOAT float;
- OPTIONAL BINARY string (STRING);
- OPTIONAL GROUP bools (LIST) {
- REPEATED GROUP list {
- OPTIONAL BOOLEAN element;
- }
- }
- REQUIRED GROUP bools_non_null (LIST) {
- REPEATED GROUP list {
- REQUIRED BOOLEAN element;
- }
- }
- OPTIONAL INT32 date (DATE);
- OPTIONAL INT32 time_milli (TIME(MILLIS,false));
- OPTIONAL INT64 time_micro (TIME_MICROS);
- OPTIONAL INT64 ts_milli (TIMESTAMP_MILLIS);
- REQUIRED INT64 ts_micro (TIMESTAMP(MICROS,false));
- REQUIRED GROUP struct {
- REQUIRED BOOLEAN bools;
- REQUIRED INT32 uint32 (INTEGER(32,false));
- REQUIRED GROUP int32 (LIST) {
- REPEATED GROUP list {
- OPTIONAL INT32 element;
- }
- }
- }
- REQUIRED BINARY dictionary_strings (STRING);
- }
- ";
- let parquet_group_type = parse_message_type(message_type).unwrap();
-
- let parquet_schema = SchemaDescriptor::new(Arc::new(parquet_group_type));
-
- let arrow_fields = vec![
- Field::new("boolean", DataType::Boolean, false),
- Field::new("int8", DataType::Int8, false),
- Field::new("int16", DataType::Int16, false),
- Field::new("int32", DataType::Int32, false),
- Field::new("int64", DataType::Int64, false),
- Field::new("double", DataType::Float64, true),
- Field::new("float", DataType::Float32, true),
- Field::new("string", DataType::Utf8, true),
- Field::new(
- "bools",
- DataType::List(Box::new(Field::new("element", DataType::Boolean, true))),
- true,
- ),
- Field::new(
- "bools_non_null",
- DataType::List(Box::new(Field::new("element", DataType::Boolean, false))),
- false,
- ),
- Field::new("date", DataType::Date32, true),
- Field::new("time_milli", DataType::Time32(TimeUnit::Millisecond), true),
- Field::new("time_micro", DataType::Time64(TimeUnit::Microsecond), true),
- Field::new(
- "ts_milli",
- DataType::Timestamp(TimeUnit::Millisecond, None),
- true,
- ),
- Field::new(
- "ts_micro",
- DataType::Timestamp(TimeUnit::Microsecond, None),
- false,
- ),
- Field::new(
- "struct",
- DataType::Struct(vec![
- Field::new("bools", DataType::Boolean, false),
- Field::new("uint32", DataType::UInt32, false),
- Field::new(
- "int32",
- DataType::List(Box::new(Field::new(
- "element",
- DataType::Int32,
- true,
- ))),
- false,
- ),
- ]),
- false,
- ),
- Field::new(
- "dictionary_strings",
- DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
- false,
- ),
- ];
- let arrow_schema = Schema::new(arrow_fields);
- let converted_arrow_schema = arrow_to_parquet_schema(&arrow_schema).unwrap();
-
- assert_eq!(
- parquet_schema.columns().len(),
- converted_arrow_schema.columns().len()
- );
- parquet_schema
- .columns()
- .iter()
- .zip(converted_arrow_schema.columns())
- .for_each(|(a, b)| {
- // Only check logical type if it's set on the Parquet side.
- // This is because the Arrow conversion always sets logical type,
- // even if there wasn't originally one.
- // This is not an issue, but is an inconvenience for this test.
- match a.logical_type() {
- Some(_) => {
- assert_eq!(a, b)
- }
- None => {
- assert_eq!(a.name(), b.name());
- assert_eq!(a.physical_type(), b.physical_type());
- assert_eq!(a.converted_type(), b.converted_type());
- }
- };
- });
- }
-
- #[test]
- #[should_panic(expected = "Parquet does not support writing empty structs")]
- fn test_empty_struct_field() {
- let arrow_fields = vec![Field::new("struct", DataType::Struct(vec![]), false)];
- let arrow_schema = Schema::new(arrow_fields);
- let converted_arrow_schema = arrow_to_parquet_schema(&arrow_schema);
-
- assert!(converted_arrow_schema.is_err());
- converted_arrow_schema.unwrap();
- }
-
- #[test]
- fn test_metadata() {
- let message_type = "
- message test_schema {
- OPTIONAL BINARY string (STRING);
- }
- ";
- let parquet_group_type = parse_message_type(message_type).unwrap();
-
- let mut key_value_metadata: Vec<KeyValue> = Vec::new();
- key_value_metadata.push(KeyValue::new("foo".to_owned(), Some("bar".to_owned())));
- key_value_metadata.push(KeyValue::new("baz".to_owned(), None));
-
- let mut expected_metadata: HashMap<String, String> = HashMap::new();
- expected_metadata.insert("foo".to_owned(), "bar".to_owned());
-
- let parquet_schema = SchemaDescriptor::new(Arc::new(parquet_group_type));
- let converted_arrow_schema =
- parquet_to_arrow_schema(&parquet_schema, &Some(key_value_metadata)).unwrap();
-
- assert_eq!(converted_arrow_schema.metadata(), &expected_metadata);
- }
-
- #[test]
- fn test_arrow_schema_roundtrip() -> Result<()> {
- // This tests the roundtrip of an Arrow schema
- // Fields that are commented out fail roundtrip tests or are unsupported by the writer
- let metadata: HashMap<String, String> =
- [("Key".to_string(), "Value".to_string())]
- .iter()
- .cloned()
- .collect();
-
- let schema = Schema::new_with_metadata(
- vec![
- Field::new("c1", DataType::Utf8, false),
- Field::new("c2", DataType::Binary, false),
- Field::new("c3", DataType::FixedSizeBinary(3), false),
- Field::new("c4", DataType::Boolean, false),
- Field::new("c5", DataType::Date32, false),
- Field::new("c6", DataType::Date64, false),
- Field::new("c7", DataType::Time32(TimeUnit::Second), false),
- Field::new("c8", DataType::Time32(TimeUnit::Millisecond), false),
- Field::new("c13", DataType::Time64(TimeUnit::Microsecond), false),
- Field::new("c14", DataType::Time64(TimeUnit::Nanosecond), false),
- Field::new("c15", DataType::Timestamp(TimeUnit::Second, None), false),
- Field::new(
- "c16",
- DataType::Timestamp(TimeUnit::Millisecond, Some("UTC".to_string())),
- false,
- ),
- Field::new(
- "c17",
- DataType::Timestamp(
- TimeUnit::Microsecond,
- Some("Africa/Johannesburg".to_string()),
- ),
- false,
- ),
- Field::new(
- "c18",
- DataType::Timestamp(TimeUnit::Nanosecond, None),
- false,
- ),
- Field::new("c19", DataType::Interval(IntervalUnit::DayTime), false),
- Field::new("c20", DataType::Interval(IntervalUnit::YearMonth), false),
- Field::new(
- "c21",
- DataType::List(Box::new(Field::new("list", DataType::Boolean, true))),
- false,
- ),
- // Field::new(
- // "c22",
- // DataType::FixedSizeList(Box::new(DataType::Boolean), 5),
- // false,
- // ),
- // Field::new(
- // "c23",
- // DataType::List(Box::new(DataType::LargeList(Box::new(
- // DataType::Struct(vec![
- // Field::new("a", DataType::Int16, true),
- // Field::new("b", DataType::Float64, false),
- // ]),
- // )))),
- // true,
- // ),
- Field::new(
- "c24",
- DataType::Struct(vec![
- Field::new("a", DataType::Utf8, false),
- Field::new("b", DataType::UInt16, false),
- ]),
- false,
- ),
- Field::new("c25", DataType::Interval(IntervalUnit::YearMonth), true),
- Field::new("c26", DataType::Interval(IntervalUnit::DayTime), true),
- // Field::new("c27", DataType::Duration(TimeUnit::Second), false),
- // Field::new("c28", DataType::Duration(TimeUnit::Millisecond), false),
- // Field::new("c29", DataType::Duration(TimeUnit::Microsecond), false),
- // Field::new("c30", DataType::Duration(TimeUnit::Nanosecond), false),
- Field::new_dict(
- "c31",
- DataType::Dictionary(
- Box::new(DataType::Int32),
- Box::new(DataType::Utf8),
- ),
- true,
- 123,
- true,
- ),
- Field::new("c32", DataType::LargeBinary, true),
- Field::new("c33", DataType::LargeUtf8, true),
- // Field::new(
- // "c34",
- // DataType::LargeList(Box::new(DataType::List(Box::new(
- // DataType::Struct(vec![
- // Field::new("a", DataType::Int16, true),
- // Field::new("b", DataType::Float64, true),
- // ]),
- // )))),
- // true,
- // ),
- Field::new("c35", DataType::Null, true),
- Field::new("c36", DataType::Decimal(2, 1), false),
- Field::new("c37", DataType::Decimal(50, 20), false),
- Field::new("c38", DataType::Decimal(18, 12), true),
- ],
- metadata,
- );
-
- // write to an empty parquet file so that schema is serialized
- let file = get_temp_file("test_arrow_schema_roundtrip.parquet", &[]);
- let mut writer = ArrowWriter::try_new(
- file.try_clone().unwrap(),
- Arc::new(schema.clone()),
- None,
- )?;
- writer.close()?;
-
- // read file back
- let parquet_reader = SerializedFileReader::try_from(file)?;
- let mut arrow_reader = ParquetFileArrowReader::new(Arc::new(parquet_reader));
- let read_schema = arrow_reader.get_schema()?;
- assert_eq!(schema, read_schema);
-
- // read all fields by columns
- let partial_read_schema =
- arrow_reader.get_schema_by_columns(0..(schema.fields().len()), false)?;
- assert_eq!(schema, partial_read_schema);
-
- Ok(())
- }
-
- #[test]
- #[ignore = "Roundtrip of lists currently fails because we don't check their types correctly in the Arrow schema"]
- fn test_arrow_schema_roundtrip_lists() -> Result<()> {
- let metadata: HashMap<String, String> =
- [("Key".to_string(), "Value".to_string())]
- .iter()
- .cloned()
- .collect();
-
- let schema = Schema::new_with_metadata(
- vec![
- Field::new(
- "c21",
- DataType::List(Box::new(Field::new(
- "array",
- DataType::Boolean,
- true,
- ))),
- false,
- ),
- Field::new(
- "c22",
- DataType::FixedSizeList(
- Box::new(Field::new("items", DataType::Boolean, false)),
- 5,
- ),
- false,
- ),
- Field::new(
- "c23",
- DataType::List(Box::new(Field::new(
- "items",
- DataType::LargeList(Box::new(Field::new(
- "items",
- DataType::Struct(vec![
- Field::new("a", DataType::Int16, true),
- Field::new("b", DataType::Float64, false),
- ]),
- true,
- ))),
- true,
- ))),
- true,
- ),
- ],
- metadata,
- );
-
- // write to an empty parquet file so that schema is serialized
- let file = get_temp_file("test_arrow_schema_roundtrip_lists.parquet", &[]);
- let mut writer = ArrowWriter::try_new(
- file.try_clone().unwrap(),
- Arc::new(schema.clone()),
- None,
- )?;
- writer.close()?;
-
- // read file back
- let parquet_reader = SerializedFileReader::try_from(file)?;
- let mut arrow_reader = ParquetFileArrowReader::new(Arc::new(parquet_reader));
- let read_schema = arrow_reader.get_schema()?;
- assert_eq!(schema, read_schema);
-
- // read all fields by columns
- let partial_read_schema =
- arrow_reader.get_schema_by_columns(0..(schema.fields().len()), false)?;
- assert_eq!(schema, partial_read_schema);
-
- Ok(())
- }
-}
diff --git a/parquet/src/basic.rs b/parquet/src/basic.rs
deleted file mode 100644
index 631257e..0000000
--- a/parquet/src/basic.rs
+++ /dev/null
@@ -1,1969 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-//! Contains Rust mappings for Thrift definition.
-//! Refer to `parquet.thrift` file to see raw definitions.
-
-use std::{convert, fmt, result, str};
-
-use parquet_format as parquet;
-
-use crate::errors::ParquetError;
-
-// Re-export parquet_format types used in this module
-pub use parquet_format::{
- BsonType, DateType, DecimalType, EnumType, IntType, JsonType, ListType, MapType,
- NullType, StringType, TimeType, TimeUnit, TimestampType, UUIDType,
-};
-
-// ----------------------------------------------------------------------
-// Types from the Thrift definition
-
-// ----------------------------------------------------------------------
-// Mirrors `parquet::Type`
-
-/// Types supported by Parquet.
-/// These physical types are intended to be used in combination with the encodings to
-/// control the on disk storage format.
-/// For example INT16 is not included as a type since a good encoding of INT32
-/// would handle this.
-#[derive(Debug, Clone, Copy, PartialEq)]
-pub enum Type {
- BOOLEAN,
- INT32,
- INT64,
- INT96,
- FLOAT,
- DOUBLE,
- BYTE_ARRAY,
- FIXED_LEN_BYTE_ARRAY,
-}
-
-// ----------------------------------------------------------------------
-// Mirrors `parquet::ConvertedType`
-
-/// Common types (converted types) used by frameworks when using Parquet.
-/// This helps map between types in those frameworks to the base types in Parquet.
-/// This is only metadata and not needed to read or write the data.
-///
-/// This struct was renamed from `LogicalType` in version 4.0.0.
-/// If targeting Parquet format 2.4.0 or above, please use [LogicalType] instead.
-#[derive(Debug, Clone, Copy, PartialEq)]
-pub enum ConvertedType {
- NONE,
- /// A BYTE_ARRAY actually contains UTF8 encoded chars.
- UTF8,
-
- /// A map is converted as an optional field containing a repeated key/value pair.
- MAP,
-
- /// A key/value pair is converted into a group of two fields.
- MAP_KEY_VALUE,
-
- /// A list is converted into an optional field containing a repeated field for its
- /// values.
- LIST,
-
- /// An enum is converted into a binary field
- ENUM,
-
- /// A decimal value.
- /// This may be used to annotate binary or fixed primitive types. The
- /// underlying byte array stores the unscaled value encoded as two's
- /// complement using big-endian byte order (the most significant byte is the
- /// zeroth element).
- ///
- /// This must be accompanied by a (maximum) precision and a scale in the
- /// SchemaElement. The precision specifies the number of digits in the decimal
- /// and the scale stores the location of the decimal point. For example 1.23
- /// would have precision 3 (3 total digits) and scale 2 (the decimal point is
- /// 2 digits over).
- DECIMAL,
-
- /// A date stored as days since Unix epoch, encoded as the INT32 physical type.
- DATE,
-
- /// The total number of milliseconds since midnight. The value is stored as an INT32
- /// physical type.
- TIME_MILLIS,
-
- /// The total number of microseconds since midnight. The value is stored as an INT64
- /// physical type.
- TIME_MICROS,
-
- /// Date and time recorded as milliseconds since the Unix epoch.
- /// Recorded as a physical type of INT64.
- TIMESTAMP_MILLIS,
-
- /// Date and time recorded as microseconds since the Unix epoch.
- /// The value is stored as an INT64 physical type.
- TIMESTAMP_MICROS,
-
- /// An unsigned 8 bit integer value stored as INT32 physical type.
- UINT_8,
-
- /// An unsigned 16 bit integer value stored as INT32 physical type.
- UINT_16,
-
- /// An unsigned 32 bit integer value stored as INT32 physical type.
- UINT_32,
-
- /// An unsigned 64 bit integer value stored as INT64 physical type.
- UINT_64,
-
- /// A signed 8 bit integer value stored as INT32 physical type.
- INT_8,
-
- /// A signed 16 bit integer value stored as INT32 physical type.
- INT_16,
-
- /// A signed 32 bit integer value stored as INT32 physical type.
- INT_32,
-
- /// A signed 64 bit integer value stored as INT64 physical type.
- INT_64,
-
- /// A JSON document embedded within a single UTF8 column.
- JSON,
-
- /// A BSON document embedded within a single BINARY column.
- BSON,
-
- /// An interval of time.
- ///
- /// This type annotates data stored as a FIXED_LEN_BYTE_ARRAY of length 12.
- /// This data is composed of three separate little endian unsigned integers.
- /// Each stores a component of a duration of time. The first integer identifies
- /// the number of months associated with the duration, the second identifies
- /// the number of days associated with the duration and the third identifies
- /// the number of milliseconds associated with the provided duration.
- /// This duration of time is independent of any particular timezone or date.
- INTERVAL,
-}
-
-// ----------------------------------------------------------------------
-// Mirrors `parquet::LogicalType`
-
-/// Logical types used by version 2.4.0+ of the Parquet format.
-///
-/// This is an *entirely new* struct as of version
-/// 4.0.0. The struct previously named `LogicalType` was renamed to
-/// [`ConvertedType`]. Please see the README.md for more details.
-#[derive(Debug, Clone, PartialEq)]
-pub enum LogicalType {
- STRING(StringType),
- MAP(MapType),
- LIST(ListType),
- ENUM(EnumType),
- DECIMAL(DecimalType),
- DATE(DateType),
- TIME(TimeType),
- TIMESTAMP(TimestampType),
- INTEGER(IntType),
- UNKNOWN(NullType),
- JSON(JsonType),
- BSON(BsonType),
- UUID(UUIDType),
-}
-
-// ----------------------------------------------------------------------
-// Mirrors `parquet::FieldRepetitionType`
-
-/// Representation of field types in schema.
-#[derive(Debug, Clone, Copy, PartialEq)]
-pub enum Repetition {
- /// Field is required (can not be null) and each record has exactly 1 value.
- REQUIRED,
- /// Field is optional (can be null) and each record has 0 or 1 values.
- OPTIONAL,
- /// Field is repeated and can contain 0 or more values.
- REPEATED,
-}
-
-// ----------------------------------------------------------------------
-// Mirrors `parquet::Encoding`
-
-/// Encodings supported by Parquet.
-/// Not all encodings are valid for all types. These enums are also used to specify the
-/// encoding of definition and repetition levels.
-#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
-pub enum Encoding {
- /// Default byte encoding.
- /// - BOOLEAN - 1 bit per value, 0 is false; 1 is true.
- /// - INT32 - 4 bytes per value, stored as little-endian.
- /// - INT64 - 8 bytes per value, stored as little-endian.
- /// - FLOAT - 4 bytes per value, stored as little-endian.
- /// - DOUBLE - 8 bytes per value, stored as little-endian.
- /// - BYTE_ARRAY - 4 byte length stored as little endian, followed by bytes.
- /// - FIXED_LEN_BYTE_ARRAY - just the bytes are stored.
- PLAIN,
-
- /// **Deprecated** dictionary encoding.
- ///
- /// The values in the dictionary are encoded using PLAIN encoding.
- /// Since it is deprecated, RLE_DICTIONARY encoding is used for a data page, and
- /// PLAIN encoding is used for dictionary page.
- PLAIN_DICTIONARY,
-
- /// Group packed run length encoding.
- ///
- /// Usable for definition/repetition levels encoding and boolean values.
- RLE,
-
- /// Bit packed encoding.
- ///
- /// This can only be used if the data has a known max width.
- /// Usable for definition/repetition levels encoding.
- BIT_PACKED,
-
- /// Delta encoding for integers, either INT32 or INT64.
- ///
- /// Works best on sorted data.
- DELTA_BINARY_PACKED,
-
- /// Encoding for byte arrays to separate the length values and the data.
- ///
- /// The lengths are encoded using DELTA_BINARY_PACKED encoding.
- DELTA_LENGTH_BYTE_ARRAY,
-
- /// Incremental encoding for byte arrays.
- ///
- /// Prefix lengths are encoded using DELTA_BINARY_PACKED encoding.
- /// Suffixes are stored using DELTA_LENGTH_BYTE_ARRAY encoding.
- DELTA_BYTE_ARRAY,
-
- /// Dictionary encoding.
- ///
- /// The ids are encoded using the RLE encoding.
- RLE_DICTIONARY,
-}
-
-// ----------------------------------------------------------------------
-// Mirrors `parquet::CompressionCodec`
-
-/// Supported compression algorithms.
-#[derive(Debug, Clone, Copy, PartialEq)]
-pub enum Compression {
- UNCOMPRESSED,
- SNAPPY,
- GZIP,
- LZO,
- BROTLI,
- LZ4,
- ZSTD,
-}
-
-// ----------------------------------------------------------------------
-// Mirrors `parquet::PageType`
-
-/// Available data pages for Parquet file format.
-/// Note that some of the page types may not be supported.
-#[derive(Debug, Clone, Copy, PartialEq)]
-pub enum PageType {
- DATA_PAGE,
- INDEX_PAGE,
- DICTIONARY_PAGE,
- DATA_PAGE_V2,
-}
-
-// ----------------------------------------------------------------------
-// Mirrors `parquet::ColumnOrder`
-
-/// Sort order for page and column statistics.
-///
-/// Types are associated with sort orders and column stats are aggregated using a sort
-/// order, and a sort order should be considered when comparing values with statistics
-/// min/max.
-///
-/// See reference in
-/// <https://github.com/apache/parquet-cpp/blob/master/src/parquet/types.h>
-#[derive(Debug, Clone, Copy, PartialEq)]
-pub enum SortOrder {
- /// Signed (either value or legacy byte-wise) comparison.
- SIGNED,
- /// Unsigned (depending on physical type either value or byte-wise) comparison.
- UNSIGNED,
- /// Comparison is undefined.
- UNDEFINED,
-}
-
-/// Column order that specifies what method was used to aggregate min/max values for
-/// statistics.
-///
-/// If column order is undefined, then it is the legacy behaviour and all values should
-/// be compared as signed values/bytes.
-#[derive(Debug, Clone, Copy, PartialEq)]
-pub enum ColumnOrder {
- /// Column uses the order defined by its logical or physical type
- /// (if there is no logical type), parquet-format 2.4.0+.
- TYPE_DEFINED_ORDER(SortOrder),
- /// Undefined column order, means legacy behaviour before parquet-format 2.4.0.
- /// Sort order is always SIGNED.
- UNDEFINED,
-}
-
-impl ColumnOrder {
- /// Returns sort order for a physical/logical type.
- pub fn get_sort_order(
- logical_type: Option<LogicalType>,
- converted_type: ConvertedType,
- physical_type: Type,
- ) -> SortOrder {
- // TODO: Should this take converted and logical type, for compatibility?
- match logical_type {
- Some(logical) => match logical {
- LogicalType::STRING(_)
- | LogicalType::ENUM(_)
- | LogicalType::JSON(_)
- | LogicalType::BSON(_) => SortOrder::UNSIGNED,
- LogicalType::INTEGER(t) => match t.is_signed {
- true => SortOrder::SIGNED,
- false => SortOrder::UNSIGNED,
- },
- LogicalType::MAP(_) | LogicalType::LIST(_) => SortOrder::UNDEFINED,
- LogicalType::DECIMAL(_) => SortOrder::SIGNED,
- LogicalType::DATE(_) => SortOrder::SIGNED,
- LogicalType::TIME(_) => SortOrder::SIGNED,
- LogicalType::TIMESTAMP(_) => SortOrder::SIGNED,
- LogicalType::UNKNOWN(_) => SortOrder::UNDEFINED,
- LogicalType::UUID(_) => SortOrder::UNSIGNED,
- },
- // Fall back to converted type
- None => Self::get_converted_sort_order(converted_type, physical_type),
- }
- }
-
- fn get_converted_sort_order(
- converted_type: ConvertedType,
- physical_type: Type,
- ) -> SortOrder {
- match converted_type {
- // Unsigned byte-wise comparison.
- ConvertedType::UTF8
- | ConvertedType::JSON
- | ConvertedType::BSON
- | ConvertedType::ENUM => SortOrder::UNSIGNED,
-
- ConvertedType::INT_8
- | ConvertedType::INT_16
- | ConvertedType::INT_32
- | ConvertedType::INT_64 => SortOrder::SIGNED,
-
- ConvertedType::UINT_8
- | ConvertedType::UINT_16
- | ConvertedType::UINT_32
- | ConvertedType::UINT_64 => SortOrder::UNSIGNED,
-
- // Signed comparison of the represented value.
- ConvertedType::DECIMAL => SortOrder::SIGNED,
-
- ConvertedType::DATE => SortOrder::SIGNED,
-
- ConvertedType::TIME_MILLIS
- | ConvertedType::TIME_MICROS
- | ConvertedType::TIMESTAMP_MILLIS
- | ConvertedType::TIMESTAMP_MICROS => SortOrder::SIGNED,
-
- ConvertedType::INTERVAL => SortOrder::UNDEFINED,
-
- ConvertedType::LIST | ConvertedType::MAP | ConvertedType::MAP_KEY_VALUE => {
- SortOrder::UNDEFINED
- }
-
- // Fall back to physical type.
- ConvertedType::NONE => Self::get_default_sort_order(physical_type),
- }
- }
-
- /// Returns default sort order based on physical type.
- fn get_default_sort_order(physical_type: Type) -> SortOrder {
- match physical_type {
- // Order: false, true
- Type::BOOLEAN => SortOrder::UNSIGNED,
- Type::INT32 | Type::INT64 => SortOrder::SIGNED,
- Type::INT96 => SortOrder::UNDEFINED,
- // Notes to remember when comparing float/double values:
- // If the min is a NaN, it should be ignored.
- // If the max is a NaN, it should be ignored.
- // If the min is +0, the row group may contain -0 values as well.
- // If the max is -0, the row group may contain +0 values as well.
- // When looking for NaN values, min and max should be ignored.
- Type::FLOAT | Type::DOUBLE => SortOrder::SIGNED,
- // Unsigned byte-wise comparison
- Type::BYTE_ARRAY | Type::FIXED_LEN_BYTE_ARRAY => SortOrder::UNSIGNED,
- }
- }
-
- /// Returns sort order associated with this column order.
- pub fn sort_order(&self) -> SortOrder {
- match *self {
- ColumnOrder::TYPE_DEFINED_ORDER(order) => order,
- ColumnOrder::UNDEFINED => SortOrder::SIGNED,
- }
- }
-}
-
-impl fmt::Display for Type {
- fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
- write!(f, "{:?}", self)
- }
-}
-
-impl fmt::Display for ConvertedType {
- fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
- write!(f, "{:?}", self)
- }
-}
-
-impl fmt::Display for Repetition {
- fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
- write!(f, "{:?}", self)
- }
-}
-
-impl fmt::Display for Encoding {
- fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
- write!(f, "{:?}", self)
- }
-}
-
-impl fmt::Display for Compression {
- fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
- write!(f, "{:?}", self)
- }
-}
-
-impl fmt::Display for PageType {
- fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
- write!(f, "{:?}", self)
- }
-}
-
-impl fmt::Display for SortOrder {
- fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
- write!(f, "{:?}", self)
- }
-}
-
-impl fmt::Display for ColumnOrder {
- fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
- write!(f, "{:?}", self)
- }
-}
-
-// ----------------------------------------------------------------------
-// parquet::Type <=> Type conversion
-
-impl convert::From<parquet::Type> for Type {
- fn from(value: parquet::Type) -> Self {
- match value {
- parquet::Type::Boolean => Type::BOOLEAN,
- parquet::Type::Int32 => Type::INT32,
- parquet::Type::Int64 => Type::INT64,
- parquet::Type::Int96 => Type::INT96,
- parquet::Type::Float => Type::FLOAT,
- parquet::Type::Double => Type::DOUBLE,
- parquet::Type::ByteArray => Type::BYTE_ARRAY,
- parquet::Type::FixedLenByteArray => Type::FIXED_LEN_BYTE_ARRAY,
- }
- }
-}
-
-impl convert::From<Type> for parquet::Type {
- fn from(value: Type) -> Self {
- match value {
- Type::BOOLEAN => parquet::Type::Boolean,
- Type::INT32 => parquet::Type::Int32,
- Type::INT64 => parquet::Type::Int64,
- Type::INT96 => parquet::Type::Int96,
- Type::FLOAT => parquet::Type::Float,
- Type::DOUBLE => parquet::Type::Double,
- Type::BYTE_ARRAY => parquet::Type::ByteArray,
- Type::FIXED_LEN_BYTE_ARRAY => parquet::Type::FixedLenByteArray,
- }
- }
-}
-
-// ----------------------------------------------------------------------
-// parquet::ConvertedType <=> ConvertedType conversion
-
-impl convert::From<Option<parquet::ConvertedType>> for ConvertedType {
- fn from(option: Option<parquet::ConvertedType>) -> Self {
- match option {
- None => ConvertedType::NONE,
- Some(value) => match value {
- parquet::ConvertedType::Utf8 => ConvertedType::UTF8,
- parquet::ConvertedType::Map => ConvertedType::MAP,
- parquet::ConvertedType::MapKeyValue => ConvertedType::MAP_KEY_VALUE,
- parquet::ConvertedType::List => ConvertedType::LIST,
- parquet::ConvertedType::Enum => ConvertedType::ENUM,
- parquet::ConvertedType::Decimal => ConvertedType::DECIMAL,
- parquet::ConvertedType::Date => ConvertedType::DATE,
- parquet::ConvertedType::TimeMillis => ConvertedType::TIME_MILLIS,
- parquet::ConvertedType::TimeMicros => ConvertedType::TIME_MICROS,
- parquet::ConvertedType::TimestampMillis => {
- ConvertedType::TIMESTAMP_MILLIS
- }
- parquet::ConvertedType::TimestampMicros => {
- ConvertedType::TIMESTAMP_MICROS
- }
- parquet::ConvertedType::Uint8 => ConvertedType::UINT_8,
- parquet::ConvertedType::Uint16 => ConvertedType::UINT_16,
- parquet::ConvertedType::Uint32 => ConvertedType::UINT_32,
- parquet::ConvertedType::Uint64 => ConvertedType::UINT_64,
- parquet::ConvertedType::Int8 => ConvertedType::INT_8,
- parquet::ConvertedType::Int16 => ConvertedType::INT_16,
- parquet::ConvertedType::Int32 => ConvertedType::INT_32,
- parquet::ConvertedType::Int64 => ConvertedType::INT_64,
- parquet::ConvertedType::Json => ConvertedType::JSON,
- parquet::ConvertedType::Bson => ConvertedType::BSON,
- parquet::ConvertedType::Interval => ConvertedType::INTERVAL,
- },
- }
- }
-}
-
-impl convert::From<ConvertedType> for Option<parquet::ConvertedType> {
- fn from(value: ConvertedType) -> Self {
- match value {
- ConvertedType::NONE => None,
- ConvertedType::UTF8 => Some(parquet::ConvertedType::Utf8),
- ConvertedType::MAP => Some(parquet::ConvertedType::Map),
- ConvertedType::MAP_KEY_VALUE => Some(parquet::ConvertedType::MapKeyValue),
- ConvertedType::LIST => Some(parquet::ConvertedType::List),
- ConvertedType::ENUM => Some(parquet::ConvertedType::Enum),
- ConvertedType::DECIMAL => Some(parquet::ConvertedType::Decimal),
- ConvertedType::DATE => Some(parquet::ConvertedType::Date),
- ConvertedType::TIME_MILLIS => Some(parquet::ConvertedType::TimeMillis),
- ConvertedType::TIME_MICROS => Some(parquet::ConvertedType::TimeMicros),
- ConvertedType::TIMESTAMP_MILLIS => {
- Some(parquet::ConvertedType::TimestampMillis)
- }
- ConvertedType::TIMESTAMP_MICROS => {
- Some(parquet::ConvertedType::TimestampMicros)
- }
- ConvertedType::UINT_8 => Some(parquet::ConvertedType::Uint8),
- ConvertedType::UINT_16 => Some(parquet::ConvertedType::Uint16),
- ConvertedType::UINT_32 => Some(parquet::ConvertedType::Uint32),
- ConvertedType::UINT_64 => Some(parquet::ConvertedType::Uint64),
- ConvertedType::INT_8 => Some(parquet::ConvertedType::Int8),
- ConvertedType::INT_16 => Some(parquet::ConvertedType::Int16),
- ConvertedType::INT_32 => Some(parquet::ConvertedType::Int32),
- ConvertedType::INT_64 => Some(parquet::ConvertedType::Int64),
- ConvertedType::JSON => Some(parquet::ConvertedType::Json),
- ConvertedType::BSON => Some(parquet::ConvertedType::Bson),
- ConvertedType::INTERVAL => Some(parquet::ConvertedType::Interval),
- }
- }
-}
-
-// ----------------------------------------------------------------------
-// parquet::LogicalType <=> LogicalType conversion
-
-impl convert::From<parquet::LogicalType> for LogicalType {
- fn from(value: parquet::LogicalType) -> Self {
- match value {
- parquet::LogicalType::STRING(t) => LogicalType::STRING(t),
- parquet::LogicalType::MAP(t) => LogicalType::MAP(t),
- parquet::LogicalType::LIST(t) => LogicalType::LIST(t),
- parquet::LogicalType::ENUM(t) => LogicalType::ENUM(t),
- parquet::LogicalType::DECIMAL(t) => LogicalType::DECIMAL(t),
- parquet::LogicalType::DATE(t) => LogicalType::DATE(t),
- parquet::LogicalType::TIME(t) => LogicalType::TIME(t),
- parquet::LogicalType::TIMESTAMP(t) => LogicalType::TIMESTAMP(t),
- parquet::LogicalType::INTEGER(t) => LogicalType::INTEGER(t),
- parquet::LogicalType::UNKNOWN(t) => LogicalType::UNKNOWN(t),
- parquet::LogicalType::JSON(t) => LogicalType::JSON(t),
- parquet::LogicalType::BSON(t) => LogicalType::BSON(t),
- parquet::LogicalType::UUID(t) => LogicalType::UUID(t),
- }
- }
-}
-
-impl convert::From<LogicalType> for parquet::LogicalType {
- fn from(value: LogicalType) -> Self {
- match value {
- LogicalType::STRING(t) => parquet::LogicalType::STRING(t),
- LogicalType::MAP(t) => parquet::LogicalType::MAP(t),
- LogicalType::LIST(t) => parquet::LogicalType::LIST(t),
- LogicalType::ENUM(t) => parquet::LogicalType::ENUM(t),
- LogicalType::DECIMAL(t) => parquet::LogicalType::DECIMAL(t),
- LogicalType::DATE(t) => parquet::LogicalType::DATE(t),
- LogicalType::TIME(t) => parquet::LogicalType::TIME(t),
- LogicalType::TIMESTAMP(t) => parquet::LogicalType::TIMESTAMP(t),
- LogicalType::INTEGER(t) => parquet::LogicalType::INTEGER(t),
- LogicalType::UNKNOWN(t) => parquet::LogicalType::UNKNOWN(t),
- LogicalType::JSON(t) => parquet::LogicalType::JSON(t),
- LogicalType::BSON(t) => parquet::LogicalType::BSON(t),
- LogicalType::UUID(t) => parquet::LogicalType::UUID(t),
- }
- }
-}
-
-// ----------------------------------------------------------------------
-// LogicalType <=> ConvertedType conversion
-
-// Note: To prevent type loss when converting from ConvertedType to LogicalType,
-// the conversion from ConvertedType -> LogicalType is not implemented.
-// Such type loss includes:
-// - Not knowing the decimal scale and precision of ConvertedType
-// - Time and timestamp nanosecond precision, that is not supported in ConvertedType.
-
-impl From<Option<LogicalType>> for ConvertedType {
- fn from(value: Option<LogicalType>) -> Self {
- match value {
- Some(value) => match value {
- LogicalType::STRING(_) => ConvertedType::UTF8,
- LogicalType::MAP(_) => ConvertedType::MAP,
- LogicalType::LIST(_) => ConvertedType::LIST,
- LogicalType::ENUM(_) => ConvertedType::ENUM,
- LogicalType::DECIMAL(_) => ConvertedType::DECIMAL,
- LogicalType::DATE(_) => ConvertedType::DATE,
- LogicalType::TIME(t) => match t.unit {
- TimeUnit::MILLIS(_) => ConvertedType::TIME_MILLIS,
- TimeUnit::MICROS(_) => ConvertedType::TIME_MICROS,
- TimeUnit::NANOS(_) => ConvertedType::NONE,
- },
- LogicalType::TIMESTAMP(t) => match t.unit {
- TimeUnit::MILLIS(_) => ConvertedType::TIMESTAMP_MILLIS,
- TimeUnit::MICROS(_) => ConvertedType::TIMESTAMP_MICROS,
- TimeUnit::NANOS(_) => ConvertedType::NONE,
- },
- LogicalType::INTEGER(t) => match (t.bit_width, t.is_signed) {
- (8, true) => ConvertedType::INT_8,
- (16, true) => ConvertedType::INT_16,
- (32, true) => ConvertedType::INT_32,
- (64, true) => ConvertedType::INT_64,
- (8, false) => ConvertedType::UINT_8,
- (16, false) => ConvertedType::UINT_16,
- (32, false) => ConvertedType::UINT_32,
- (64, false) => ConvertedType::UINT_64,
- t => panic!("Integer type {:?} is not supported", t),
- },
- LogicalType::UNKNOWN(_) => ConvertedType::NONE,
- LogicalType::JSON(_) => ConvertedType::JSON,
- LogicalType::BSON(_) => ConvertedType::BSON,
- LogicalType::UUID(_) => ConvertedType::NONE,
- },
- None => ConvertedType::NONE,
- }
- }
-}
-
-// ----------------------------------------------------------------------
-// parquet::FieldRepetitionType <=> Repetition conversion
-
-impl convert::From<parquet::FieldRepetitionType> for Repetition {
- fn from(value: parquet::FieldRepetitionType) -> Self {
- match value {
- parquet::FieldRepetitionType::Required => Repetition::REQUIRED,
- parquet::FieldRepetitionType::Optional => Repetition::OPTIONAL,
- parquet::FieldRepetitionType::Repeated => Repetition::REPEATED,
- }
- }
-}
-
-impl convert::From<Repetition> for parquet::FieldRepetitionType {
- fn from(value: Repetition) -> Self {
- match value {
- Repetition::REQUIRED => parquet::FieldRepetitionType::Required,
- Repetition::OPTIONAL => parquet::FieldRepetitionType::Optional,
- Repetition::REPEATED => parquet::FieldRepetitionType::Repeated,
- }
- }
-}
-
-// ----------------------------------------------------------------------
-// parquet::Encoding <=> Encoding conversion
-
-impl convert::From<parquet::Encoding> for Encoding {
- fn from(value: parquet::Encoding) -> Self {
- match value {
- parquet::Encoding::Plain => Encoding::PLAIN,
- parquet::Encoding::PlainDictionary => Encoding::PLAIN_DICTIONARY,
- parquet::Encoding::Rle => Encoding::RLE,
- parquet::Encoding::BitPacked => Encoding::BIT_PACKED,
- parquet::Encoding::DeltaBinaryPacked => Encoding::DELTA_BINARY_PACKED,
- parquet::Encoding::DeltaLengthByteArray => Encoding::DELTA_LENGTH_BYTE_ARRAY,
- parquet::Encoding::DeltaByteArray => Encoding::DELTA_BYTE_ARRAY,
- parquet::Encoding::RleDictionary => Encoding::RLE_DICTIONARY,
- }
- }
-}
-
-impl convert::From<Encoding> for parquet::Encoding {
- fn from(value: Encoding) -> Self {
- match value {
- Encoding::PLAIN => parquet::Encoding::Plain,
- Encoding::PLAIN_DICTIONARY => parquet::Encoding::PlainDictionary,
- Encoding::RLE => parquet::Encoding::Rle,
- Encoding::BIT_PACKED => parquet::Encoding::BitPacked,
- Encoding::DELTA_BINARY_PACKED => parquet::Encoding::DeltaBinaryPacked,
- Encoding::DELTA_LENGTH_BYTE_ARRAY => parquet::Encoding::DeltaLengthByteArray,
- Encoding::DELTA_BYTE_ARRAY => parquet::Encoding::DeltaByteArray,
- Encoding::RLE_DICTIONARY => parquet::Encoding::RleDictionary,
- }
- }
-}
-
-// ----------------------------------------------------------------------
-// parquet::CompressionCodec <=> Compression conversion
-
-impl convert::From<parquet::CompressionCodec> for Compression {
- fn from(value: parquet::CompressionCodec) -> Self {
- match value {
- parquet::CompressionCodec::Uncompressed => Compression::UNCOMPRESSED,
- parquet::CompressionCodec::Snappy => Compression::SNAPPY,
- parquet::CompressionCodec::Gzip => Compression::GZIP,
- parquet::CompressionCodec::Lzo => Compression::LZO,
- parquet::CompressionCodec::Brotli => Compression::BROTLI,
- parquet::CompressionCodec::Lz4 => Compression::LZ4,
- parquet::CompressionCodec::Zstd => Compression::ZSTD,
- }
- }
-}
-
-impl convert::From<Compression> for parquet::CompressionCodec {
- fn from(value: Compression) -> Self {
- match value {
- Compression::UNCOMPRESSED => parquet::CompressionCodec::Uncompressed,
- Compression::SNAPPY => parquet::CompressionCodec::Snappy,
- Compression::GZIP => parquet::CompressionCodec::Gzip,
- Compression::LZO => parquet::CompressionCodec::Lzo,
- Compression::BROTLI => parquet::CompressionCodec::Brotli,
- Compression::LZ4 => parquet::CompressionCodec::Lz4,
- Compression::ZSTD => parquet::CompressionCodec::Zstd,
- }
- }
-}
-
-// ----------------------------------------------------------------------
-// parquet::PageType <=> PageType conversion
-
-impl convert::From<parquet::PageType> for PageType {
- fn from(value: parquet::PageType) -> Self {
- match value {
- parquet::PageType::DataPage => PageType::DATA_PAGE,
- parquet::PageType::IndexPage => PageType::INDEX_PAGE,
- parquet::PageType::DictionaryPage => PageType::DICTIONARY_PAGE,
- parquet::PageType::DataPageV2 => PageType::DATA_PAGE_V2,
- }
- }
-}
-
-impl convert::From<PageType> for parquet::PageType {
- fn from(value: PageType) -> Self {
- match value {
- PageType::DATA_PAGE => parquet::PageType::DataPage,
- PageType::INDEX_PAGE => parquet::PageType::IndexPage,
- PageType::DICTIONARY_PAGE => parquet::PageType::DictionaryPage,
- PageType::DATA_PAGE_V2 => parquet::PageType::DataPageV2,
- }
- }
-}
-
-// ----------------------------------------------------------------------
-// String conversions for schema parsing.
-
-impl str::FromStr for Repetition {
- type Err = ParquetError;
-
- fn from_str(s: &str) -> result::Result<Self, Self::Err> {
- match s {
- "REQUIRED" => Ok(Repetition::REQUIRED),
- "OPTIONAL" => Ok(Repetition::OPTIONAL),
- "REPEATED" => Ok(Repetition::REPEATED),
- other => Err(general_err!("Invalid repetition {}", other)),
- }
- }
-}
-
-impl str::FromStr for Type {
- type Err = ParquetError;
-
- fn from_str(s: &str) -> result::Result<Self, Self::Err> {
- match s {
- "BOOLEAN" => Ok(Type::BOOLEAN),
- "INT32" => Ok(Type::INT32),
- "INT64" => Ok(Type::INT64),
- "INT96" => Ok(Type::INT96),
- "FLOAT" => Ok(Type::FLOAT),
- "DOUBLE" => Ok(Type::DOUBLE),
- "BYTE_ARRAY" | "BINARY" => Ok(Type::BYTE_ARRAY),
- "FIXED_LEN_BYTE_ARRAY" => Ok(Type::FIXED_LEN_BYTE_ARRAY),
- other => Err(general_err!("Invalid type {}", other)),
- }
- }
-}
-
-impl str::FromStr for ConvertedType {
- type Err = ParquetError;
-
- fn from_str(s: &str) -> result::Result<Self, Self::Err> {
- match s {
- "NONE" => Ok(ConvertedType::NONE),
- "UTF8" => Ok(ConvertedType::UTF8),
- "MAP" => Ok(ConvertedType::MAP),
- "MAP_KEY_VALUE" => Ok(ConvertedType::MAP_KEY_VALUE),
- "LIST" => Ok(ConvertedType::LIST),
- "ENUM" => Ok(ConvertedType::ENUM),
- "DECIMAL" => Ok(ConvertedType::DECIMAL),
- "DATE" => Ok(ConvertedType::DATE),
- "TIME_MILLIS" => Ok(ConvertedType::TIME_MILLIS),
- "TIME_MICROS" => Ok(ConvertedType::TIME_MICROS),
- "TIMESTAMP_MILLIS" => Ok(ConvertedType::TIMESTAMP_MILLIS),
- "TIMESTAMP_MICROS" => Ok(ConvertedType::TIMESTAMP_MICROS),
- "UINT_8" => Ok(ConvertedType::UINT_8),
- "UINT_16" => Ok(ConvertedType::UINT_16),
- "UINT_32" => Ok(ConvertedType::UINT_32),
- "UINT_64" => Ok(ConvertedType::UINT_64),
- "INT_8" => Ok(ConvertedType::INT_8),
- "INT_16" => Ok(ConvertedType::INT_16),
- "INT_32" => Ok(ConvertedType::INT_32),
- "INT_64" => Ok(ConvertedType::INT_64),
- "JSON" => Ok(ConvertedType::JSON),
- "BSON" => Ok(ConvertedType::BSON),
- "INTERVAL" => Ok(ConvertedType::INTERVAL),
- other => Err(general_err!("Invalid converted type {}", other)),
- }
- }
-}
-
-impl str::FromStr for LogicalType {
- type Err = ParquetError;
-
- fn from_str(s: &str) -> result::Result<Self, Self::Err> {
- match s {
- // The type is a placeholder that gets updated elsewhere
- "INTEGER" => Ok(LogicalType::INTEGER(IntType {
- bit_width: 8,
- is_signed: false,
- })),
- "MAP" => Ok(LogicalType::MAP(MapType {})),
- "LIST" => Ok(LogicalType::LIST(ListType {})),
- "ENUM" => Ok(LogicalType::ENUM(EnumType {})),
- "DECIMAL" => Ok(LogicalType::DECIMAL(DecimalType {
- precision: -1,
- scale: -1,
- })),
- "DATE" => Ok(LogicalType::DATE(DateType {})),
- "TIME" => Ok(LogicalType::TIME(TimeType {
- is_adjusted_to_u_t_c: false,
- unit: TimeUnit::MILLIS(parquet::MilliSeconds {}),
- })),
- "TIMESTAMP" => Ok(LogicalType::TIMESTAMP(TimestampType {
- is_adjusted_to_u_t_c: false,
- unit: TimeUnit::MILLIS(parquet::MilliSeconds {}),
- })),
- "STRING" => Ok(LogicalType::STRING(StringType {})),
- "JSON" => Ok(LogicalType::JSON(JsonType {})),
- "BSON" => Ok(LogicalType::BSON(BsonType {})),
- "UUID" => Ok(LogicalType::UUID(UUIDType {})),
- "UNKNOWN" => Ok(LogicalType::UNKNOWN(NullType {})),
- "INTERVAL" => Err(general_err!("Interval logical type not yet supported")),
- other => Err(general_err!("Invalid logical type {}", other)),
- }
- }
-}
-
-#[cfg(test)]
-mod tests {
- use super::*;
-
- #[test]
- fn test_display_type() {
- assert_eq!(Type::BOOLEAN.to_string(), "BOOLEAN");
- assert_eq!(Type::INT32.to_string(), "INT32");
- assert_eq!(Type::INT64.to_string(), "INT64");
- assert_eq!(Type::INT96.to_string(), "INT96");
- assert_eq!(Type::FLOAT.to_string(), "FLOAT");
- assert_eq!(Type::DOUBLE.to_string(), "DOUBLE");
- assert_eq!(Type::BYTE_ARRAY.to_string(), "BYTE_ARRAY");
- assert_eq!(
- Type::FIXED_LEN_BYTE_ARRAY.to_string(),
- "FIXED_LEN_BYTE_ARRAY"
- );
- }
-
- #[test]
- fn test_from_type() {
- assert_eq!(Type::from(parquet::Type::Boolean), Type::BOOLEAN);
- assert_eq!(Type::from(parquet::Type::Int32), Type::INT32);
- assert_eq!(Type::from(parquet::Type::Int64), Type::INT64);
- assert_eq!(Type::from(parquet::Type::Int96), Type::INT96);
- assert_eq!(Type::from(parquet::Type::Float), Type::FLOAT);
- assert_eq!(Type::from(parquet::Type::Double), Type::DOUBLE);
- assert_eq!(Type::from(parquet::Type::ByteArray), Type::BYTE_ARRAY);
- assert_eq!(
- Type::from(parquet::Type::FixedLenByteArray),
- Type::FIXED_LEN_BYTE_ARRAY
- );
- }
-
- #[test]
- fn test_into_type() {
- assert_eq!(parquet::Type::Boolean, Type::BOOLEAN.into());
- assert_eq!(parquet::Type::Int32, Type::INT32.into());
- assert_eq!(parquet::Type::Int64, Type::INT64.into());
- assert_eq!(parquet::Type::Int96, Type::INT96.into());
- assert_eq!(parquet::Type::Float, Type::FLOAT.into());
- assert_eq!(parquet::Type::Double, Type::DOUBLE.into());
- assert_eq!(parquet::Type::ByteArray, Type::BYTE_ARRAY.into());
- assert_eq!(
- parquet::Type::FixedLenByteArray,
- Type::FIXED_LEN_BYTE_ARRAY.into()
- );
- }
-
- #[test]
- fn test_from_string_into_type() {
- assert_eq!(
- Type::BOOLEAN.to_string().parse::<Type>().unwrap(),
- Type::BOOLEAN
- );
- assert_eq!(
- Type::INT32.to_string().parse::<Type>().unwrap(),
- Type::INT32
- );
- assert_eq!(
- Type::INT64.to_string().parse::<Type>().unwrap(),
- Type::INT64
- );
- assert_eq!(
- Type::INT96.to_string().parse::<Type>().unwrap(),
- Type::INT96
- );
- assert_eq!(
- Type::FLOAT.to_string().parse::<Type>().unwrap(),
- Type::FLOAT
- );
- assert_eq!(
- Type::DOUBLE.to_string().parse::<Type>().unwrap(),
- Type::DOUBLE
- );
- assert_eq!(
- Type::BYTE_ARRAY.to_string().parse::<Type>().unwrap(),
- Type::BYTE_ARRAY
- );
- assert_eq!("BINARY".parse::<Type>().unwrap(), Type::BYTE_ARRAY);
- assert_eq!(
- Type::FIXED_LEN_BYTE_ARRAY
- .to_string()
- .parse::<Type>()
- .unwrap(),
- Type::FIXED_LEN_BYTE_ARRAY
- );
- }
-
- #[test]
- fn test_display_converted_type() {
- assert_eq!(ConvertedType::NONE.to_string(), "NONE");
- assert_eq!(ConvertedType::UTF8.to_string(), "UTF8");
- assert_eq!(ConvertedType::MAP.to_string(), "MAP");
- assert_eq!(ConvertedType::MAP_KEY_VALUE.to_string(), "MAP_KEY_VALUE");
- assert_eq!(ConvertedType::LIST.to_string(), "LIST");
- assert_eq!(ConvertedType::ENUM.to_string(), "ENUM");
- assert_eq!(ConvertedType::DECIMAL.to_string(), "DECIMAL");
- assert_eq!(ConvertedType::DATE.to_string(), "DATE");
- assert_eq!(ConvertedType::TIME_MILLIS.to_string(), "TIME_MILLIS");
- assert_eq!(ConvertedType::DATE.to_string(), "DATE");
- assert_eq!(ConvertedType::TIME_MICROS.to_string(), "TIME_MICROS");
- assert_eq!(
- ConvertedType::TIMESTAMP_MILLIS.to_string(),
- "TIMESTAMP_MILLIS"
- );
- assert_eq!(
- ConvertedType::TIMESTAMP_MICROS.to_string(),
- "TIMESTAMP_MICROS"
- );
- assert_eq!(ConvertedType::UINT_8.to_string(), "UINT_8");
- assert_eq!(ConvertedType::UINT_16.to_string(), "UINT_16");
- assert_eq!(ConvertedType::UINT_32.to_string(), "UINT_32");
- assert_eq!(ConvertedType::UINT_64.to_string(), "UINT_64");
- assert_eq!(ConvertedType::INT_8.to_string(), "INT_8");
- assert_eq!(ConvertedType::INT_16.to_string(), "INT_16");
- assert_eq!(ConvertedType::INT_32.to_string(), "INT_32");
- assert_eq!(ConvertedType::INT_64.to_string(), "INT_64");
- assert_eq!(ConvertedType::JSON.to_string(), "JSON");
- assert_eq!(ConvertedType::BSON.to_string(), "BSON");
- assert_eq!(ConvertedType::INTERVAL.to_string(), "INTERVAL");
- }
-
- #[test]
- fn test_from_converted_type() {
- let parquet_conv_none: Option<parquet::ConvertedType> = None;
- assert_eq!(ConvertedType::from(parquet_conv_none), ConvertedType::NONE);
- assert_eq!(
- ConvertedType::from(Some(parquet::ConvertedType::Utf8)),
- ConvertedType::UTF8
- );
- assert_eq!(
- ConvertedType::from(Some(parquet::ConvertedType::Map)),
- ConvertedType::MAP
- );
- assert_eq!(
- ConvertedType::from(Some(parquet::ConvertedType::MapKeyValue)),
- ConvertedType::MAP_KEY_VALUE
- );
- assert_eq!(
- ConvertedType::from(Some(parquet::ConvertedType::List)),
- ConvertedType::LIST
- );
- assert_eq!(
- ConvertedType::from(Some(parquet::ConvertedType::Enum)),
- ConvertedType::ENUM
- );
- assert_eq!(
- ConvertedType::from(Some(parquet::ConvertedType::Decimal)),
- ConvertedType::DECIMAL
- );
- assert_eq!(
- ConvertedType::from(Some(parquet::ConvertedType::Date)),
- ConvertedType::DATE
- );
- assert_eq!(
- ConvertedType::from(Some(parquet::ConvertedType::TimeMillis)),
- ConvertedType::TIME_MILLIS
- );
- assert_eq!(
- ConvertedType::from(Some(parquet::ConvertedType::TimeMicros)),
- ConvertedType::TIME_MICROS
- );
- assert_eq!(
- ConvertedType::from(Some(parquet::ConvertedType::TimestampMillis)),
- ConvertedType::TIMESTAMP_MILLIS
- );
- assert_eq!(
- ConvertedType::from(Some(parquet::ConvertedType::TimestampMicros)),
- ConvertedType::TIMESTAMP_MICROS
- );
- assert_eq!(
- ConvertedType::from(Some(parquet::ConvertedType::Uint8)),
- ConvertedType::UINT_8
- );
- assert_eq!(
- ConvertedType::from(Some(parquet::ConvertedType::Uint16)),
- ConvertedType::UINT_16
- );
- assert_eq!(
- ConvertedType::from(Some(parquet::ConvertedType::Uint32)),
- ConvertedType::UINT_32
- );
- assert_eq!(
- ConvertedType::from(Some(parquet::ConvertedType::Uint64)),
- ConvertedType::UINT_64
- );
- assert_eq!(
- ConvertedType::from(Some(parquet::ConvertedType::Int8)),
- ConvertedType::INT_8
- );
- assert_eq!(
- ConvertedType::from(Some(parquet::ConvertedType::Int16)),
- ConvertedType::INT_16
- );
- assert_eq!(
- ConvertedType::from(Some(parquet::ConvertedType::Int32)),
- ConvertedType::INT_32
- );
- assert_eq!(
- ConvertedType::from(Some(parquet::ConvertedType::Int64)),
- ConvertedType::INT_64
- );
- assert_eq!(
- ConvertedType::from(Some(parquet::ConvertedType::Json)),
- ConvertedType::JSON
- );
- assert_eq!(
- ConvertedType::from(Some(parquet::ConvertedType::Bson)),
- ConvertedType::BSON
- );
- assert_eq!(
- ConvertedType::from(Some(parquet::ConvertedType::Interval)),
- ConvertedType::INTERVAL
- );
- }
-
- #[test]
- fn test_into_converted_type() {
- let converted_type: Option<parquet::ConvertedType> = None;
- assert_eq!(converted_type, ConvertedType::NONE.into());
- assert_eq!(
- Some(parquet::ConvertedType::Utf8),
- ConvertedType::UTF8.into()
- );
- assert_eq!(Some(parquet::ConvertedType::Map), ConvertedType::MAP.into());
- assert_eq!(
- Some(parquet::ConvertedType::MapKeyValue),
- ConvertedType::MAP_KEY_VALUE.into()
- );
- assert_eq!(
- Some(parquet::ConvertedType::List),
- ConvertedType::LIST.into()
- );
- assert_eq!(
- Some(parquet::ConvertedType::Enum),
- ConvertedType::ENUM.into()
- );
- assert_eq!(
- Some(parquet::ConvertedType::Decimal),
- ConvertedType::DECIMAL.into()
- );
- assert_eq!(
- Some(parquet::ConvertedType::Date),
- ConvertedType::DATE.into()
- );
- assert_eq!(
- Some(parquet::ConvertedType::TimeMillis),
- ConvertedType::TIME_MILLIS.into()
- );
- assert_eq!(
- Some(parquet::ConvertedType::TimeMicros),
- ConvertedType::TIME_MICROS.into()
- );
- assert_eq!(
- Some(parquet::ConvertedType::TimestampMillis),
- ConvertedType::TIMESTAMP_MILLIS.into()
- );
- assert_eq!(
- Some(parquet::ConvertedType::TimestampMicros),
- ConvertedType::TIMESTAMP_MICROS.into()
- );
- assert_eq!(
- Some(parquet::ConvertedType::Uint8),
- ConvertedType::UINT_8.into()
- );
- assert_eq!(
- Some(parquet::ConvertedType::Uint16),
- ConvertedType::UINT_16.into()
- );
- assert_eq!(
- Some(parquet::ConvertedType::Uint32),
- ConvertedType::UINT_32.into()
- );
- assert_eq!(
- Some(parquet::ConvertedType::Uint64),
- ConvertedType::UINT_64.into()
- );
- assert_eq!(
- Some(parquet::ConvertedType::Int8),
- ConvertedType::INT_8.into()
- );
- assert_eq!(
- Some(parquet::ConvertedType::Int16),
- ConvertedType::INT_16.into()
- );
- assert_eq!(
- Some(parquet::ConvertedType::Int32),
- ConvertedType::INT_32.into()
- );
- assert_eq!(
- Some(parquet::ConvertedType::Int64),
- ConvertedType::INT_64.into()
- );
- assert_eq!(
- Some(parquet::ConvertedType::Json),
- ConvertedType::JSON.into()
- );
- assert_eq!(
- Some(parquet::ConvertedType::Bson),
- ConvertedType::BSON.into()
- );
- assert_eq!(
- Some(parquet::ConvertedType::Interval),
- ConvertedType::INTERVAL.into()
- );
- }
-
- #[test]
- fn test_from_string_into_converted_type() {
- assert_eq!(
- ConvertedType::NONE
- .to_string()
- .parse::<ConvertedType>()
- .unwrap(),
- ConvertedType::NONE
- );
- assert_eq!(
- ConvertedType::UTF8
- .to_string()
- .parse::<ConvertedType>()
- .unwrap(),
- ConvertedType::UTF8
- );
- assert_eq!(
- ConvertedType::MAP
- .to_string()
- .parse::<ConvertedType>()
- .unwrap(),
- ConvertedType::MAP
- );
- assert_eq!(
- ConvertedType::MAP_KEY_VALUE
- .to_string()
- .parse::<ConvertedType>()
- .unwrap(),
- ConvertedType::MAP_KEY_VALUE
- );
- assert_eq!(
- ConvertedType::LIST
- .to_string()
- .parse::<ConvertedType>()
- .unwrap(),
- ConvertedType::LIST
- );
- assert_eq!(
- ConvertedType::ENUM
- .to_string()
- .parse::<ConvertedType>()
- .unwrap(),
- ConvertedType::ENUM
- );
- assert_eq!(
- ConvertedType::DECIMAL
- .to_string()
- .parse::<ConvertedType>()
- .unwrap(),
- ConvertedType::DECIMAL
- );
- assert_eq!(
- ConvertedType::DATE
- .to_string()
- .parse::<ConvertedType>()
- .unwrap(),
- ConvertedType::DATE
- );
- assert_eq!(
- ConvertedType::TIME_MILLIS
- .to_string()
- .parse::<ConvertedType>()
- .unwrap(),
- ConvertedType::TIME_MILLIS
- );
- assert_eq!(
- ConvertedType::TIME_MICROS
- .to_string()
- .parse::<ConvertedType>()
- .unwrap(),
- ConvertedType::TIME_MICROS
- );
- assert_eq!(
- ConvertedType::TIMESTAMP_MILLIS
- .to_string()
- .parse::<ConvertedType>()
- .unwrap(),
- ConvertedType::TIMESTAMP_MILLIS
- );
- assert_eq!(
- ConvertedType::TIMESTAMP_MICROS
- .to_string()
- .parse::<ConvertedType>()
- .unwrap(),
- ConvertedType::TIMESTAMP_MICROS
- );
- assert_eq!(
- ConvertedType::UINT_8
- .to_string()
- .parse::<ConvertedType>()
- .unwrap(),
- ConvertedType::UINT_8
- );
- assert_eq!(
- ConvertedType::UINT_16
- .to_string()
- .parse::<ConvertedType>()
- .unwrap(),
- ConvertedType::UINT_16
- );
- assert_eq!(
- ConvertedType::UINT_32
- .to_string()
- .parse::<ConvertedType>()
- .unwrap(),
- ConvertedType::UINT_32
- );
- assert_eq!(
- ConvertedType::UINT_64
- .to_string()
- .parse::<ConvertedType>()
- .unwrap(),
- ConvertedType::UINT_64
- );
- assert_eq!(
- ConvertedType::INT_8
- .to_string()
- .parse::<ConvertedType>()
- .unwrap(),
- ConvertedType::INT_8
- );
- assert_eq!(
- ConvertedType::INT_16
- .to_string()
- .parse::<ConvertedType>()
- .unwrap(),
- ConvertedType::INT_16
- );
- assert_eq!(
- ConvertedType::INT_32
- .to_string()
- .parse::<ConvertedType>()
- .unwrap(),
- ConvertedType::INT_32
- );
- assert_eq!(
- ConvertedType::INT_64
- .to_string()
- .parse::<ConvertedType>()
- .unwrap(),
- ConvertedType::INT_64
- );
- assert_eq!(
- ConvertedType::JSON
- .to_string()
- .parse::<ConvertedType>()
- .unwrap(),
- ConvertedType::JSON
- );
- assert_eq!(
- ConvertedType::BSON
- .to_string()
- .parse::<ConvertedType>()
- .unwrap(),
- ConvertedType::BSON
- );
- assert_eq!(
- ConvertedType::INTERVAL
- .to_string()
- .parse::<ConvertedType>()
- .unwrap(),
- ConvertedType::INTERVAL
- );
- }
-
- #[test]
- fn test_logical_to_converted_type() {
- let logical_none: Option<LogicalType> = None;
- assert_eq!(ConvertedType::from(logical_none), ConvertedType::NONE);
- assert_eq!(
- ConvertedType::from(Some(LogicalType::DECIMAL(DecimalType {
- precision: 20,
- scale: 5
- }))),
- ConvertedType::DECIMAL
- );
- assert_eq!(
- ConvertedType::from(Some(LogicalType::BSON(Default::default()))),
- ConvertedType::BSON
- );
- assert_eq!(
- ConvertedType::from(Some(LogicalType::JSON(Default::default()))),
- ConvertedType::JSON
- );
- assert_eq!(
- ConvertedType::from(Some(LogicalType::STRING(Default::default()))),
- ConvertedType::UTF8
- );
- assert_eq!(
- ConvertedType::from(Some(LogicalType::DATE(Default::default()))),
- ConvertedType::DATE
- );
- assert_eq!(
- ConvertedType::from(Some(LogicalType::TIME(TimeType {
- unit: TimeUnit::MILLIS(Default::default()),
- is_adjusted_to_u_t_c: true,
- }))),
- ConvertedType::TIME_MILLIS
- );
- assert_eq!(
- ConvertedType::from(Some(LogicalType::TIME(TimeType {
- unit: TimeUnit::MICROS(Default::default()),
- is_adjusted_to_u_t_c: true,
- }))),
- ConvertedType::TIME_MICROS
- );
- assert_eq!(
- ConvertedType::from(Some(LogicalType::TIME(TimeType {
- unit: TimeUnit::NANOS(Default::default()),
- is_adjusted_to_u_t_c: false,
- }))),
- ConvertedType::NONE
- );
- assert_eq!(
- ConvertedType::from(Some(LogicalType::TIMESTAMP(TimestampType {
- unit: TimeUnit::MILLIS(Default::default()),
- is_adjusted_to_u_t_c: true,
- }))),
- ConvertedType::TIMESTAMP_MILLIS
- );
- assert_eq!(
- ConvertedType::from(Some(LogicalType::TIMESTAMP(TimestampType {
- unit: TimeUnit::MICROS(Default::default()),
- is_adjusted_to_u_t_c: false,
- }))),
- ConvertedType::TIMESTAMP_MICROS
- );
- assert_eq!(
- ConvertedType::from(Some(LogicalType::TIMESTAMP(TimestampType {
- unit: TimeUnit::NANOS(Default::default()),
- is_adjusted_to_u_t_c: false,
- }))),
- ConvertedType::NONE
- );
- assert_eq!(
- ConvertedType::from(Some(LogicalType::INTEGER(IntType {
- bit_width: 8,
- is_signed: false
- }))),
- ConvertedType::UINT_8
- );
- assert_eq!(
- ConvertedType::from(Some(LogicalType::INTEGER(IntType {
- bit_width: 8,
- is_signed: true
- }))),
- ConvertedType::INT_8
- );
- assert_eq!(
- ConvertedType::from(Some(LogicalType::INTEGER(IntType {
- bit_width: 16,
- is_signed: false
- }))),
- ConvertedType::UINT_16
- );
- assert_eq!(
- ConvertedType::from(Some(LogicalType::INTEGER(IntType {
- bit_width: 16,
- is_signed: true
- }))),
- ConvertedType::INT_16
- );
- assert_eq!(
- ConvertedType::from(Some(LogicalType::INTEGER(IntType {
- bit_width: 32,
- is_signed: false
- }))),
- ConvertedType::UINT_32
- );
- assert_eq!(
- ConvertedType::from(Some(LogicalType::INTEGER(IntType {
- bit_width: 32,
- is_signed: true
- }))),
- ConvertedType::INT_32
- );
- assert_eq!(
- ConvertedType::from(Some(LogicalType::INTEGER(IntType {
- bit_width: 64,
- is_signed: false
- }))),
- ConvertedType::UINT_64
- );
- assert_eq!(
- ConvertedType::from(Some(LogicalType::INTEGER(IntType {
- bit_width: 64,
- is_signed: true
- }))),
- ConvertedType::INT_64
- );
- assert_eq!(
- ConvertedType::from(Some(LogicalType::LIST(Default::default()))),
- ConvertedType::LIST
- );
- assert_eq!(
- ConvertedType::from(Some(LogicalType::MAP(Default::default()))),
- ConvertedType::MAP
- );
- assert_eq!(
- ConvertedType::from(Some(LogicalType::UUID(Default::default()))),
- ConvertedType::NONE
- );
- assert_eq!(
- ConvertedType::from(Some(LogicalType::ENUM(Default::default()))),
- ConvertedType::ENUM
- );
- assert_eq!(
- ConvertedType::from(Some(LogicalType::UNKNOWN(Default::default()))),
- ConvertedType::NONE
- );
- }
-
- #[test]
- fn test_display_repetition() {
- assert_eq!(Repetition::REQUIRED.to_string(), "REQUIRED");
- assert_eq!(Repetition::OPTIONAL.to_string(), "OPTIONAL");
- assert_eq!(Repetition::REPEATED.to_string(), "REPEATED");
- }
-
- #[test]
- fn test_from_repetition() {
- assert_eq!(
- Repetition::from(parquet::FieldRepetitionType::Required),
- Repetition::REQUIRED
- );
- assert_eq!(
- Repetition::from(parquet::FieldRepetitionType::Optional),
- Repetition::OPTIONAL
- );
- assert_eq!(
- Repetition::from(parquet::FieldRepetitionType::Repeated),
- Repetition::REPEATED
- );
- }
-
- #[test]
- fn test_into_repetition() {
- assert_eq!(
- parquet::FieldRepetitionType::Required,
- Repetition::REQUIRED.into()
- );
- assert_eq!(
- parquet::FieldRepetitionType::Optional,
- Repetition::OPTIONAL.into()
- );
- assert_eq!(
- parquet::FieldRepetitionType::Repeated,
- Repetition::REPEATED.into()
- );
- }
-
- #[test]
- fn test_from_string_into_repetition() {
- assert_eq!(
- Repetition::REQUIRED
- .to_string()
- .parse::<Repetition>()
- .unwrap(),
- Repetition::REQUIRED
- );
- assert_eq!(
- Repetition::OPTIONAL
- .to_string()
- .parse::<Repetition>()
- .unwrap(),
- Repetition::OPTIONAL
- );
- assert_eq!(
- Repetition::REPEATED
- .to_string()
- .parse::<Repetition>()
- .unwrap(),
- Repetition::REPEATED
- );
- }
-
- #[test]
- fn test_display_encoding() {
- assert_eq!(Encoding::PLAIN.to_string(), "PLAIN");
- assert_eq!(Encoding::PLAIN_DICTIONARY.to_string(), "PLAIN_DICTIONARY");
- assert_eq!(Encoding::RLE.to_string(), "RLE");
- assert_eq!(Encoding::BIT_PACKED.to_string(), "BIT_PACKED");
- assert_eq!(
- Encoding::DELTA_BINARY_PACKED.to_string(),
- "DELTA_BINARY_PACKED"
- );
- assert_eq!(
- Encoding::DELTA_LENGTH_BYTE_ARRAY.to_string(),
- "DELTA_LENGTH_BYTE_ARRAY"
- );
- assert_eq!(Encoding::DELTA_BYTE_ARRAY.to_string(), "DELTA_BYTE_ARRAY");
- assert_eq!(Encoding::RLE_DICTIONARY.to_string(), "RLE_DICTIONARY");
- }
-
- #[test]
- fn test_from_encoding() {
- assert_eq!(Encoding::from(parquet::Encoding::Plain), Encoding::PLAIN);
- assert_eq!(
- Encoding::from(parquet::Encoding::PlainDictionary),
- Encoding::PLAIN_DICTIONARY
- );
- assert_eq!(Encoding::from(parquet::Encoding::Rle), Encoding::RLE);
- assert_eq!(
- Encoding::from(parquet::Encoding::BitPacked),
- Encoding::BIT_PACKED
- );
- assert_eq!(
- Encoding::from(parquet::Encoding::DeltaBinaryPacked),
- Encoding::DELTA_BINARY_PACKED
- );
- assert_eq!(
- Encoding::from(parquet::Encoding::DeltaLengthByteArray),
- Encoding::DELTA_LENGTH_BYTE_ARRAY
- );
- assert_eq!(
- Encoding::from(parquet::Encoding::DeltaByteArray),
- Encoding::DELTA_BYTE_ARRAY
- );
- }
-
- #[test]
- fn test_into_encoding() {
- assert_eq!(parquet::Encoding::Plain, Encoding::PLAIN.into());
- assert_eq!(
- parquet::Encoding::PlainDictionary,
- Encoding::PLAIN_DICTIONARY.into()
- );
- assert_eq!(parquet::Encoding::Rle, Encoding::RLE.into());
- assert_eq!(parquet::Encoding::BitPacked, Encoding::BIT_PACKED.into());
- assert_eq!(
- parquet::Encoding::DeltaBinaryPacked,
- Encoding::DELTA_BINARY_PACKED.into()
- );
- assert_eq!(
- parquet::Encoding::DeltaLengthByteArray,
- Encoding::DELTA_LENGTH_BYTE_ARRAY.into()
- );
- assert_eq!(
- parquet::Encoding::DeltaByteArray,
- Encoding::DELTA_BYTE_ARRAY.into()
- );
- }
-
- #[test]
- fn test_display_compression() {
- assert_eq!(Compression::UNCOMPRESSED.to_string(), "UNCOMPRESSED");
- assert_eq!(Compression::SNAPPY.to_string(), "SNAPPY");
- assert_eq!(Compression::GZIP.to_string(), "GZIP");
- assert_eq!(Compression::LZO.to_string(), "LZO");
- assert_eq!(Compression::BROTLI.to_string(), "BROTLI");
- assert_eq!(Compression::LZ4.to_string(), "LZ4");
- assert_eq!(Compression::ZSTD.to_string(), "ZSTD");
- }
-
- #[test]
- fn test_from_compression() {
- assert_eq!(
- Compression::from(parquet::CompressionCodec::Uncompressed),
- Compression::UNCOMPRESSED
- );
- assert_eq!(
- Compression::from(parquet::CompressionCodec::Snappy),
- Compression::SNAPPY
- );
- assert_eq!(
- Compression::from(parquet::CompressionCodec::Gzip),
- Compression::GZIP
- );
- assert_eq!(
- Compression::from(parquet::CompressionCodec::Lzo),
- Compression::LZO
- );
- assert_eq!(
- Compression::from(parquet::CompressionCodec::Brotli),
- Compression::BROTLI
- );
- assert_eq!(
- Compression::from(parquet::CompressionCodec::Lz4),
- Compression::LZ4
- );
- assert_eq!(
- Compression::from(parquet::CompressionCodec::Zstd),
- Compression::ZSTD
- );
- }
-
- #[test]
- fn test_into_compression() {
- assert_eq!(
- parquet::CompressionCodec::Uncompressed,
- Compression::UNCOMPRESSED.into()
- );
- assert_eq!(
- parquet::CompressionCodec::Snappy,
- Compression::SNAPPY.into()
- );
- assert_eq!(parquet::CompressionCodec::Gzip, Compression::GZIP.into());
- assert_eq!(parquet::CompressionCodec::Lzo, Compression::LZO.into());
- assert_eq!(
- parquet::CompressionCodec::Brotli,
- Compression::BROTLI.into()
- );
- assert_eq!(parquet::CompressionCodec::Lz4, Compression::LZ4.into());
- assert_eq!(parquet::CompressionCodec::Zstd, Compression::ZSTD.into());
- }
-
- #[test]
- fn test_display_page_type() {
- assert_eq!(PageType::DATA_PAGE.to_string(), "DATA_PAGE");
- assert_eq!(PageType::INDEX_PAGE.to_string(), "INDEX_PAGE");
- assert_eq!(PageType::DICTIONARY_PAGE.to_string(), "DICTIONARY_PAGE");
- assert_eq!(PageType::DATA_PAGE_V2.to_string(), "DATA_PAGE_V2");
- }
-
- #[test]
- fn test_from_page_type() {
- assert_eq!(
- PageType::from(parquet::PageType::DataPage),
- PageType::DATA_PAGE
- );
- assert_eq!(
- PageType::from(parquet::PageType::IndexPage),
- PageType::INDEX_PAGE
- );
- assert_eq!(
- PageType::from(parquet::PageType::DictionaryPage),
- PageType::DICTIONARY_PAGE
- );
- assert_eq!(
- PageType::from(parquet::PageType::DataPageV2),
- PageType::DATA_PAGE_V2
- );
- }
-
- #[test]
- fn test_into_page_type() {
- assert_eq!(parquet::PageType::DataPage, PageType::DATA_PAGE.into());
- assert_eq!(parquet::PageType::IndexPage, PageType::INDEX_PAGE.into());
- assert_eq!(
- parquet::PageType::DictionaryPage,
- PageType::DICTIONARY_PAGE.into()
- );
- assert_eq!(parquet::PageType::DataPageV2, PageType::DATA_PAGE_V2.into());
- }
-
- #[test]
- fn test_display_sort_order() {
- assert_eq!(SortOrder::SIGNED.to_string(), "SIGNED");
- assert_eq!(SortOrder::UNSIGNED.to_string(), "UNSIGNED");
- assert_eq!(SortOrder::UNDEFINED.to_string(), "UNDEFINED");
- }
-
- #[test]
- fn test_display_column_order() {
- assert_eq!(
- ColumnOrder::TYPE_DEFINED_ORDER(SortOrder::SIGNED).to_string(),
- "TYPE_DEFINED_ORDER(SIGNED)"
- );
- assert_eq!(
- ColumnOrder::TYPE_DEFINED_ORDER(SortOrder::UNSIGNED).to_string(),
- "TYPE_DEFINED_ORDER(UNSIGNED)"
- );
- assert_eq!(
- ColumnOrder::TYPE_DEFINED_ORDER(SortOrder::UNDEFINED).to_string(),
- "TYPE_DEFINED_ORDER(UNDEFINED)"
- );
- assert_eq!(ColumnOrder::UNDEFINED.to_string(), "UNDEFINED");
- }
-
- #[test]
- fn test_column_order_get_logical_type_sort_order() {
- // Helper to check the order in a list of values.
- // Only logical type is checked.
- fn check_sort_order(types: Vec<LogicalType>, expected_order: SortOrder) {
- for tpe in types {
- assert_eq!(
- ColumnOrder::get_sort_order(
- Some(tpe),
- ConvertedType::NONE,
- Type::BYTE_ARRAY
- ),
- expected_order
- );
- }
- }
-
- // Unsigned comparison (physical type does not matter)
- let unsigned = vec![
- LogicalType::STRING(Default::default()),
- LogicalType::JSON(Default::default()),
- LogicalType::BSON(Default::default()),
- LogicalType::ENUM(Default::default()),
- LogicalType::UUID(Default::default()),
- LogicalType::INTEGER(IntType {
- bit_width: 8,
- is_signed: false,
- }),
- LogicalType::INTEGER(IntType {
- bit_width: 16,
- is_signed: false,
- }),
- LogicalType::INTEGER(IntType {
- bit_width: 32,
- is_signed: false,
- }),
- LogicalType::INTEGER(IntType {
- bit_width: 64,
- is_signed: false,
- }),
- ];
- check_sort_order(unsigned, SortOrder::UNSIGNED);
-
- // Signed comparison (physical type does not matter)
- let signed = vec![
- LogicalType::INTEGER(IntType {
- bit_width: 8,
- is_signed: true,
- }),
- LogicalType::INTEGER(IntType {
- bit_width: 8,
- is_signed: true,
- }),
- LogicalType::INTEGER(IntType {
- bit_width: 8,
- is_signed: true,
- }),
- LogicalType::INTEGER(IntType {
- bit_width: 8,
- is_signed: true,
- }),
- LogicalType::DECIMAL(DecimalType {
- scale: 20,
- precision: 4,
- }),
- LogicalType::DATE(Default::default()),
- LogicalType::TIME(TimeType {
- is_adjusted_to_u_t_c: false,
- unit: TimeUnit::MILLIS(Default::default()),
- }),
- LogicalType::TIME(TimeType {
- is_adjusted_to_u_t_c: false,
- unit: TimeUnit::MICROS(Default::default()),
- }),
- LogicalType::TIME(TimeType {
- is_adjusted_to_u_t_c: true,
- unit: TimeUnit::NANOS(Default::default()),
- }),
- LogicalType::TIMESTAMP(TimestampType {
- is_adjusted_to_u_t_c: false,
- unit: TimeUnit::MILLIS(Default::default()),
- }),
- LogicalType::TIMESTAMP(TimestampType {
- is_adjusted_to_u_t_c: false,
- unit: TimeUnit::MICROS(Default::default()),
- }),
- LogicalType::TIMESTAMP(TimestampType {
- is_adjusted_to_u_t_c: true,
- unit: TimeUnit::NANOS(Default::default()),
- }),
- ];
- check_sort_order(signed, SortOrder::SIGNED);
-
- // Undefined comparison
- let undefined = vec![
- LogicalType::LIST(Default::default()),
- LogicalType::MAP(Default::default()),
- ];
- check_sort_order(undefined, SortOrder::UNDEFINED);
- }
-
- #[test]
- fn test_column_order_get_coverted_type_sort_order() {
- // Helper to check the order in a list of values.
- // Only converted type is checked.
- fn check_sort_order(types: Vec<ConvertedType>, expected_order: SortOrder) {
- for tpe in types {
- assert_eq!(
- ColumnOrder::get_sort_order(None, tpe, Type::BYTE_ARRAY),
- expected_order
- );
- }
- }
-
- // Unsigned comparison (physical type does not matter)
- let unsigned = vec![
- ConvertedType::UTF8,
- ConvertedType::JSON,
- ConvertedType::BSON,
- ConvertedType::ENUM,
- ConvertedType::UINT_8,
- ConvertedType::UINT_16,
- ConvertedType::UINT_32,
- ConvertedType::UINT_64,
- ];
- check_sort_order(unsigned, SortOrder::UNSIGNED);
-
- // Signed comparison (physical type does not matter)
- let signed = vec![
- ConvertedType::INT_8,
- ConvertedType::INT_16,
- ConvertedType::INT_32,
- ConvertedType::INT_64,
- ConvertedType::DECIMAL,
- ConvertedType::DATE,
- ConvertedType::TIME_MILLIS,
- ConvertedType::TIME_MICROS,
- ConvertedType::TIMESTAMP_MILLIS,
- ConvertedType::TIMESTAMP_MICROS,
- ];
- check_sort_order(signed, SortOrder::SIGNED);
-
- // Undefined comparison
- let undefined = vec![
- ConvertedType::LIST,
- ConvertedType::MAP,
- ConvertedType::MAP_KEY_VALUE,
- ConvertedType::INTERVAL,
- ];
- check_sort_order(undefined, SortOrder::UNDEFINED);
-
- // Check None logical type
- // This should return a sort order for byte array type.
- check_sort_order(vec![ConvertedType::NONE], SortOrder::UNSIGNED);
- }
-
- #[test]
- fn test_column_order_get_default_sort_order() {
- // Comparison based on physical type
- assert_eq!(
- ColumnOrder::get_default_sort_order(Type::BOOLEAN),
- SortOrder::UNSIGNED
- );
- assert_eq!(
- ColumnOrder::get_default_sort_order(Type::INT32),
- SortOrder::SIGNED
- );
- assert_eq!(
- ColumnOrder::get_default_sort_order(Type::INT64),
- SortOrder::SIGNED
- );
- assert_eq!(
- ColumnOrder::get_default_sort_order(Type::INT96),
- SortOrder::UNDEFINED
- );
- assert_eq!(
- ColumnOrder::get_default_sort_order(Type::FLOAT),
- SortOrder::SIGNED
- );
- assert_eq!(
- ColumnOrder::get_default_sort_order(Type::DOUBLE),
- SortOrder::SIGNED
- );
- assert_eq!(
- ColumnOrder::get_default_sort_order(Type::BYTE_ARRAY),
- SortOrder::UNSIGNED
- );
- assert_eq!(
- ColumnOrder::get_default_sort_order(Type::FIXED_LEN_BYTE_ARRAY),
- SortOrder::UNSIGNED
- );
- }
-
- #[test]
- fn test_column_order_sort_order() {
- assert_eq!(
- ColumnOrder::TYPE_DEFINED_ORDER(SortOrder::SIGNED).sort_order(),
- SortOrder::SIGNED
- );
- assert_eq!(
- ColumnOrder::TYPE_DEFINED_ORDER(SortOrder::UNSIGNED).sort_order(),
- SortOrder::UNSIGNED
- );
- assert_eq!(
- ColumnOrder::TYPE_DEFINED_ORDER(SortOrder::UNDEFINED).sort_order(),
- SortOrder::UNDEFINED
- );
- assert_eq!(ColumnOrder::UNDEFINED.sort_order(), SortOrder::SIGNED);
- }
-}
diff --git a/parquet/src/bin/parquet-read.rs b/parquet/src/bin/parquet-read.rs
deleted file mode 100644
index aa3b827..0000000
--- a/parquet/src/bin/parquet-read.rs
+++ /dev/null
@@ -1,128 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-//! Binary file to read data from a Parquet file.
-//!
-//! # Install
-//!
-//! `parquet-read` can be installed using `cargo`:
-//! ```
-//! cargo install parquet
-//! ```
-//! After this `parquet-read` should be globally available:
-//! ```
-//! parquet-read XYZ.parquet
-//! ```
-//!
-//! The binary can also be built from the source code and run as follows:
-//! ```
-//! cargo run --bin parquet-read XYZ.parquet
-//! ```
-//!
-//! # Usage
-//! ```
-//! parquet-read <file-path> [num-records]
-//! ```
-//!
-//! ## Flags
-//! -h, --help Prints help information
-//! -j, --json Print Parquet file in JSON lines Format
-//! -V, --version Prints version information
-//!
-//! ## Args
-//! <file-path> Path to a Parquet file
-//! <num-records> Number of records to read. When not provided, all records are read.
-//!
-//! Note that `parquet-read` reads full file schema, no projection or filtering is
-//! applied.
-
-extern crate parquet;
-
-use std::{env, fs::File, path::Path};
-
-use clap::{crate_authors, crate_version, App, Arg};
-
-use parquet::file::reader::{FileReader, SerializedFileReader};
-use parquet::record::Row;
-
-fn main() {
- let app = App::new("parquet-read")
- .version(crate_version!())
- .author(crate_authors!())
- .about("Read data from a Parquet file and print output in console, in either built-in or JSON format")
- .arg(
- Arg::with_name("file_path")
- .value_name("file-path")
- .required(true)
- .index(1)
- .help("Path to a parquet file"),
- )
- .arg(
- Arg::with_name("num_records")
- .value_name("num-records")
- .index(2)
- .help(
- "Number of records to read. When not provided, all records are read.",
- ),
- )
- .arg(
- Arg::with_name("json")
- .short("j")
- .long("json")
- .takes_value(false)
- .help("Print Parquet file in JSON lines format"),
- );
-
- let matches = app.get_matches();
- let filename = matches.value_of("file_path").unwrap();
- let num_records: Option<usize> = if matches.is_present("num_records") {
- match matches.value_of("num_records").unwrap().parse() {
- Ok(value) => Some(value),
- Err(e) => panic!("Error when reading value for [num-records], {}", e),
- }
- } else {
- None
- };
-
- let json = matches.is_present("json");
- let path = Path::new(&filename);
- let file = File::open(&path).unwrap();
- let parquet_reader = SerializedFileReader::new(file).unwrap();
-
- // Use full schema as projected schema
- let mut iter = parquet_reader.get_row_iter(None).unwrap();
-
- let mut start = 0;
- let end = num_records.unwrap_or(0);
- let all_records = num_records.is_none();
-
- while all_records || start < end {
- match iter.next() {
- Some(row) => print_row(&row, json),
- None => break,
- }
- start += 1;
- }
-}
-
-fn print_row(row: &Row, json: bool) {
- if json {
- println!("{}", row.to_json_value())
- } else {
- println!("{}", row.to_string());
- }
-}
diff --git a/parquet/src/bin/parquet-rowcount.rs b/parquet/src/bin/parquet-rowcount.rs
deleted file mode 100644
index 3c61bab..0000000
--- a/parquet/src/bin/parquet-rowcount.rs
+++ /dev/null
@@ -1,87 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-//! Binary file to return the number of rows found from Parquet file(s).
-//!
-//! # Install
-//!
-//! `parquet-rowcount` can be installed using `cargo`:
-//! ```
-//! cargo install parquet
-//! ```
-//! After this `parquet-rowcount` should be globally available:
-//! ```
-//! parquet-rowcount XYZ.parquet
-//! ```
-//!
-//! The binary can also be built from the source code and run as follows:
-//! ```
-//! cargo run --bin parquet-rowcount XYZ.parquet ABC.parquet ZXC.parquet
-//! ```
-//!
-//! # Usage
-//! ```
-//! parquet-rowcount <file-paths>...
-//! ```
-//!
-//! ## Flags
-//! -h, --help Prints help information
-//! -V, --version Prints version information
-//!
-//! ## Args
-//! <file-paths>... List of Parquet files to read from
-//!
-//! Note that `parquet-rowcount` reads full file schema, no projection or filtering is
-//! applied.
-
-extern crate parquet;
-
-use std::{env, fs::File, path::Path};
-
-use clap::{crate_authors, crate_version, App, Arg};
-
-use parquet::file::reader::{FileReader, SerializedFileReader};
-
-fn main() {
- let matches = App::new("parquet-rowcount")
- .version(crate_version!())
- .author(crate_authors!())
- .about("Return number of rows in Parquet file")
- .arg(
- Arg::with_name("file_paths")
- .value_name("file-paths")
- .required(true)
- .multiple(true)
- .help("List of Parquet files to read from separated by space"),
- )
- .get_matches();
-
- let filenames: Vec<&str> = matches.values_of("file_paths").unwrap().collect();
- for filename in &filenames {
- let path = Path::new(filename);
- let file = File::open(path).unwrap();
- let parquet_reader = SerializedFileReader::new(file).unwrap();
- let row_group_metadata = parquet_reader.metadata().row_groups();
- let mut total_num_rows = 0;
-
- for group_metadata in row_group_metadata {
- total_num_rows += group_metadata.num_rows();
- }
-
- eprintln!("File {}: rowcount={}", filename, total_num_rows);
- }
-}
diff --git a/parquet/src/bin/parquet-schema.rs b/parquet/src/bin/parquet-schema.rs
deleted file mode 100644
index 1b80637..0000000
--- a/parquet/src/bin/parquet-schema.rs
+++ /dev/null
@@ -1,104 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-//! Binary file to print the schema and metadata of a Parquet file.
-//!
-//! # Install
-//!
-//! `parquet-schema` can be installed using `cargo`:
-//! ```
-//! cargo install parquet
-//! ```
-//! After this `parquet-schema` should be globally available:
-//! ```
-//! parquet-schema XYZ.parquet
-//! ```
-//!
-//! The binary can also be built from the source code and run as follows:
-//! ```
-//! cargo run --bin parquet-schema XYZ.parquet
-//! ```
-//!
-//! # Usage
-//! ```
-//! parquet-schema [FLAGS] <file-path>
-//! ```
-//!
-//! ## Flags
-//! -h, --help Prints help information
-//! -V, --version Prints version information
-//! -v, --verbose Enable printing full file metadata
-//!
-//! ## Args
-//! <file-path> Path to a Parquet file
-//!
-//! Note that `verbose` is an optional boolean flag that allows to print schema only,
-//! when not provided or print full file metadata when provided.
-
-extern crate parquet;
-
-use std::{env, fs::File, path::Path};
-
-use clap::{crate_authors, crate_version, App, Arg};
-
-use parquet::{
- file::reader::{FileReader, SerializedFileReader},
- schema::printer::{print_file_metadata, print_parquet_metadata},
-};
-
-fn main() {
- let matches = App::new("parquet-schema")
- .version(crate_version!())
- .author(crate_authors!())
- .arg(
- Arg::with_name("file_path")
- .value_name("file-path")
- .required(true)
- .index(1)
- .help("Path to a Parquet file"),
- )
- .arg(
- Arg::with_name("verbose")
- .short("v")
- .long("verbose")
- .takes_value(false)
- .help("Enable printing full file metadata"),
- )
- .get_matches();
-
- let filename = matches.value_of("file_path").unwrap();
- let path = Path::new(&filename);
- let file = match File::open(&path) {
- Err(e) => panic!("Error when opening file {}: {}", path.display(), e),
- Ok(f) => f,
- };
- let verbose = matches.is_present("verbose");
-
- match SerializedFileReader::new(file) {
- Err(e) => panic!("Error when parsing Parquet file: {}", e),
- Ok(parquet_reader) => {
- let metadata = parquet_reader.metadata();
- println!("Metadata for file: {}", &filename);
- println!();
- if verbose {
- print_parquet_metadata(&mut std::io::stdout(), &metadata);
- } else {
- print_file_metadata(&mut std::io::stdout(), &metadata.file_metadata());
- }
- }
- }
-}
diff --git a/parquet/src/column/mod.rs b/parquet/src/column/mod.rs
deleted file mode 100644
index 7ed7bfc..0000000
--- a/parquet/src/column/mod.rs
+++ /dev/null
@@ -1,124 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-//! Low level column reader and writer APIs.
-//!
-//! This API is designed for reading and writing column values, definition and repetition
-//! levels directly.
-//!
-//! # Example of writing and reading data
-//!
-//! Data has the following format:
-//! ```text
-//! +---------------+
-//! | values|
-//! +---------------+
-//! |[1, 2] |
-//! |[3, null, null]|
-//! +---------------+
-//! ```
-//!
-//! The example uses column writer and reader APIs to write raw values, definition and
-//! repetition levels and read them to verify write/read correctness.
-//!
-//! ```rust,no_run
-//! use std::{fs, path::Path, sync::Arc};
-//!
-//! use parquet::{
-//! column::{reader::ColumnReader, writer::ColumnWriter},
-//! file::{
-//! properties::WriterProperties,
-//! reader::{FileReader, SerializedFileReader},
-//! writer::{FileWriter, SerializedFileWriter},
-//! },
-//! schema::parser::parse_message_type,
-//! };
-//!
-//! let path = Path::new("/path/to/column_sample.parquet");
-//!
-//! // Writing data using column writer API.
-//!
-//! let message_type = "
-//! message schema {
-//! optional group values (LIST) {
-//! repeated group list {
-//! optional INT32 element;
-//! }
-//! }
-//! }
-//! ";
-//! let schema = Arc::new(parse_message_type(message_type).unwrap());
-//! let props = Arc::new(WriterProperties::builder().build());
-//! let file = fs::File::create(path).unwrap();
-//! let mut writer = SerializedFileWriter::new(file, schema, props).unwrap();
-//! let mut row_group_writer = writer.next_row_group().unwrap();
-//! while let Some(mut col_writer) = row_group_writer.next_column().unwrap() {
-//! match col_writer {
-//! // You can also use `get_typed_column_writer` method to extract typed writer.
-//! ColumnWriter::Int32ColumnWriter(ref mut typed_writer) => {
-//! typed_writer
-//! .write_batch(&[1, 2, 3], Some(&[3, 3, 3, 2, 2]), Some(&[0, 1, 0, 1, 1]))
-//! .unwrap();
-//! }
-//! _ => {}
-//! }
-//! row_group_writer.close_column(col_writer).unwrap();
-//! }
-//! writer.close_row_group(row_group_writer).unwrap();
-//! writer.close().unwrap();
-//!
-//! // Reading data using column reader API.
-//!
-//! let file = fs::File::open(path).unwrap();
-//! let reader = SerializedFileReader::new(file).unwrap();
-//! let metadata = reader.metadata();
-//!
-//! let mut res = Ok((0, 0));
-//! let mut values = vec![0; 8];
-//! let mut def_levels = vec![0; 8];
-//! let mut rep_levels = vec![0; 8];
-//!
-//! for i in 0..metadata.num_row_groups() {
-//! let row_group_reader = reader.get_row_group(i).unwrap();
-//! let row_group_metadata = metadata.row_group(i);
-//!
-//! for j in 0..row_group_metadata.num_columns() {
-//! let mut column_reader = row_group_reader.get_column_reader(j).unwrap();
-//! match column_reader {
-//! // You can also use `get_typed_column_reader` method to extract typed reader.
-//! ColumnReader::Int32ColumnReader(ref mut typed_reader) => {
-//! res = typed_reader.read_batch(
-//! 8, // batch size
-//! Some(&mut def_levels),
-//! Some(&mut rep_levels),
-//! &mut values,
-//! );
-//! }
-//! _ => {}
-//! }
-//! }
-//! }
-//!
-//! assert_eq!(res, Ok((3, 5)));
-//! assert_eq!(values, vec![1, 2, 3, 0, 0, 0, 0, 0]);
-//! assert_eq!(def_levels, vec![3, 3, 3, 2, 2, 0, 0, 0]);
-//! assert_eq!(rep_levels, vec![0, 1, 0, 1, 1, 0, 0, 0]);
-//! ```
-
-pub mod page;
-pub mod reader;
-pub mod writer;
diff --git a/parquet/src/column/page.rs b/parquet/src/column/page.rs
deleted file mode 100644
index b75d3b5..0000000
--- a/parquet/src/column/page.rs
+++ /dev/null
@@ -1,307 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-//! Contains Parquet Page definitions and page reader interface.
-
-use crate::basic::{Encoding, PageType};
-use crate::errors::Result;
-use crate::file::{metadata::ColumnChunkMetaData, statistics::Statistics};
-use crate::schema::types::{ColumnDescPtr, SchemaDescPtr};
-use crate::util::memory::ByteBufferPtr;
-
-/// Parquet Page definition.
-///
-/// List of supported pages.
-/// These are 1-to-1 mapped from the equivalent Thrift definitions, except `buf` which
-/// used to store uncompressed bytes of the page.
-#[derive(Clone)]
-pub enum Page {
- DataPage {
- buf: ByteBufferPtr,
- num_values: u32,
- encoding: Encoding,
- def_level_encoding: Encoding,
- rep_level_encoding: Encoding,
- statistics: Option<Statistics>,
- },
- DataPageV2 {
- buf: ByteBufferPtr,
- num_values: u32,
- encoding: Encoding,
- num_nulls: u32,
- num_rows: u32,
- def_levels_byte_len: u32,
- rep_levels_byte_len: u32,
- is_compressed: bool,
- statistics: Option<Statistics>,
- },
- DictionaryPage {
- buf: ByteBufferPtr,
- num_values: u32,
- encoding: Encoding,
- is_sorted: bool,
- },
-}
-
-impl Page {
- /// Returns [`PageType`](crate::basic::PageType) for this page.
- pub fn page_type(&self) -> PageType {
- match self {
- Page::DataPage { .. } => PageType::DATA_PAGE,
- Page::DataPageV2 { .. } => PageType::DATA_PAGE_V2,
- Page::DictionaryPage { .. } => PageType::DICTIONARY_PAGE,
- }
- }
-
- /// Returns internal byte buffer reference for this page.
- pub fn buffer(&self) -> &ByteBufferPtr {
- match self {
- Page::DataPage { ref buf, .. } => &buf,
- Page::DataPageV2 { ref buf, .. } => &buf,
- Page::DictionaryPage { ref buf, .. } => &buf,
- }
- }
-
- /// Returns number of values in this page.
- pub fn num_values(&self) -> u32 {
- match self {
- Page::DataPage { num_values, .. } => *num_values,
- Page::DataPageV2 { num_values, .. } => *num_values,
- Page::DictionaryPage { num_values, .. } => *num_values,
- }
- }
-
- /// Returns this page [`Encoding`](crate::basic::Encoding).
- pub fn encoding(&self) -> Encoding {
- match self {
- Page::DataPage { encoding, .. } => *encoding,
- Page::DataPageV2 { encoding, .. } => *encoding,
- Page::DictionaryPage { encoding, .. } => *encoding,
- }
- }
-
- /// Returns optional [`Statistics`](crate::file::statistics::Statistics).
- pub fn statistics(&self) -> Option<&Statistics> {
- match self {
- Page::DataPage { ref statistics, .. } => statistics.as_ref(),
- Page::DataPageV2 { ref statistics, .. } => statistics.as_ref(),
- Page::DictionaryPage { .. } => None,
- }
- }
-}
-
-/// Helper struct to represent pages with potentially compressed buffer (data page v1) or
-/// compressed and concatenated buffer (def levels + rep levels + compressed values for
-/// data page v2).
-///
-/// The difference with `Page` is that `Page` buffer is always uncompressed.
-pub struct CompressedPage {
- compressed_page: Page,
- uncompressed_size: usize,
-}
-
-impl CompressedPage {
- /// Creates `CompressedPage` from a page with potentially compressed buffer and
- /// uncompressed size.
- pub fn new(compressed_page: Page, uncompressed_size: usize) -> Self {
- Self {
- compressed_page,
- uncompressed_size,
- }
- }
-
- /// Returns page type.
- pub fn page_type(&self) -> PageType {
- self.compressed_page.page_type()
- }
-
- /// Returns underlying page with potentially compressed buffer.
- pub fn compressed_page(&self) -> &Page {
- &self.compressed_page
- }
-
- /// Returns uncompressed size in bytes.
- pub fn uncompressed_size(&self) -> usize {
- self.uncompressed_size
- }
-
- /// Returns compressed size in bytes.
- ///
- /// Note that it is assumed that buffer is compressed, but it may not be. In this
- /// case compressed size will be equal to uncompressed size.
- pub fn compressed_size(&self) -> usize {
- self.compressed_page.buffer().len()
- }
-
- /// Number of values in page.
- pub fn num_values(&self) -> u32 {
- self.compressed_page.num_values()
- }
-
- /// Returns encoding for values in page.
- pub fn encoding(&self) -> Encoding {
- self.compressed_page.encoding()
- }
-
- /// Returns slice of compressed buffer in the page.
- pub fn data(&self) -> &[u8] {
- self.compressed_page.buffer().data()
- }
-}
-
-/// Contains page write metrics.
-pub struct PageWriteSpec {
- pub page_type: PageType,
- pub uncompressed_size: usize,
- pub compressed_size: usize,
- pub num_values: u32,
- pub offset: u64,
- pub bytes_written: u64,
-}
-
-impl PageWriteSpec {
- /// Creates new spec with default page write metrics.
- pub fn new() -> Self {
- Self {
- page_type: PageType::DATA_PAGE,
- uncompressed_size: 0,
- compressed_size: 0,
- num_values: 0,
- offset: 0,
- bytes_written: 0,
- }
- }
-}
-
-/// API for reading pages from a column chunk.
-/// This offers a iterator like API to get the next page.
-pub trait PageReader: Iterator<Item = Result<Page>> {
- /// Gets the next page in the column chunk associated with this reader.
- /// Returns `None` if there are no pages left.
- fn get_next_page(&mut self) -> Result<Option<Page>>;
-}
-
-/// API for writing pages in a column chunk.
-///
-/// It is reasonable to assume that all pages will be written in the correct order, e.g.
-/// dictionary page followed by data pages, or a set of data pages, etc.
-pub trait PageWriter {
- /// Writes a page into the output stream/sink.
- /// Returns `PageWriteSpec` that contains information about written page metrics,
- /// including number of bytes, size, number of values, offset, etc.
- ///
- /// This method is called for every compressed page we write into underlying buffer,
- /// either data page or dictionary page.
- fn write_page(&mut self, page: CompressedPage) -> Result<PageWriteSpec>;
-
- /// Writes column chunk metadata into the output stream/sink.
- ///
- /// This method is called once before page writer is closed, normally when writes are
- /// finalised in column writer.
- fn write_metadata(&mut self, metadata: &ColumnChunkMetaData) -> Result<()>;
-
- /// Closes resources and flushes underlying sink.
- /// Page writer should not be used after this method is called.
- fn close(&mut self) -> Result<()>;
-}
-
-/// An iterator over pages of some specific column in a parquet file.
-pub trait PageIterator: Iterator<Item = Result<Box<dyn PageReader>>> {
- /// Get schema of parquet file.
- fn schema(&mut self) -> Result<SchemaDescPtr>;
-
- /// Get column schema of this page iterator.
- fn column_schema(&mut self) -> Result<ColumnDescPtr>;
-}
-
-#[cfg(test)]
-mod tests {
- use super::*;
-
- #[test]
- fn test_page() {
- let data_page = Page::DataPage {
- buf: ByteBufferPtr::new(vec![0, 1, 2]),
- num_values: 10,
- encoding: Encoding::PLAIN,
- def_level_encoding: Encoding::RLE,
- rep_level_encoding: Encoding::RLE,
- statistics: Some(Statistics::int32(Some(1), Some(2), None, 1, true)),
- };
- assert_eq!(data_page.page_type(), PageType::DATA_PAGE);
- assert_eq!(data_page.buffer().data(), vec![0, 1, 2].as_slice());
- assert_eq!(data_page.num_values(), 10);
- assert_eq!(data_page.encoding(), Encoding::PLAIN);
- assert_eq!(
- data_page.statistics(),
- Some(&Statistics::int32(Some(1), Some(2), None, 1, true))
- );
-
- let data_page_v2 = Page::DataPageV2 {
- buf: ByteBufferPtr::new(vec![0, 1, 2]),
- num_values: 10,
- encoding: Encoding::PLAIN,
- num_nulls: 5,
- num_rows: 20,
- def_levels_byte_len: 30,
- rep_levels_byte_len: 40,
- is_compressed: false,
- statistics: Some(Statistics::int32(Some(1), Some(2), None, 1, true)),
- };
- assert_eq!(data_page_v2.page_type(), PageType::DATA_PAGE_V2);
- assert_eq!(data_page_v2.buffer().data(), vec![0, 1, 2].as_slice());
- assert_eq!(data_page_v2.num_values(), 10);
- assert_eq!(data_page_v2.encoding(), Encoding::PLAIN);
- assert_eq!(
- data_page_v2.statistics(),
- Some(&Statistics::int32(Some(1), Some(2), None, 1, true))
- );
-
- let dict_page = Page::DictionaryPage {
- buf: ByteBufferPtr::new(vec![0, 1, 2]),
- num_values: 10,
- encoding: Encoding::PLAIN,
- is_sorted: false,
- };
- assert_eq!(dict_page.page_type(), PageType::DICTIONARY_PAGE);
- assert_eq!(dict_page.buffer().data(), vec![0, 1, 2].as_slice());
- assert_eq!(dict_page.num_values(), 10);
- assert_eq!(dict_page.encoding(), Encoding::PLAIN);
- assert_eq!(dict_page.statistics(), None);
- }
-
- #[test]
- fn test_compressed_page() {
- let data_page = Page::DataPage {
- buf: ByteBufferPtr::new(vec![0, 1, 2]),
- num_values: 10,
- encoding: Encoding::PLAIN,
- def_level_encoding: Encoding::RLE,
- rep_level_encoding: Encoding::RLE,
- statistics: Some(Statistics::int32(Some(1), Some(2), None, 1, true)),
- };
-
- let cpage = CompressedPage::new(data_page, 5);
-
- assert_eq!(cpage.page_type(), PageType::DATA_PAGE);
- assert_eq!(cpage.uncompressed_size(), 5);
- assert_eq!(cpage.compressed_size(), 3);
- assert_eq!(cpage.num_values(), 10);
- assert_eq!(cpage.encoding(), Encoding::PLAIN);
- assert_eq!(cpage.data(), &[0, 1, 2]);
- }
-}
diff --git a/parquet/src/column/reader.rs b/parquet/src/column/reader.rs
deleted file mode 100644
index 63be17b..0000000
--- a/parquet/src/column/reader.rs
+++ /dev/null
@@ -1,1364 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-//! Contains column reader API.
-
-use std::{
- cmp::{max, min},
- collections::HashMap,
-};
-
-use super::page::{Page, PageReader};
-use crate::basic::*;
-use crate::data_type::*;
-use crate::encodings::{
- decoding::{get_decoder, Decoder, DictDecoder, PlainDecoder},
- levels::LevelDecoder,
-};
-use crate::errors::{ParquetError, Result};
-use crate::schema::types::ColumnDescPtr;
-use crate::util::memory::ByteBufferPtr;
-
-/// Column reader for a Parquet type.
-pub enum ColumnReader {
- BoolColumnReader(ColumnReaderImpl<BoolType>),
- Int32ColumnReader(ColumnReaderImpl<Int32Type>),
- Int64ColumnReader(ColumnReaderImpl<Int64Type>),
- Int96ColumnReader(ColumnReaderImpl<Int96Type>),
- FloatColumnReader(ColumnReaderImpl<FloatType>),
- DoubleColumnReader(ColumnReaderImpl<DoubleType>),
- ByteArrayColumnReader(ColumnReaderImpl<ByteArrayType>),
- FixedLenByteArrayColumnReader(ColumnReaderImpl<FixedLenByteArrayType>),
-}
-
-/// Gets a specific column reader corresponding to column descriptor `col_descr`. The
-/// column reader will read from pages in `col_page_reader`.
-pub fn get_column_reader(
- col_descr: ColumnDescPtr,
- col_page_reader: Box<dyn PageReader>,
-) -> ColumnReader {
- match col_descr.physical_type() {
- Type::BOOLEAN => ColumnReader::BoolColumnReader(ColumnReaderImpl::new(
- col_descr,
- col_page_reader,
- )),
- Type::INT32 => ColumnReader::Int32ColumnReader(ColumnReaderImpl::new(
- col_descr,
- col_page_reader,
- )),
- Type::INT64 => ColumnReader::Int64ColumnReader(ColumnReaderImpl::new(
- col_descr,
- col_page_reader,
- )),
- Type::INT96 => ColumnReader::Int96ColumnReader(ColumnReaderImpl::new(
- col_descr,
- col_page_reader,
- )),
- Type::FLOAT => ColumnReader::FloatColumnReader(ColumnReaderImpl::new(
- col_descr,
- col_page_reader,
- )),
- Type::DOUBLE => ColumnReader::DoubleColumnReader(ColumnReaderImpl::new(
- col_descr,
- col_page_reader,
- )),
- Type::BYTE_ARRAY => ColumnReader::ByteArrayColumnReader(ColumnReaderImpl::new(
- col_descr,
- col_page_reader,
- )),
- Type::FIXED_LEN_BYTE_ARRAY => ColumnReader::FixedLenByteArrayColumnReader(
- ColumnReaderImpl::new(col_descr, col_page_reader),
- ),
- }
-}
-
-/// Gets a typed column reader for the specific type `T`, by "up-casting" `col_reader` of
-/// non-generic type to a generic column reader type `ColumnReaderImpl`.
-///
-/// Panics if actual enum value for `col_reader` does not match the type `T`.
-pub fn get_typed_column_reader<T: DataType>(
- col_reader: ColumnReader,
-) -> ColumnReaderImpl<T> {
- T::get_column_reader(col_reader).unwrap_or_else(|| {
- panic!(
- "Failed to convert column reader into a typed column reader for `{}` type",
- T::get_physical_type()
- )
- })
-}
-
-/// Typed value reader for a particular primitive column.
-pub struct ColumnReaderImpl<T: DataType> {
- descr: ColumnDescPtr,
- def_level_decoder: Option<LevelDecoder>,
- rep_level_decoder: Option<LevelDecoder>,
- page_reader: Box<dyn PageReader>,
- current_encoding: Option<Encoding>,
-
- // The total number of values stored in the data page.
- num_buffered_values: u32,
-
- // The number of values from the current data page that has been decoded into memory
- // so far.
- num_decoded_values: u32,
-
- // Cache of decoders for existing encodings
- decoders: HashMap<Encoding, Box<dyn Decoder<T>>>,
-}
-
-impl<T: DataType> ColumnReaderImpl<T> {
- /// Creates new column reader based on column descriptor and page reader.
- pub fn new(descr: ColumnDescPtr, page_reader: Box<dyn PageReader>) -> Self {
- Self {
- descr,
- def_level_decoder: None,
- rep_level_decoder: None,
- page_reader,
- current_encoding: None,
- num_buffered_values: 0,
- num_decoded_values: 0,
- decoders: HashMap::new(),
- }
- }
-
- /// Reads a batch of values of at most `batch_size`.
- ///
- /// This will try to read from the row group, and fills up at most `batch_size` values
- /// for `def_levels`, `rep_levels` and `values`. It will stop either when the row
- /// group is depleted or `batch_size` values has been read, or there is no space
- /// in the input slices (values/definition levels/repetition levels).
- ///
- /// Note that in case the field being read is not required, `values` could contain
- /// less values than `def_levels`. Also note that this will skip reading def / rep
- /// levels if the field is required / not repeated, respectively.
- ///
- /// If `def_levels` or `rep_levels` is `None`, this will also skip reading the
- /// respective levels. This is useful when the caller of this function knows in
- /// advance that the field is required and non-repeated, therefore can avoid
- /// allocating memory for the levels data. Note that if field has definition
- /// levels, but caller provides None, there might be inconsistency between
- /// levels/values (see comments below).
- ///
- /// Returns a tuple where the first element is the actual number of values read,
- /// and the second element is the actual number of levels read.
- #[inline]
- pub fn read_batch(
- &mut self,
- batch_size: usize,
- mut def_levels: Option<&mut [i16]>,
- mut rep_levels: Option<&mut [i16]>,
- values: &mut [T::T],
- ) -> Result<(usize, usize)> {
- let mut values_read = 0;
- let mut levels_read = 0;
-
- // Compute the smallest batch size we can read based on provided slices
- let mut batch_size = min(batch_size, values.len());
- if let Some(ref levels) = def_levels {
- batch_size = min(batch_size, levels.len());
- }
- if let Some(ref levels) = rep_levels {
- batch_size = min(batch_size, levels.len());
- }
-
- // Read exhaustively all pages until we read all batch_size values/levels
- // or there are no more values/levels to read.
- while max(values_read, levels_read) < batch_size {
- if !self.has_next()? {
- break;
- }
-
- // Batch size for the current iteration
- let iter_batch_size = {
- // Compute approximate value based on values decoded so far
- let mut adjusted_size = min(
- batch_size,
- (self.num_buffered_values - self.num_decoded_values) as usize,
- );
-
- // Adjust batch size by taking into account how much data there
- // to read. As batch_size is also smaller than value and level
- // slices (if available), this ensures that available space is not
- // exceeded.
- adjusted_size = min(adjusted_size, batch_size - values_read);
- adjusted_size = min(adjusted_size, batch_size - levels_read);
-
- adjusted_size
- };
-
- let mut values_to_read = 0;
- let mut num_def_levels = 0;
- let mut num_rep_levels = 0;
-
- // If the field is required and non-repeated, there are no definition levels
- if self.descr.max_def_level() > 0 && def_levels.as_ref().is_some() {
- if let Some(ref mut levels) = def_levels {
- num_def_levels = self.read_def_levels(
- &mut levels[levels_read..levels_read + iter_batch_size],
- )?;
- for i in levels_read..levels_read + num_def_levels {
- if levels[i] == self.descr.max_def_level() {
- values_to_read += 1;
- }
- }
- }
- } else {
- // If max definition level == 0, then it is REQUIRED field, read all
- // values. If definition levels are not provided, we still
- // read all values.
- values_to_read = iter_batch_size;
- }
-
- if self.descr.max_rep_level() > 0 && rep_levels.is_some() {
- if let Some(ref mut levels) = rep_levels {
- num_rep_levels = self.read_rep_levels(
- &mut levels[levels_read..levels_read + iter_batch_size],
- )?;
-
- // If definition levels are defined, check that rep levels == def
- // levels
- if def_levels.is_some() {
- assert_eq!(
- num_def_levels, num_rep_levels,
- "Number of decoded rep / def levels did not match"
- );
- }
- }
- }
-
- // At this point we have read values, definition and repetition levels.
- // If both definition and repetition levels are defined, their counts
- // should be equal. Values count is always less or equal to definition levels.
- //
- // Note that if field is not required, but no definition levels are provided,
- // we would read values of batch size and (if provided, of course) repetition
- // levels of batch size - [!] they will not be synced, because only definition
- // levels enforce number of non-null values to read.
-
- let curr_values_read =
- self.read_values(&mut values[values_read..values_read + values_to_read])?;
-
- // Update all "return" counters and internal state.
-
- // This is to account for when def or rep levels are not provided
- let curr_levels_read = max(num_def_levels, num_rep_levels);
- self.num_decoded_values += max(curr_levels_read, curr_values_read) as u32;
- levels_read += curr_levels_read;
- values_read += curr_values_read;
- }
-
- Ok((values_read, levels_read))
- }
-
- /// Reads a new page and set up the decoders for levels, values or dictionary.
- /// Returns false if there's no page left.
- fn read_new_page(&mut self) -> Result<bool> {
- #[allow(while_true)]
- while true {
- match self.page_reader.get_next_page()? {
- // No more page to read
- None => return Ok(false),
- Some(current_page) => {
- match current_page {
- // 1. Dictionary page: configure dictionary for this page.
- p @ Page::DictionaryPage { .. } => {
- self.configure_dictionary(p)?;
- continue;
- }
- // 2. Data page v1
- Page::DataPage {
- buf,
- num_values,
- encoding,
- def_level_encoding,
- rep_level_encoding,
- statistics: _,
- } => {
- self.num_buffered_values = num_values;
- self.num_decoded_values = 0;
-
- let mut buffer_ptr = buf;
-
- if self.descr.max_rep_level() > 0 {
- let mut rep_decoder = LevelDecoder::v1(
- rep_level_encoding,
- self.descr.max_rep_level(),
- );
- let total_bytes = rep_decoder.set_data(
- self.num_buffered_values as usize,
- buffer_ptr.all(),
- );
- buffer_ptr = buffer_ptr.start_from(total_bytes);
- self.rep_level_decoder = Some(rep_decoder);
- }
-
- if self.descr.max_def_level() > 0 {
- let mut def_decoder = LevelDecoder::v1(
- def_level_encoding,
- self.descr.max_def_level(),
- );
- let total_bytes = def_decoder.set_data(
- self.num_buffered_values as usize,
- buffer_ptr.all(),
- );
- buffer_ptr = buffer_ptr.start_from(total_bytes);
- self.def_level_decoder = Some(def_decoder);
- }
-
- // Data page v1 does not have offset, all content of buffer
- // should be passed
- self.set_current_page_encoding(
- encoding,
- &buffer_ptr,
- 0,
- num_values as usize,
- )?;
- return Ok(true);
- }
- // 3. Data page v2
- Page::DataPageV2 {
- buf,
- num_values,
- encoding,
- num_nulls: _,
- num_rows: _,
- def_levels_byte_len,
- rep_levels_byte_len,
- is_compressed: _,
- statistics: _,
- } => {
- self.num_buffered_values = num_values;
- self.num_decoded_values = 0;
-
- let mut offset = 0;
-
- // DataPage v2 only supports RLE encoding for repetition
- // levels
- if self.descr.max_rep_level() > 0 {
- let mut rep_decoder =
- LevelDecoder::v2(self.descr.max_rep_level());
- let bytes_read = rep_decoder.set_data_range(
- self.num_buffered_values as usize,
- &buf,
- offset,
- rep_levels_byte_len as usize,
- );
- offset += bytes_read;
- self.rep_level_decoder = Some(rep_decoder);
- }
-
- // DataPage v2 only supports RLE encoding for definition
- // levels
- if self.descr.max_def_level() > 0 {
- let mut def_decoder =
- LevelDecoder::v2(self.descr.max_def_level());
- let bytes_read = def_decoder.set_data_range(
- self.num_buffered_values as usize,
- &buf,
- offset,
- def_levels_byte_len as usize,
- );
- offset += bytes_read;
- self.def_level_decoder = Some(def_decoder);
- }
-
- self.set_current_page_encoding(
- encoding,
- &buf,
- offset,
- num_values as usize,
- )?;
- return Ok(true);
- }
- };
- }
- }
- }
-
- Ok(true)
- }
-
- /// Resolves and updates encoding and set decoder for the current page
- fn set_current_page_encoding(
- &mut self,
- mut encoding: Encoding,
- buffer_ptr: &ByteBufferPtr,
- offset: usize,
- len: usize,
- ) -> Result<()> {
- if encoding == Encoding::PLAIN_DICTIONARY {
- encoding = Encoding::RLE_DICTIONARY;
- }
-
- let decoder = if encoding == Encoding::RLE_DICTIONARY {
- self.decoders
- .get_mut(&encoding)
- .expect("Decoder for dict should have been set")
- } else {
- // Search cache for data page decoder
- #[allow(clippy::map_entry)]
- if !self.decoders.contains_key(&encoding) {
- // Initialize decoder for this page
- let data_decoder = get_decoder::<T>(self.descr.clone(), encoding)?;
- self.decoders.insert(encoding, data_decoder);
- }
- self.decoders.get_mut(&encoding).unwrap()
- };
-
- decoder.set_data(buffer_ptr.start_from(offset), len as usize)?;
- self.current_encoding = Some(encoding);
- Ok(())
- }
-
- #[inline]
- fn has_next(&mut self) -> Result<bool> {
- if self.num_buffered_values == 0
- || self.num_buffered_values == self.num_decoded_values
- {
- // TODO: should we return false if read_new_page() = true and
- // num_buffered_values = 0?
- if !self.read_new_page()? {
- Ok(false)
- } else {
- Ok(self.num_buffered_values != 0)
- }
- } else {
- Ok(true)
- }
- }
-
- #[inline]
- fn read_rep_levels(&mut self, buffer: &mut [i16]) -> Result<usize> {
- let level_decoder = self
- .rep_level_decoder
- .as_mut()
- .expect("rep_level_decoder be set");
- level_decoder.get(buffer)
- }
-
- #[inline]
- fn read_def_levels(&mut self, buffer: &mut [i16]) -> Result<usize> {
- let level_decoder = self
- .def_level_decoder
- .as_mut()
- .expect("def_level_decoder be set");
- level_decoder.get(buffer)
- }
-
- #[inline]
- fn read_values(&mut self, buffer: &mut [T::T]) -> Result<usize> {
- let encoding = self
- .current_encoding
- .expect("current_encoding should be set");
- let current_decoder = self
- .decoders
- .get_mut(&encoding)
- .unwrap_or_else(|| panic!("decoder for encoding {} should be set", encoding));
- current_decoder.get(buffer)
- }
-
- #[inline]
- fn configure_dictionary(&mut self, page: Page) -> Result<bool> {
- let mut encoding = page.encoding();
- if encoding == Encoding::PLAIN || encoding == Encoding::PLAIN_DICTIONARY {
- encoding = Encoding::RLE_DICTIONARY
- }
-
- if self.decoders.contains_key(&encoding) {
- return Err(general_err!("Column cannot have more than one dictionary"));
- }
-
- if encoding == Encoding::RLE_DICTIONARY {
- let mut dictionary = PlainDecoder::<T>::new(self.descr.type_length());
- let num_values = page.num_values();
- dictionary.set_data(page.buffer().clone(), num_values as usize)?;
-
- let mut decoder = DictDecoder::new();
- decoder.set_dict(Box::new(dictionary))?;
- self.decoders.insert(encoding, Box::new(decoder));
- Ok(true)
- } else {
- Err(nyi_err!(
- "Invalid/Unsupported encoding type for dictionary: {}",
- encoding
- ))
- }
- }
-}
-
-#[cfg(test)]
-mod tests {
- use super::*;
-
- use rand::distributions::uniform::SampleUniform;
- use std::{collections::VecDeque, sync::Arc, vec::IntoIter};
-
- use crate::basic::Type as PhysicalType;
- use crate::column::page::Page;
- use crate::schema::types::{ColumnDescriptor, ColumnPath, Type as SchemaType};
- use crate::util::test_common::make_pages;
-
- const NUM_LEVELS: usize = 128;
- const NUM_PAGES: usize = 2;
- const MAX_DEF_LEVEL: i16 = 5;
- const MAX_REP_LEVEL: i16 = 5;
-
- // Macro to generate test cases
- macro_rules! test {
- // branch for generating i32 cases
- ($test_func:ident, i32, $func:ident, $def_level:expr, $rep_level:expr,
- $num_pages:expr, $num_levels:expr, $batch_size:expr, $min:expr, $max:expr) => {
- test_internal!(
- $test_func,
- Int32Type,
- get_test_int32_type,
- $func,
- $def_level,
- $rep_level,
- $num_pages,
- $num_levels,
- $batch_size,
- $min,
- $max
- );
- };
- // branch for generating i64 cases
- ($test_func:ident, i64, $func:ident, $def_level:expr, $rep_level:expr,
- $num_pages:expr, $num_levels:expr, $batch_size:expr, $min:expr, $max:expr) => {
- test_internal!(
- $test_func,
- Int64Type,
- get_test_int64_type,
- $func,
- $def_level,
- $rep_level,
- $num_pages,
- $num_levels,
- $batch_size,
- $min,
- $max
- );
- };
- }
-
- macro_rules! test_internal {
- ($test_func:ident, $ty:ident, $pty:ident, $func:ident, $def_level:expr,
- $rep_level:expr, $num_pages:expr, $num_levels:expr, $batch_size:expr,
- $min:expr, $max:expr) => {
- #[test]
- fn $test_func() {
- let desc = Arc::new(ColumnDescriptor::new(
- Arc::new($pty()),
- $def_level,
- $rep_level,
- ColumnPath::new(Vec::new()),
- ));
- let mut tester = ColumnReaderTester::<$ty>::new();
- tester.$func(desc, $num_pages, $num_levels, $batch_size, $min, $max);
- }
- };
- }
-
- test!(
- test_read_plain_v1_int32,
- i32,
- plain_v1,
- MAX_DEF_LEVEL,
- MAX_REP_LEVEL,
- NUM_PAGES,
- NUM_LEVELS,
- 16,
- std::i32::MIN,
- std::i32::MAX
- );
- test!(
- test_read_plain_v2_int32,
- i32,
- plain_v2,
- MAX_DEF_LEVEL,
- MAX_REP_LEVEL,
- NUM_PAGES,
- NUM_LEVELS,
- 16,
- std::i32::MIN,
- std::i32::MAX
- );
-
- test!(
- test_read_plain_v1_int32_uneven,
- i32,
- plain_v1,
- MAX_DEF_LEVEL,
- MAX_REP_LEVEL,
- NUM_PAGES,
- NUM_LEVELS,
- 17,
- std::i32::MIN,
- std::i32::MAX
- );
- test!(
- test_read_plain_v2_int32_uneven,
- i32,
- plain_v2,
- MAX_DEF_LEVEL,
- MAX_REP_LEVEL,
- NUM_PAGES,
- NUM_LEVELS,
- 17,
- std::i32::MIN,
- std::i32::MAX
- );
-
- test!(
- test_read_plain_v1_int32_multi_page,
- i32,
- plain_v1,
- MAX_DEF_LEVEL,
- MAX_REP_LEVEL,
- NUM_PAGES,
- NUM_LEVELS,
- 512,
- std::i32::MIN,
- std::i32::MAX
- );
- test!(
- test_read_plain_v2_int32_multi_page,
- i32,
- plain_v2,
- MAX_DEF_LEVEL,
- MAX_REP_LEVEL,
- NUM_PAGES,
- NUM_LEVELS,
- 512,
- std::i32::MIN,
- std::i32::MAX
- );
-
- // test cases when column descriptor has MAX_DEF_LEVEL = 0 and MAX_REP_LEVEL = 0
- test!(
- test_read_plain_v1_int32_required_non_repeated,
- i32,
- plain_v1,
- 0,
- 0,
- NUM_PAGES,
- NUM_LEVELS,
- 16,
- std::i32::MIN,
- std::i32::MAX
- );
- test!(
- test_read_plain_v2_int32_required_non_repeated,
- i32,
- plain_v2,
- 0,
- 0,
- NUM_PAGES,
- NUM_LEVELS,
- 16,
- std::i32::MIN,
- std::i32::MAX
- );
-
- test!(
- test_read_plain_v1_int64,
- i64,
- plain_v1,
- 1,
- 1,
- NUM_PAGES,
- NUM_LEVELS,
- 16,
- std::i64::MIN,
- std::i64::MAX
- );
- test!(
- test_read_plain_v2_int64,
- i64,
- plain_v2,
- 1,
- 1,
- NUM_PAGES,
- NUM_LEVELS,
- 16,
- std::i64::MIN,
- std::i64::MAX
- );
-
- test!(
- test_read_plain_v1_int64_uneven,
- i64,
- plain_v1,
- 1,
- 1,
- NUM_PAGES,
- NUM_LEVELS,
- 17,
- std::i64::MIN,
- std::i64::MAX
- );
- test!(
- test_read_plain_v2_int64_uneven,
- i64,
- plain_v2,
- 1,
- 1,
- NUM_PAGES,
- NUM_LEVELS,
- 17,
- std::i64::MIN,
- std::i64::MAX
- );
-
- test!(
- test_read_plain_v1_int64_multi_page,
- i64,
- plain_v1,
- 1,
- 1,
- NUM_PAGES,
- NUM_LEVELS,
- 512,
- std::i64::MIN,
- std::i64::MAX
- );
- test!(
- test_read_plain_v2_int64_multi_page,
- i64,
- plain_v2,
- 1,
- 1,
- NUM_PAGES,
- NUM_LEVELS,
- 512,
- std::i64::MIN,
- std::i64::MAX
- );
-
- // test cases when column descriptor has MAX_DEF_LEVEL = 0 and MAX_REP_LEVEL = 0
- test!(
- test_read_plain_v1_int64_required_non_repeated,
- i64,
- plain_v1,
- 0,
- 0,
- NUM_PAGES,
- NUM_LEVELS,
- 16,
- std::i64::MIN,
- std::i64::MAX
- );
- test!(
- test_read_plain_v2_int64_required_non_repeated,
- i64,
- plain_v2,
- 0,
- 0,
- NUM_PAGES,
- NUM_LEVELS,
- 16,
- std::i64::MIN,
- std::i64::MAX
- );
-
- test!(
- test_read_dict_v1_int32_small,
- i32,
- dict_v1,
- MAX_DEF_LEVEL,
- MAX_REP_LEVEL,
- 2,
- 2,
- 16,
- 0,
- 3
- );
- test!(
- test_read_dict_v2_int32_small,
- i32,
- dict_v2,
- MAX_DEF_LEVEL,
- MAX_REP_LEVEL,
- 2,
- 2,
- 16,
- 0,
- 3
- );
-
- test!(
- test_read_dict_v1_int32,
- i32,
- dict_v1,
- MAX_DEF_LEVEL,
- MAX_REP_LEVEL,
- NUM_PAGES,
- NUM_LEVELS,
- 16,
- 0,
- 3
- );
- test!(
- test_read_dict_v2_int32,
- i32,
- dict_v2,
- MAX_DEF_LEVEL,
- MAX_REP_LEVEL,
- NUM_PAGES,
- NUM_LEVELS,
- 16,
- 0,
- 3
- );
-
- test!(
- test_read_dict_v1_int32_uneven,
- i32,
- dict_v1,
- MAX_DEF_LEVEL,
- MAX_REP_LEVEL,
- NUM_PAGES,
- NUM_LEVELS,
- 17,
- 0,
- 3
- );
- test!(
- test_read_dict_v2_int32_uneven,
- i32,
- dict_v2,
- MAX_DEF_LEVEL,
- MAX_REP_LEVEL,
- NUM_PAGES,
- NUM_LEVELS,
- 17,
- 0,
- 3
- );
-
- test!(
- test_read_dict_v1_int32_multi_page,
- i32,
- dict_v1,
- MAX_DEF_LEVEL,
- MAX_REP_LEVEL,
- NUM_PAGES,
- NUM_LEVELS,
- 512,
- 0,
- 3
- );
- test!(
- test_read_dict_v2_int32_multi_page,
- i32,
- dict_v2,
- MAX_DEF_LEVEL,
- MAX_REP_LEVEL,
- NUM_PAGES,
- NUM_LEVELS,
- 512,
- 0,
- 3
- );
-
- test!(
- test_read_dict_v1_int64,
- i64,
- dict_v1,
- MAX_DEF_LEVEL,
- MAX_REP_LEVEL,
- NUM_PAGES,
- NUM_LEVELS,
- 16,
- 0,
- 3
- );
- test!(
- test_read_dict_v2_int64,
- i64,
- dict_v2,
- MAX_DEF_LEVEL,
- MAX_REP_LEVEL,
- NUM_PAGES,
- NUM_LEVELS,
- 16,
- 0,
- 3
- );
-
- #[test]
- fn test_read_batch_values_only() {
- test_read_batch_int32(16, &mut [0; 10], None, None); // < batch_size
- test_read_batch_int32(16, &mut [0; 16], None, None); // == batch_size
- test_read_batch_int32(16, &mut [0; 51], None, None); // > batch_size
- }
-
- #[test]
- fn test_read_batch_values_def_levels() {
- test_read_batch_int32(16, &mut [0; 10], Some(&mut [0; 10]), None);
- test_read_batch_int32(16, &mut [0; 16], Some(&mut [0; 16]), None);
- test_read_batch_int32(16, &mut [0; 51], Some(&mut [0; 51]), None);
- }
-
- #[test]
- fn test_read_batch_values_rep_levels() {
- test_read_batch_int32(16, &mut [0; 10], None, Some(&mut [0; 10]));
- test_read_batch_int32(16, &mut [0; 16], None, Some(&mut [0; 16]));
- test_read_batch_int32(16, &mut [0; 51], None, Some(&mut [0; 51]));
- }
-
- #[test]
- fn test_read_batch_different_buf_sizes() {
- test_read_batch_int32(16, &mut [0; 8], Some(&mut [0; 9]), Some(&mut [0; 7]));
- test_read_batch_int32(16, &mut [0; 1], Some(&mut [0; 9]), Some(&mut [0; 3]));
- }
-
- #[test]
- fn test_read_batch_values_def_rep_levels() {
- test_read_batch_int32(
- 128,
- &mut [0; 128],
- Some(&mut [0; 128]),
- Some(&mut [0; 128]),
- );
- }
-
- #[test]
- fn test_read_batch_adjust_after_buffering_page() {
- // This test covers scenario when buffering new page results in setting number
- // of decoded values to 0, resulting on reading `batch_size` of values, but it is
- // larger than we can insert into slice (affects values and levels).
- //
- // Note: values are chosen to reproduce the issue.
- //
- let primitive_type = get_test_int32_type();
- let desc = Arc::new(ColumnDescriptor::new(
- Arc::new(primitive_type),
- 1,
- 1,
- ColumnPath::new(Vec::new()),
- ));
-
- let num_pages = 2;
- let num_levels = 4;
- let batch_size = 5;
- let values = &mut vec![0; 7];
- let def_levels = &mut vec![0; 7];
- let rep_levels = &mut vec![0; 7];
-
- let mut tester = ColumnReaderTester::<Int32Type>::new();
- tester.test_read_batch(
- desc,
- Encoding::RLE_DICTIONARY,
- num_pages,
- num_levels,
- batch_size,
- std::i32::MIN,
- std::i32::MAX,
- values,
- Some(def_levels),
- Some(rep_levels),
- false,
- );
- }
-
- // ----------------------------------------------------------------------
- // Helper methods to make pages and test
- //
- // # Overview
- //
- // Most of the test functionality is implemented in `ColumnReaderTester`, which
- // provides some general data page test methods:
- // - `test_read_batch_general`
- // - `test_read_batch`
- //
- // There are also some high level wrappers that are part of `ColumnReaderTester`:
- // - `plain_v1` -> call `test_read_batch_general` with data page v1 and plain encoding
- // - `plain_v2` -> call `test_read_batch_general` with data page v2 and plain encoding
- // - `dict_v1` -> call `test_read_batch_general` with data page v1 + dictionary page
- // - `dict_v2` -> call `test_read_batch_general` with data page v2 + dictionary page
- //
- // And even higher level wrappers that simplify testing of almost the same test cases:
- // - `get_test_int32_type`, provides dummy schema type
- // - `get_test_int64_type`, provides dummy schema type
- // - `test_read_batch_int32`, wrapper for `read_batch` tests, since they are basically
- // the same, just different def/rep levels and batch size.
- //
- // # Page assembly
- //
- // Page construction and generation of values, definition and repetition levels
- // happens in `make_pages` function.
- // All values are randomly generated based on provided min/max, levels are calculated
- // based on provided max level for column descriptor (which is basically either int32
- // or int64 type in tests) and `levels_per_page` variable.
- //
- // We use `DataPageBuilder` and its implementation `DataPageBuilderImpl` to actually
- // turn values, definition and repetition levels into data pages (either v1 or v2).
- //
- // Those data pages are then stored as part of `TestPageReader` (we just pass vector
- // of generated pages directly), which implements `PageReader` interface.
- //
- // # Comparison
- //
- // This allows us to pass test page reader into column reader, so we can test
- // functionality of column reader - see `test_read_batch`, where we create column
- // reader -> typed column reader, buffer values in `read_batch` method and compare
- // output with generated data.
-
- // Returns dummy Parquet `Type` for primitive field, because most of our tests use
- // INT32 physical type.
- fn get_test_int32_type() -> SchemaType {
- SchemaType::primitive_type_builder("a", PhysicalType::INT32)
- .with_repetition(Repetition::REQUIRED)
- .with_converted_type(ConvertedType::INT_32)
- .with_length(-1)
- .build()
- .expect("build() should be OK")
- }
-
- // Returns dummy Parquet `Type` for INT64 physical type.
- fn get_test_int64_type() -> SchemaType {
- SchemaType::primitive_type_builder("a", PhysicalType::INT64)
- .with_repetition(Repetition::REQUIRED)
- .with_converted_type(ConvertedType::INT_64)
- .with_length(-1)
- .build()
- .expect("build() should be OK")
- }
-
- // Tests `read_batch()` functionality for INT32.
- //
- // This is a high level wrapper on `ColumnReaderTester` that allows us to specify some
- // boilerplate code for setting up definition/repetition levels and column descriptor.
- fn test_read_batch_int32(
- batch_size: usize,
- values: &mut [i32],
- def_levels: Option<&mut [i16]>,
- rep_levels: Option<&mut [i16]>,
- ) {
- let primitive_type = get_test_int32_type();
- // make field is required based on provided slices of levels
- let max_def_level = if def_levels.is_some() {
- MAX_DEF_LEVEL
- } else {
- 0
- };
- let max_rep_level = if def_levels.is_some() {
- MAX_REP_LEVEL
- } else {
- 0
- };
-
- let desc = Arc::new(ColumnDescriptor::new(
- Arc::new(primitive_type),
- max_def_level,
- max_rep_level,
- ColumnPath::new(Vec::new()),
- ));
- let mut tester = ColumnReaderTester::<Int32Type>::new();
- tester.test_read_batch(
- desc,
- Encoding::RLE_DICTIONARY,
- NUM_PAGES,
- NUM_LEVELS,
- batch_size,
- std::i32::MIN,
- std::i32::MAX,
- values,
- def_levels,
- rep_levels,
- false,
- );
- }
-
- struct ColumnReaderTester<T: DataType>
- where
- T::T: PartialOrd + SampleUniform + Copy,
- {
- rep_levels: Vec<i16>,
- def_levels: Vec<i16>,
- values: Vec<T::T>,
- }
-
- impl<T: DataType> ColumnReaderTester<T>
- where
- T::T: PartialOrd + SampleUniform + Copy,
- {
- pub fn new() -> Self {
- Self {
- rep_levels: Vec::new(),
- def_levels: Vec::new(),
- values: Vec::new(),
- }
- }
-
- // Method to generate and test data pages v1
- fn plain_v1(
- &mut self,
- desc: ColumnDescPtr,
- num_pages: usize,
- num_levels: usize,
- batch_size: usize,
- min: T::T,
- max: T::T,
- ) {
- self.test_read_batch_general(
- desc,
- Encoding::PLAIN,
- num_pages,
- num_levels,
- batch_size,
- min,
- max,
- false,
- );
- }
-
- // Method to generate and test data pages v2
- fn plain_v2(
- &mut self,
- desc: ColumnDescPtr,
- num_pages: usize,
- num_levels: usize,
- batch_size: usize,
- min: T::T,
- max: T::T,
- ) {
- self.test_read_batch_general(
- desc,
- Encoding::PLAIN,
- num_pages,
- num_levels,
- batch_size,
- min,
- max,
- true,
- );
- }
-
- // Method to generate and test dictionary page + data pages v1
- fn dict_v1(
- &mut self,
- desc: ColumnDescPtr,
- num_pages: usize,
- num_levels: usize,
- batch_size: usize,
- min: T::T,
- max: T::T,
- ) {
- self.test_read_batch_general(
- desc,
- Encoding::RLE_DICTIONARY,
- num_pages,
- num_levels,
- batch_size,
- min,
- max,
- false,
- );
- }
-
- // Method to generate and test dictionary page + data pages v2
- fn dict_v2(
- &mut self,
- desc: ColumnDescPtr,
- num_pages: usize,
- num_levels: usize,
- batch_size: usize,
- min: T::T,
- max: T::T,
- ) {
- self.test_read_batch_general(
- desc,
- Encoding::RLE_DICTIONARY,
- num_pages,
- num_levels,
- batch_size,
- min,
- max,
- true,
- );
- }
-
- // Helper function for the general case of `read_batch()` where `values`,
- // `def_levels` and `rep_levels` are always provided with enough space.
- fn test_read_batch_general(
- &mut self,
- desc: ColumnDescPtr,
- encoding: Encoding,
- num_pages: usize,
- num_levels: usize,
- batch_size: usize,
- min: T::T,
- max: T::T,
- use_v2: bool,
- ) {
- let mut def_levels = vec![0; num_levels * num_pages];
- let mut rep_levels = vec![0; num_levels * num_pages];
- let mut values = vec![T::T::default(); num_levels * num_pages];
- self.test_read_batch(
- desc,
- encoding,
- num_pages,
- num_levels,
- batch_size,
- min,
- max,
- &mut values,
- Some(&mut def_levels),
- Some(&mut rep_levels),
- use_v2,
- );
- }
-
- // Helper function to test `read_batch()` method with custom buffers for values,
- // definition and repetition levels.
- fn test_read_batch(
- &mut self,
- desc: ColumnDescPtr,
- encoding: Encoding,
- num_pages: usize,
- num_levels: usize,
- batch_size: usize,
- min: T::T,
- max: T::T,
- values: &mut [T::T],
- mut def_levels: Option<&mut [i16]>,
- mut rep_levels: Option<&mut [i16]>,
- use_v2: bool,
- ) {
- let mut pages = VecDeque::new();
- make_pages::<T>(
- desc.clone(),
- encoding,
- num_pages,
- num_levels,
- min,
- max,
- &mut self.def_levels,
- &mut self.rep_levels,
- &mut self.values,
- &mut pages,
- use_v2,
- );
- let max_def_level = desc.max_def_level();
- let page_reader = TestPageReader::new(Vec::from(pages));
- let column_reader: ColumnReader =
- get_column_reader(desc, Box::new(page_reader));
- let mut typed_column_reader = get_typed_column_reader::<T>(column_reader);
-
- let mut curr_values_read = 0;
- let mut curr_levels_read = 0;
- let mut done = false;
- while !done {
- let actual_def_levels =
- def_levels.as_mut().map(|vec| &mut vec[curr_levels_read..]);
- let actual_rep_levels =
- rep_levels.as_mut().map(|vec| &mut vec[curr_levels_read..]);
-
- let (values_read, levels_read) = typed_column_reader
- .read_batch(
- batch_size,
- actual_def_levels,
- actual_rep_levels,
- &mut values[curr_values_read..],
- )
- .expect("read_batch() should be OK");
-
- if values_read == 0 && levels_read == 0 {
- done = true;
- }
-
- curr_values_read += values_read;
- curr_levels_read += levels_read;
- }
-
- assert!(
- values.len() >= curr_values_read,
- "values.len() >= values_read"
- );
- assert_eq!(
- &values[0..curr_values_read],
- &self.values[0..curr_values_read],
- "values content doesn't match"
- );
-
- if let Some(ref levels) = def_levels {
- assert!(
- levels.len() >= curr_levels_read,
- "def_levels.len() >= levels_read"
- );
- assert_eq!(
- &levels[0..curr_levels_read],
- &self.def_levels[0..curr_levels_read],
- "definition levels content doesn't match"
- );
- }
-
- if let Some(ref levels) = rep_levels {
- assert!(
- levels.len() >= curr_levels_read,
- "rep_levels.len() >= levels_read"
- );
- assert_eq!(
- &levels[0..curr_levels_read],
- &self.rep_levels[0..curr_levels_read],
- "repetition levels content doesn't match"
- );
- }
-
- if def_levels.is_none() && rep_levels.is_none() {
- assert!(
- curr_levels_read == 0,
- "expected to read 0 levels, found {}",
- curr_levels_read
- );
- } else if def_levels.is_some() && max_def_level > 0 {
- assert!(
- curr_levels_read >= curr_values_read,
- "expected levels read to be greater than values read"
- );
- }
- }
- }
-
- struct TestPageReader {
- pages: IntoIter<Page>,
- }
-
- impl TestPageReader {
- pub fn new(pages: Vec<Page>) -> Self {
- Self {
- pages: pages.into_iter(),
- }
- }
- }
-
- impl PageReader for TestPageReader {
- fn get_next_page(&mut self) -> Result<Option<Page>> {
- Ok(self.pages.next())
- }
- }
-
- impl Iterator for TestPageReader {
- type Item = Result<Page>;
-
- fn next(&mut self) -> Option<Self::Item> {
- self.get_next_page().transpose()
- }
- }
-}
diff --git a/parquet/src/column/writer.rs b/parquet/src/column/writer.rs
deleted file mode 100644
index 910a9ed..0000000
--- a/parquet/src/column/writer.rs
+++ /dev/null
@@ -1,2020 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-//! Contains column writer API.
-use std::{cmp, collections::VecDeque, convert::TryFrom, marker::PhantomData, sync::Arc};
-
-use crate::basic::{Compression, Encoding, LogicalType, PageType, Type};
-use crate::column::page::{CompressedPage, Page, PageWriteSpec, PageWriter};
-use crate::compression::{create_codec, Codec};
-use crate::data_type::private::ParquetValueType;
-use crate::data_type::AsBytes;
-use crate::data_type::*;
-use crate::encodings::{
- encoding::{get_encoder, DictEncoder, Encoder},
- levels::{max_buffer_size, LevelEncoder},
-};
-use crate::errors::{ParquetError, Result};
-use crate::file::statistics::Statistics;
-use crate::file::{
- metadata::ColumnChunkMetaData,
- properties::{WriterProperties, WriterPropertiesPtr, WriterVersion},
-};
-use crate::schema::types::ColumnDescPtr;
-use crate::util::bit_util::FromBytes;
-use crate::util::memory::{ByteBufferPtr, MemTracker};
-
-/// Column writer for a Parquet type.
-pub enum ColumnWriter {
- BoolColumnWriter(ColumnWriterImpl<BoolType>),
- Int32ColumnWriter(ColumnWriterImpl<Int32Type>),
- Int64ColumnWriter(ColumnWriterImpl<Int64Type>),
- Int96ColumnWriter(ColumnWriterImpl<Int96Type>),
- FloatColumnWriter(ColumnWriterImpl<FloatType>),
- DoubleColumnWriter(ColumnWriterImpl<DoubleType>),
- ByteArrayColumnWriter(ColumnWriterImpl<ByteArrayType>),
- FixedLenByteArrayColumnWriter(ColumnWriterImpl<FixedLenByteArrayType>),
-}
-
-pub enum Level {
- Page,
- Column,
-}
-
-macro_rules! gen_stats_section {
- ($physical_ty: ty, $stat_fn: ident, $min: ident, $max: ident, $distinct: ident, $nulls: ident) => {{
- let min = $min.as_ref().and_then(|v| {
- Some(read_num_bytes!(
- $physical_ty,
- v.as_bytes().len(),
- &v.as_bytes()
- ))
- });
- let max = $max.as_ref().and_then(|v| {
- Some(read_num_bytes!(
- $physical_ty,
- v.as_bytes().len(),
- &v.as_bytes()
- ))
- });
- Statistics::$stat_fn(min, max, $distinct, $nulls, false)
- }};
-}
-
-/// Gets a specific column writer corresponding to column descriptor `descr`.
-pub fn get_column_writer(
- descr: ColumnDescPtr,
- props: WriterPropertiesPtr,
- page_writer: Box<dyn PageWriter>,
-) -> ColumnWriter {
- match descr.physical_type() {
- Type::BOOLEAN => ColumnWriter::BoolColumnWriter(ColumnWriterImpl::new(
- descr,
- props,
- page_writer,
- )),
- Type::INT32 => ColumnWriter::Int32ColumnWriter(ColumnWriterImpl::new(
- descr,
- props,
- page_writer,
- )),
- Type::INT64 => ColumnWriter::Int64ColumnWriter(ColumnWriterImpl::new(
- descr,
- props,
- page_writer,
- )),
- Type::INT96 => ColumnWriter::Int96ColumnWriter(ColumnWriterImpl::new(
- descr,
- props,
- page_writer,
- )),
- Type::FLOAT => ColumnWriter::FloatColumnWriter(ColumnWriterImpl::new(
- descr,
- props,
- page_writer,
- )),
- Type::DOUBLE => ColumnWriter::DoubleColumnWriter(ColumnWriterImpl::new(
- descr,
- props,
- page_writer,
- )),
- Type::BYTE_ARRAY => ColumnWriter::ByteArrayColumnWriter(ColumnWriterImpl::new(
- descr,
- props,
- page_writer,
- )),
- Type::FIXED_LEN_BYTE_ARRAY => ColumnWriter::FixedLenByteArrayColumnWriter(
- ColumnWriterImpl::new(descr, props, page_writer),
- ),
- }
-}
-
-/// Gets a typed column writer for the specific type `T`, by "up-casting" `col_writer` of
-/// non-generic type to a generic column writer type `ColumnWriterImpl`.
-///
-/// Panics if actual enum value for `col_writer` does not match the type `T`.
-pub fn get_typed_column_writer<T: DataType>(
- col_writer: ColumnWriter,
-) -> ColumnWriterImpl<T> {
- T::get_column_writer(col_writer).unwrap_or_else(|| {
- panic!(
- "Failed to convert column writer into a typed column writer for `{}` type",
- T::get_physical_type()
- )
- })
-}
-
-/// Similar to `get_typed_column_writer` but returns a reference.
-pub fn get_typed_column_writer_ref<T: DataType>(
- col_writer: &ColumnWriter,
-) -> &ColumnWriterImpl<T> {
- T::get_column_writer_ref(col_writer).unwrap_or_else(|| {
- panic!(
- "Failed to convert column writer into a typed column writer for `{}` type",
- T::get_physical_type()
- )
- })
-}
-
-/// Similar to `get_typed_column_writer` but returns a reference.
-pub fn get_typed_column_writer_mut<T: DataType>(
- col_writer: &mut ColumnWriter,
-) -> &mut ColumnWriterImpl<T> {
- T::get_column_writer_mut(col_writer).unwrap_or_else(|| {
- panic!(
- "Failed to convert column writer into a typed column writer for `{}` type",
- T::get_physical_type()
- )
- })
-}
-
-/// Typed column writer for a primitive column.
-pub struct ColumnWriterImpl<T: DataType> {
- // Column writer properties
- descr: ColumnDescPtr,
- props: WriterPropertiesPtr,
- page_writer: Box<dyn PageWriter>,
- has_dictionary: bool,
- dict_encoder: Option<DictEncoder<T>>,
- encoder: Box<dyn Encoder<T>>,
- codec: Compression,
- compressor: Option<Box<dyn Codec>>,
- // Metrics per page
- num_buffered_values: u32,
- num_buffered_encoded_values: u32,
- num_buffered_rows: u32,
- min_page_value: Option<T::T>,
- max_page_value: Option<T::T>,
- num_page_nulls: u64,
- page_distinct_count: Option<u64>,
- // Metrics per column writer
- total_bytes_written: u64,
- total_rows_written: u64,
- total_uncompressed_size: u64,
- total_compressed_size: u64,
- total_num_values: u64,
- dictionary_page_offset: Option<u64>,
- data_page_offset: Option<u64>,
- min_column_value: Option<T::T>,
- max_column_value: Option<T::T>,
- num_column_nulls: u64,
- column_distinct_count: Option<u64>,
- // Reused buffers
- def_levels_sink: Vec<i16>,
- rep_levels_sink: Vec<i16>,
- data_pages: VecDeque<CompressedPage>,
- _phantom: PhantomData<T>,
-}
-
-impl<T: DataType> ColumnWriterImpl<T> {
- pub fn new(
- descr: ColumnDescPtr,
- props: WriterPropertiesPtr,
- page_writer: Box<dyn PageWriter>,
- ) -> Self {
- let codec = props.compression(descr.path());
- let compressor = create_codec(codec).unwrap();
-
- // Optionally set dictionary encoder.
- let dict_encoder = if props.dictionary_enabled(descr.path())
- && has_dictionary_support(T::get_physical_type(), &props)
- {
- Some(DictEncoder::new(descr.clone(), Arc::new(MemTracker::new())))
- } else {
- None
- };
-
- // Whether or not this column writer has a dictionary encoding.
- let has_dictionary = dict_encoder.is_some();
-
- // Set either main encoder or fallback encoder.
- let fallback_encoder = get_encoder(
- descr.clone(),
- props
- .encoding(descr.path())
- .unwrap_or_else(|| fallback_encoding(T::get_physical_type(), &props)),
- Arc::new(MemTracker::new()),
- )
- .unwrap();
-
- Self {
- descr,
- props,
- page_writer,
- has_dictionary,
- dict_encoder,
- encoder: fallback_encoder,
- codec,
- compressor,
- num_buffered_values: 0,
- num_buffered_encoded_values: 0,
- num_buffered_rows: 0,
- total_bytes_written: 0,
- total_rows_written: 0,
- total_uncompressed_size: 0,
- total_compressed_size: 0,
- total_num_values: 0,
- dictionary_page_offset: None,
- data_page_offset: None,
- def_levels_sink: vec![],
- rep_levels_sink: vec![],
- data_pages: VecDeque::new(),
- min_page_value: None,
- max_page_value: None,
- num_page_nulls: 0,
- page_distinct_count: None,
- min_column_value: None,
- max_column_value: None,
- num_column_nulls: 0,
- column_distinct_count: None,
- _phantom: PhantomData,
- }
- }
-
- fn write_batch_internal(
- &mut self,
- values: &[T::T],
- def_levels: Option<&[i16]>,
- rep_levels: Option<&[i16]>,
- min: &Option<T::T>,
- max: &Option<T::T>,
- null_count: Option<u64>,
- distinct_count: Option<u64>,
- ) -> Result<usize> {
- // We check for DataPage limits only after we have inserted the values. If a user
- // writes a large number of values, the DataPage size can be well above the limit.
- //
- // The purpose of this chunking is to bound this. Even if a user writes large
- // number of values, the chunking will ensure that we add data page at a
- // reasonable pagesize limit.
-
- // TODO: find out why we don't account for size of levels when we estimate page
- // size.
-
- // Find out the minimal length to prevent index out of bound errors.
- let mut min_len = values.len();
- if let Some(levels) = def_levels {
- min_len = cmp::min(min_len, levels.len());
- }
- if let Some(levels) = rep_levels {
- min_len = cmp::min(min_len, levels.len());
- }
-
- // Find out number of batches to process.
- let write_batch_size = self.props.write_batch_size();
- let num_batches = min_len / write_batch_size;
-
- // Process pre-calculated statistics
- match (min, max) {
- (Some(min), Some(max)) => {
- if self
- .min_column_value
- .as_ref()
- .map_or(true, |v| self.compare_greater(v, min))
- {
- self.min_column_value = Some(min.clone());
- }
- if self
- .max_column_value
- .as_ref()
- .map_or(true, |v| self.compare_greater(max, v))
- {
- self.max_column_value = Some(max.clone());
- }
- }
- (None, Some(_)) | (Some(_), None) => {
- panic!("min/max should be both set or both None")
- }
- (None, None) => {}
- }
-
- if let Some(distinct) = distinct_count {
- self.column_distinct_count =
- Some(self.column_distinct_count.unwrap_or(0) + distinct);
- }
-
- if let Some(nulls) = null_count {
- self.num_column_nulls += nulls;
- }
-
- let calculate_page_stats = (min.is_none() || max.is_none())
- && null_count.is_none()
- && distinct_count.is_none();
-
- let mut values_offset = 0;
- let mut levels_offset = 0;
- for _ in 0..num_batches {
- values_offset += self.write_mini_batch(
- &values[values_offset..values_offset + write_batch_size],
- def_levels.map(|lv| &lv[levels_offset..levels_offset + write_batch_size]),
- rep_levels.map(|lv| &lv[levels_offset..levels_offset + write_batch_size]),
- calculate_page_stats,
- )?;
- levels_offset += write_batch_size;
- }
-
- values_offset += self.write_mini_batch(
- &values[values_offset..],
- def_levels.map(|lv| &lv[levels_offset..]),
- rep_levels.map(|lv| &lv[levels_offset..]),
- calculate_page_stats,
- )?;
-
- // Return total number of values processed.
- Ok(values_offset)
- }
-
- /// Writes batch of values, definition levels and repetition levels.
- /// Returns number of values processed (written).
- ///
- /// If definition and repetition levels are provided, we write fully those levels and
- /// select how many values to write (this number will be returned), since number of
- /// actual written values may be smaller than provided values.
- ///
- /// If only values are provided, then all values are written and the length of
- /// of the values buffer is returned.
- ///
- /// Definition and/or repetition levels can be omitted, if values are
- /// non-nullable and/or non-repeated.
- pub fn write_batch(
- &mut self,
- values: &[T::T],
- def_levels: Option<&[i16]>,
- rep_levels: Option<&[i16]>,
- ) -> Result<usize> {
- self.write_batch_internal(
- values, def_levels, rep_levels, &None, &None, None, None,
- )
- }
-
- /// Writer may optionally provide pre-calculated statistics for this batch, in which case we do
- /// not calculate page level statistics as this will defeat the purpose of speeding up the write
- /// process with pre-calculated statistics.
- pub fn write_batch_with_statistics(
- &mut self,
- values: &[T::T],
- def_levels: Option<&[i16]>,
- rep_levels: Option<&[i16]>,
- min: &Option<T::T>,
- max: &Option<T::T>,
- nulls_count: Option<u64>,
- distinct_count: Option<u64>,
- ) -> Result<usize> {
- self.write_batch_internal(
- values,
- def_levels,
- rep_levels,
- min,
- max,
- nulls_count,
- distinct_count,
- )
- }
-
- /// Returns total number of bytes written by this column writer so far.
- /// This value is also returned when column writer is closed.
- pub fn get_total_bytes_written(&self) -> u64 {
- self.total_bytes_written
- }
-
- /// Returns total number of rows written by this column writer so far.
- /// This value is also returned when column writer is closed.
- pub fn get_total_rows_written(&self) -> u64 {
- self.total_rows_written
- }
-
- /// Finalises writes and closes the column writer.
- /// Returns total bytes written, total rows written and column chunk metadata.
- pub fn close(mut self) -> Result<(u64, u64, ColumnChunkMetaData)> {
- if self.dict_encoder.is_some() {
- self.write_dictionary_page()?;
- }
- self.flush_data_pages()?;
- let metadata = self.write_column_metadata()?;
- self.dict_encoder = None;
- self.page_writer.close()?;
-
- Ok((self.total_bytes_written, self.total_rows_written, metadata))
- }
-
- /// Writes mini batch of values, definition and repetition levels.
- /// This allows fine-grained processing of values and maintaining a reasonable
- /// page size.
- fn write_mini_batch(
- &mut self,
- values: &[T::T],
- def_levels: Option<&[i16]>,
- rep_levels: Option<&[i16]>,
- calculate_page_stats: bool,
- ) -> Result<usize> {
- let mut values_to_write = 0;
-
- // Check if number of definition levels is the same as number of repetition
- // levels.
- if let (Some(def), Some(rep)) = (def_levels, rep_levels) {
- if def.len() != rep.len() {
- return Err(general_err!(
- "Inconsistent length of definition and repetition levels: {} != {}",
- def.len(),
- rep.len()
- ));
- }
- }
-
- // Process definition levels and determine how many values to write.
- let num_values = if self.descr.max_def_level() > 0 {
- let levels = def_levels.ok_or_else(|| {
- general_err!(
- "Definition levels are required, because max definition level = {}",
- self.descr.max_def_level()
- )
- })?;
-
- for &level in levels {
- if level == self.descr.max_def_level() {
- values_to_write += 1;
- } else if calculate_page_stats {
- self.num_page_nulls += 1
- }
- }
-
- self.write_definition_levels(levels);
- u32::try_from(levels.len()).unwrap()
- } else {
- values_to_write = values.len();
- u32::try_from(values_to_write).unwrap()
- };
-
- // Process repetition levels and determine how many rows we are about to process.
- if self.descr.max_rep_level() > 0 {
- // A row could contain more than one value.
- let levels = rep_levels.ok_or_else(|| {
- general_err!(
- "Repetition levels are required, because max repetition level = {}",
- self.descr.max_rep_level()
- )
- })?;
-
- // Count the occasions where we start a new row
- for &level in levels {
- self.num_buffered_rows += (level == 0) as u32
- }
-
- self.write_repetition_levels(levels);
- } else {
- // Each value is exactly one row.
- // Equals to the number of values, we count nulls as well.
- self.num_buffered_rows += num_values;
- }
-
- // Check that we have enough values to write.
- let values_to_write = values.get(0..values_to_write).ok_or_else(|| {
- general_err!(
- "Expected to write {} values, but have only {}",
- values_to_write,
- values.len()
- )
- })?;
-
- if calculate_page_stats {
- for val in values_to_write {
- self.update_page_min_max(val);
- }
- }
-
- self.write_values(values_to_write)?;
-
- self.num_buffered_values += num_values;
- self.num_buffered_encoded_values += u32::try_from(values_to_write.len()).unwrap();
-
- if self.should_add_data_page() {
- self.add_data_page(calculate_page_stats)?;
- }
-
- if self.should_dict_fallback() {
- self.dict_fallback()?;
- }
-
- Ok(values_to_write.len())
- }
-
- #[inline]
- fn write_definition_levels(&mut self, def_levels: &[i16]) {
- self.def_levels_sink.extend_from_slice(def_levels);
- }
-
- #[inline]
- fn write_repetition_levels(&mut self, rep_levels: &[i16]) {
- self.rep_levels_sink.extend_from_slice(rep_levels);
- }
-
- #[inline]
- fn write_values(&mut self, values: &[T::T]) -> Result<()> {
- match self.dict_encoder {
- Some(ref mut encoder) => encoder.put(values),
- None => self.encoder.put(values),
- }
- }
-
- /// Returns true if we need to fall back to non-dictionary encoding.
- ///
- /// We can only fall back if dictionary encoder is set and we have exceeded dictionary
- /// size.
- #[inline]
- fn should_dict_fallback(&self) -> bool {
- match self.dict_encoder {
- Some(ref encoder) => {
- encoder.dict_encoded_size() >= self.props.dictionary_pagesize_limit()
- }
- None => false,
- }
- }
-
- /// Returns true if there is enough data for a data page, false otherwise.
- #[inline]
- fn should_add_data_page(&self) -> bool {
- match self.dict_encoder {
- Some(ref encoder) => {
- encoder.estimated_data_encoded_size() >= self.props.data_pagesize_limit()
- }
- None => {
- self.encoder.estimated_data_encoded_size()
- >= self.props.data_pagesize_limit()
- }
- }
- }
-
- /// Performs dictionary fallback.
- /// Prepares and writes dictionary and all data pages into page writer.
- fn dict_fallback(&mut self) -> Result<()> {
- // At this point we know that we need to fall back.
- self.write_dictionary_page()?;
- self.flush_data_pages()?;
- self.dict_encoder = None;
- Ok(())
- }
-
- /// Adds data page.
- /// Data page is either buffered in case of dictionary encoding or written directly.
- fn add_data_page(&mut self, calculate_page_stat: bool) -> Result<()> {
- // Extract encoded values
- let value_bytes = match self.dict_encoder {
- Some(ref mut encoder) => encoder.write_indices()?,
- None => self.encoder.flush_buffer()?,
- };
-
- // Select encoding based on current encoder and writer version (v1 or v2).
- let encoding = if self.dict_encoder.is_some() {
- self.props.dictionary_data_page_encoding()
- } else {
- self.encoder.encoding()
- };
-
- let max_def_level = self.descr.max_def_level();
- let max_rep_level = self.descr.max_rep_level();
-
- // always update column NULL count, no matter if page stats are used
- self.num_column_nulls += self.num_page_nulls;
-
- let page_statistics = if calculate_page_stat {
- self.update_column_min_max();
- Some(self.make_page_statistics())
- } else {
- None
- };
-
- let compressed_page = match self.props.writer_version() {
- WriterVersion::PARQUET_1_0 => {
- let mut buffer = vec![];
-
- if max_rep_level > 0 {
- buffer.extend_from_slice(
- &self.encode_levels_v1(
- Encoding::RLE,
- &self.rep_levels_sink[..],
- max_rep_level,
- )?[..],
- );
- }
-
- if max_def_level > 0 {
- buffer.extend_from_slice(
- &self.encode_levels_v1(
- Encoding::RLE,
- &self.def_levels_sink[..],
- max_def_level,
- )?[..],
- );
- }
-
- buffer.extend_from_slice(value_bytes.data());
- let uncompressed_size = buffer.len();
-
- if let Some(ref mut cmpr) = self.compressor {
- let mut compressed_buf = Vec::with_capacity(value_bytes.data().len());
- cmpr.compress(&buffer[..], &mut compressed_buf)?;
- buffer = compressed_buf;
- }
-
- let data_page = Page::DataPage {
- buf: ByteBufferPtr::new(buffer),
- num_values: self.num_buffered_values,
- encoding,
- def_level_encoding: Encoding::RLE,
- rep_level_encoding: Encoding::RLE,
- statistics: page_statistics,
- };
-
- CompressedPage::new(data_page, uncompressed_size)
- }
- WriterVersion::PARQUET_2_0 => {
- let mut rep_levels_byte_len = 0;
- let mut def_levels_byte_len = 0;
- let mut buffer = vec![];
-
- if max_rep_level > 0 {
- let levels =
- self.encode_levels_v2(&self.rep_levels_sink[..], max_rep_level)?;
- rep_levels_byte_len = levels.len();
- buffer.extend_from_slice(&levels[..]);
- }
-
- if max_def_level > 0 {
- let levels =
- self.encode_levels_v2(&self.def_levels_sink[..], max_def_level)?;
- def_levels_byte_len = levels.len();
- buffer.extend_from_slice(&levels[..]);
- }
-
- let uncompressed_size =
- rep_levels_byte_len + def_levels_byte_len + value_bytes.len();
-
- // Data Page v2 compresses values only.
- match self.compressor {
- Some(ref mut cmpr) => {
- cmpr.compress(value_bytes.data(), &mut buffer)?;
- }
- None => buffer.extend_from_slice(value_bytes.data()),
- }
-
- let data_page = Page::DataPageV2 {
- buf: ByteBufferPtr::new(buffer),
- num_values: self.num_buffered_values,
- encoding,
- num_nulls: self.num_buffered_values
- - self.num_buffered_encoded_values,
- num_rows: self.num_buffered_rows,
- def_levels_byte_len: def_levels_byte_len as u32,
- rep_levels_byte_len: rep_levels_byte_len as u32,
- is_compressed: self.compressor.is_some(),
- statistics: page_statistics,
- };
-
- CompressedPage::new(data_page, uncompressed_size)
- }
- };
-
- // Check if we need to buffer data page or flush it to the sink directly.
- if self.dict_encoder.is_some() {
- self.data_pages.push_back(compressed_page);
- } else {
- self.write_data_page(compressed_page)?;
- }
-
- // Update total number of rows.
- self.total_rows_written += self.num_buffered_rows as u64;
-
- // Reset state.
- self.rep_levels_sink.clear();
- self.def_levels_sink.clear();
- self.num_buffered_values = 0;
- self.num_buffered_encoded_values = 0;
- self.num_buffered_rows = 0;
- self.min_page_value = None;
- self.max_page_value = None;
- self.num_page_nulls = 0;
- self.page_distinct_count = None;
-
- Ok(())
- }
-
- /// Finalises any outstanding data pages and flushes buffered data pages from
- /// dictionary encoding into underlying sink.
- #[inline]
- fn flush_data_pages(&mut self) -> Result<()> {
- // Write all outstanding data to a new page.
- let calculate_page_stats =
- self.min_page_value.is_some() && self.max_page_value.is_some();
- if self.num_buffered_values > 0 {
- self.add_data_page(calculate_page_stats)?;
- }
-
- while let Some(page) = self.data_pages.pop_front() {
- self.write_data_page(page)?;
- }
-
- Ok(())
- }
-
- /// Assembles and writes column chunk metadata.
- fn write_column_metadata(&mut self) -> Result<ColumnChunkMetaData> {
- let total_compressed_size = self.total_compressed_size as i64;
- let total_uncompressed_size = self.total_uncompressed_size as i64;
- let num_values = self.total_num_values as i64;
- let dict_page_offset = self.dictionary_page_offset.map(|v| v as i64);
- // If data page offset is not set, then no pages have been written
- let data_page_offset = self.data_page_offset.unwrap_or(0) as i64;
-
- let file_offset;
- let mut encodings = Vec::new();
-
- if self.has_dictionary {
- assert!(dict_page_offset.is_some(), "Dictionary offset is not set");
- file_offset = dict_page_offset.unwrap() + total_compressed_size;
- // NOTE: This should be in sync with writing dictionary pages.
- encodings.push(self.props.dictionary_page_encoding());
- encodings.push(self.props.dictionary_data_page_encoding());
- // Fallback to alternative encoding, add it to the list.
- if self.dict_encoder.is_none() {
- encodings.push(self.encoder.encoding());
- }
- } else {
- file_offset = data_page_offset + total_compressed_size;
- encodings.push(self.encoder.encoding());
- }
- // We use only RLE level encoding for data page v1 and data page v2.
- encodings.push(Encoding::RLE);
-
- let statistics = self.make_column_statistics();
- let metadata = ColumnChunkMetaData::builder(self.descr.clone())
- .set_compression(self.codec)
- .set_encodings(encodings)
- .set_file_offset(file_offset)
- .set_total_compressed_size(total_compressed_size)
- .set_total_uncompressed_size(total_uncompressed_size)
- .set_num_values(num_values)
- .set_data_page_offset(data_page_offset)
- .set_dictionary_page_offset(dict_page_offset)
- .set_statistics(statistics)
- .build()?;
-
- self.page_writer.write_metadata(&metadata)?;
-
- Ok(metadata)
- }
-
- /// Encodes definition or repetition levels for Data Page v1.
- #[inline]
- fn encode_levels_v1(
- &self,
- encoding: Encoding,
- levels: &[i16],
- max_level: i16,
- ) -> Result<Vec<u8>> {
- let size = max_buffer_size(encoding, max_level, levels.len());
- let mut encoder = LevelEncoder::v1(encoding, max_level, vec![0; size]);
- encoder.put(&levels)?;
- encoder.consume()
- }
-
- /// Encodes definition or repetition levels for Data Page v2.
- /// Encoding is always RLE.
- #[inline]
- fn encode_levels_v2(&self, levels: &[i16], max_level: i16) -> Result<Vec<u8>> {
- let size = max_buffer_size(Encoding::RLE, max_level, levels.len());
- let mut encoder = LevelEncoder::v2(max_level, vec![0; size]);
- encoder.put(&levels)?;
- encoder.consume()
- }
-
- /// Writes compressed data page into underlying sink and updates global metrics.
- #[inline]
- fn write_data_page(&mut self, page: CompressedPage) -> Result<()> {
- let page_spec = self.page_writer.write_page(page)?;
- self.update_metrics_for_page(page_spec);
- Ok(())
- }
-
- /// Writes dictionary page into underlying sink.
- #[inline]
- fn write_dictionary_page(&mut self) -> Result<()> {
- let compressed_page = {
- let encoder = self
- .dict_encoder
- .as_ref()
- .ok_or_else(|| general_err!("Dictionary encoder is not set"))?;
-
- let is_sorted = encoder.is_sorted();
- let num_values = encoder.num_entries();
- let mut values_buf = encoder.write_dict()?;
- let uncompressed_size = values_buf.len();
-
- if let Some(ref mut cmpr) = self.compressor {
- let mut output_buf = Vec::with_capacity(uncompressed_size);
- cmpr.compress(values_buf.data(), &mut output_buf)?;
- values_buf = ByteBufferPtr::new(output_buf);
- }
-
- let dict_page = Page::DictionaryPage {
- buf: values_buf,
- num_values: num_values as u32,
- encoding: self.props.dictionary_page_encoding(),
- is_sorted,
- };
- CompressedPage::new(dict_page, uncompressed_size)
- };
-
- let page_spec = self.page_writer.write_page(compressed_page)?;
- self.update_metrics_for_page(page_spec);
- Ok(())
- }
-
- /// Updates column writer metrics with each page metadata.
- #[inline]
- fn update_metrics_for_page(&mut self, page_spec: PageWriteSpec) {
- self.total_uncompressed_size += page_spec.uncompressed_size as u64;
- self.total_compressed_size += page_spec.compressed_size as u64;
- self.total_num_values += page_spec.num_values as u64;
- self.total_bytes_written += page_spec.bytes_written;
-
- match page_spec.page_type {
- PageType::DATA_PAGE | PageType::DATA_PAGE_V2 => {
- if self.data_page_offset.is_none() {
- self.data_page_offset = Some(page_spec.offset);
- }
- }
- PageType::DICTIONARY_PAGE => {
- assert!(
- self.dictionary_page_offset.is_none(),
- "Dictionary offset is already set"
- );
- self.dictionary_page_offset = Some(page_spec.offset);
- }
- _ => {}
- }
- }
-
- /// Returns reference to the underlying page writer.
- /// This method is intended to use in tests only.
- fn get_page_writer_ref(&self) -> &dyn PageWriter {
- self.page_writer.as_ref()
- }
-
- fn make_column_statistics(&self) -> Statistics {
- self.make_typed_statistics(Level::Column)
- }
-
- fn make_page_statistics(&self) -> Statistics {
- self.make_typed_statistics(Level::Page)
- }
-
- pub fn make_typed_statistics(&self, level: Level) -> Statistics {
- let (min, max, distinct, nulls) = match level {
- Level::Page => (
- self.min_page_value.as_ref(),
- self.max_page_value.as_ref(),
- self.page_distinct_count,
- self.num_page_nulls,
- ),
- Level::Column => (
- self.min_column_value.as_ref(),
- self.max_column_value.as_ref(),
- self.column_distinct_count,
- self.num_column_nulls,
- ),
- };
- match self.descr.physical_type() {
- Type::INT32 => gen_stats_section!(i32, int32, min, max, distinct, nulls),
- Type::BOOLEAN => gen_stats_section!(i32, int32, min, max, distinct, nulls),
- Type::INT64 => gen_stats_section!(i64, int64, min, max, distinct, nulls),
- Type::INT96 => gen_stats_section!(Int96, int96, min, max, distinct, nulls),
- Type::FLOAT => gen_stats_section!(f32, float, min, max, distinct, nulls),
- Type::DOUBLE => gen_stats_section!(f64, double, min, max, distinct, nulls),
- Type::BYTE_ARRAY | Type::FIXED_LEN_BYTE_ARRAY => {
- let min = min.as_ref().map(|v| ByteArray::from(v.as_bytes().to_vec()));
- let max = max.as_ref().map(|v| ByteArray::from(v.as_bytes().to_vec()));
- Statistics::byte_array(min, max, distinct, nulls, false)
- }
- }
- }
-
- #[allow(clippy::eq_op)]
- fn update_page_min_max(&mut self, val: &T::T) {
- // simple "isNaN" check that works for all types
- if val == val {
- if self
- .min_page_value
- .as_ref()
- .map_or(true, |min| self.compare_greater(min, val))
- {
- self.min_page_value = Some(val.clone());
- }
- if self
- .max_page_value
- .as_ref()
- .map_or(true, |max| self.compare_greater(val, max))
- {
- self.max_page_value = Some(val.clone());
- }
- }
- }
-
- fn update_column_min_max(&mut self) {
- let update_min = self.min_column_value.as_ref().map_or(true, |min| {
- let page_value = self.min_page_value.as_ref().unwrap();
- self.compare_greater(min, page_value)
- });
- if update_min {
- self.min_column_value = self.min_page_value.clone();
- }
-
- let update_max = self.max_column_value.as_ref().map_or(true, |max| {
- let page_value = self.max_page_value.as_ref().unwrap();
- self.compare_greater(page_value, max)
- });
- if update_max {
- self.max_column_value = self.max_page_value.clone();
- }
- }
-
- /// Evaluate `a > b` according to underlying logical type.
- fn compare_greater(&self, a: &T::T, b: &T::T) -> bool {
- if let Some(LogicalType::INTEGER(int_type)) = self.descr.logical_type() {
- if !int_type.is_signed {
- // need to compare unsigned
- return a.as_u64().unwrap() > b.as_u64().unwrap();
- }
- }
- a > b
- }
-}
-
-// ----------------------------------------------------------------------
-// Encoding support for column writer.
-// This mirrors parquet-mr default encodings for writes. See:
-// https://github.com/apache/parquet-mr/blob/master/parquet-column/src/main/java/org/apache/parquet/column/values/factory/DefaultV1ValuesWriterFactory.java
-// https://github.com/apache/parquet-mr/blob/master/parquet-column/src/main/java/org/apache/parquet/column/values/factory/DefaultV2ValuesWriterFactory.java
-
-/// Trait to define default encoding for types, including whether or not the type
-/// supports dictionary encoding.
-trait EncodingWriteSupport {
- /// Returns true if dictionary is supported for column writer, false otherwise.
- fn has_dictionary_support(props: &WriterProperties) -> bool;
-}
-
-/// Returns encoding for a column when no other encoding is provided in writer properties.
-fn fallback_encoding(kind: Type, props: &WriterProperties) -> Encoding {
- match (kind, props.writer_version()) {
- (Type::BOOLEAN, WriterVersion::PARQUET_2_0) => Encoding::RLE,
- (Type::INT32, WriterVersion::PARQUET_2_0) => Encoding::DELTA_BINARY_PACKED,
- (Type::INT64, WriterVersion::PARQUET_2_0) => Encoding::DELTA_BINARY_PACKED,
- (Type::BYTE_ARRAY, WriterVersion::PARQUET_2_0) => Encoding::DELTA_BYTE_ARRAY,
- (Type::FIXED_LEN_BYTE_ARRAY, WriterVersion::PARQUET_2_0) => {
- Encoding::DELTA_BYTE_ARRAY
- }
- _ => Encoding::PLAIN,
- }
-}
-
-/// Returns true if dictionary is supported for column writer, false otherwise.
-fn has_dictionary_support(kind: Type, props: &WriterProperties) -> bool {
- match (kind, props.writer_version()) {
- // Booleans do not support dict encoding and should use a fallback encoding.
- (Type::BOOLEAN, _) => false,
- // Dictionary encoding was not enabled in PARQUET 1.0
- (Type::FIXED_LEN_BYTE_ARRAY, WriterVersion::PARQUET_1_0) => false,
- (Type::FIXED_LEN_BYTE_ARRAY, WriterVersion::PARQUET_2_0) => true,
- _ => true,
- }
-}
-
-#[cfg(test)]
-mod tests {
- use rand::distributions::uniform::SampleUniform;
-
- use crate::column::{
- page::PageReader,
- reader::{get_column_reader, get_typed_column_reader, ColumnReaderImpl},
- };
- use crate::file::{
- properties::WriterProperties, reader::SerializedPageReader,
- writer::SerializedPageWriter,
- };
- use crate::schema::types::{ColumnDescriptor, ColumnPath, Type as SchemaType};
- use crate::util::{
- io::{FileSink, FileSource},
- test_common::{get_temp_file, random_numbers_range},
- };
-
- use super::*;
-
- #[test]
- fn test_column_writer_inconsistent_def_rep_length() {
- let page_writer = get_test_page_writer();
- let props = Arc::new(WriterProperties::builder().build());
- let mut writer = get_test_column_writer::<Int32Type>(page_writer, 1, 1, props);
- let res = writer.write_batch(&[1, 2, 3, 4], Some(&[1, 1, 1]), Some(&[0, 0]));
- assert!(res.is_err());
- if let Err(err) = res {
- assert_eq!(
- format!("{}", err),
- "Parquet error: Inconsistent length of definition and repetition levels: 3 != 2"
- );
- }
- }
-
- #[test]
- fn test_column_writer_invalid_def_levels() {
- let page_writer = get_test_page_writer();
- let props = Arc::new(WriterProperties::builder().build());
- let mut writer = get_test_column_writer::<Int32Type>(page_writer, 1, 0, props);
- let res = writer.write_batch(&[1, 2, 3, 4], None, None);
- assert!(res.is_err());
- if let Err(err) = res {
- assert_eq!(
- format!("{}", err),
- "Parquet error: Definition levels are required, because max definition level = 1"
- );
- }
- }
-
- #[test]
- fn test_column_writer_invalid_rep_levels() {
- let page_writer = get_test_page_writer();
- let props = Arc::new(WriterProperties::builder().build());
- let mut writer = get_test_column_writer::<Int32Type>(page_writer, 0, 1, props);
- let res = writer.write_batch(&[1, 2, 3, 4], None, None);
- assert!(res.is_err());
- if let Err(err) = res {
- assert_eq!(
- format!("{}", err),
- "Parquet error: Repetition levels are required, because max repetition level = 1"
- );
- }
- }
-
- #[test]
- fn test_column_writer_not_enough_values_to_write() {
- let page_writer = get_test_page_writer();
- let props = Arc::new(WriterProperties::builder().build());
- let mut writer = get_test_column_writer::<Int32Type>(page_writer, 1, 0, props);
- let res = writer.write_batch(&[1, 2], Some(&[1, 1, 1, 1]), None);
- assert!(res.is_err());
- if let Err(err) = res {
- assert_eq!(
- format!("{}", err),
- "Parquet error: Expected to write 4 values, but have only 2"
- );
- }
- }
-
- #[test]
- #[should_panic(expected = "Dictionary offset is already set")]
- fn test_column_writer_write_only_one_dictionary_page() {
- let page_writer = get_test_page_writer();
- let props = Arc::new(WriterProperties::builder().build());
- let mut writer = get_test_column_writer::<Int32Type>(page_writer, 0, 0, props);
- writer.write_batch(&[1, 2, 3, 4], None, None).unwrap();
- // First page should be correctly written.
- let res = writer.write_dictionary_page();
- assert!(res.is_ok());
- writer.write_dictionary_page().unwrap();
- }
-
- #[test]
- fn test_column_writer_error_when_writing_disabled_dictionary() {
- let page_writer = get_test_page_writer();
- let props = Arc::new(
- WriterProperties::builder()
- .set_dictionary_enabled(false)
- .build(),
- );
- let mut writer = get_test_column_writer::<Int32Type>(page_writer, 0, 0, props);
- writer.write_batch(&[1, 2, 3, 4], None, None).unwrap();
- let res = writer.write_dictionary_page();
- assert!(res.is_err());
- if let Err(err) = res {
- assert_eq!(
- format!("{}", err),
- "Parquet error: Dictionary encoder is not set"
- );
- }
- }
-
- #[test]
- fn test_column_writer_boolean_type_does_not_support_dictionary() {
- let page_writer = get_test_page_writer();
- let props = Arc::new(
- WriterProperties::builder()
- .set_dictionary_enabled(true)
- .build(),
- );
- let mut writer = get_test_column_writer::<BoolType>(page_writer, 0, 0, props);
- writer
- .write_batch(&[true, false, true, false], None, None)
- .unwrap();
-
- let (bytes_written, rows_written, metadata) = writer.close().unwrap();
- // PlainEncoder uses bit writer to write boolean values, which all fit into 1
- // byte.
- assert_eq!(bytes_written, 1);
- assert_eq!(rows_written, 4);
- assert_eq!(metadata.encodings(), &vec![Encoding::PLAIN, Encoding::RLE]);
- assert_eq!(metadata.num_values(), 4); // just values
- assert_eq!(metadata.dictionary_page_offset(), None);
- }
-
- #[test]
- fn test_column_writer_default_encoding_support_bool() {
- check_encoding_write_support::<BoolType>(
- WriterVersion::PARQUET_1_0,
- true,
- &[true, false],
- None,
- &[Encoding::PLAIN, Encoding::RLE],
- );
- check_encoding_write_support::<BoolType>(
- WriterVersion::PARQUET_1_0,
- false,
- &[true, false],
- None,
- &[Encoding::PLAIN, Encoding::RLE],
- );
- check_encoding_write_support::<BoolType>(
- WriterVersion::PARQUET_2_0,
- true,
- &[true, false],
- None,
- &[Encoding::RLE, Encoding::RLE],
- );
- check_encoding_write_support::<BoolType>(
- WriterVersion::PARQUET_2_0,
- false,
- &[true, false],
- None,
- &[Encoding::RLE, Encoding::RLE],
- );
- }
-
- #[test]
- fn test_column_writer_default_encoding_support_int32() {
- check_encoding_write_support::<Int32Type>(
- WriterVersion::PARQUET_1_0,
- true,
- &[1, 2],
- Some(0),
- &[Encoding::PLAIN, Encoding::RLE_DICTIONARY, Encoding::RLE],
- );
- check_encoding_write_support::<Int32Type>(
- WriterVersion::PARQUET_1_0,
- false,
- &[1, 2],
- None,
- &[Encoding::PLAIN, Encoding::RLE],
- );
- check_encoding_write_support::<Int32Type>(
- WriterVersion::PARQUET_2_0,
- true,
- &[1, 2],
- Some(0),
- &[Encoding::PLAIN, Encoding::RLE_DICTIONARY, Encoding::RLE],
- );
- check_encoding_write_support::<Int32Type>(
- WriterVersion::PARQUET_2_0,
- false,
- &[1, 2],
- None,
- &[Encoding::DELTA_BINARY_PACKED, Encoding::RLE],
- );
- }
-
- #[test]
- fn test_column_writer_default_encoding_support_int64() {
- check_encoding_write_support::<Int64Type>(
- WriterVersion::PARQUET_1_0,
- true,
- &[1, 2],
- Some(0),
- &[Encoding::PLAIN, Encoding::RLE_DICTIONARY, Encoding::RLE],
- );
- check_encoding_write_support::<Int64Type>(
- WriterVersion::PARQUET_1_0,
- false,
- &[1, 2],
- None,
- &[Encoding::PLAIN, Encoding::RLE],
- );
- check_encoding_write_support::<Int64Type>(
- WriterVersion::PARQUET_2_0,
- true,
- &[1, 2],
- Some(0),
- &[Encoding::PLAIN, Encoding::RLE_DICTIONARY, Encoding::RLE],
- );
- check_encoding_write_support::<Int64Type>(
- WriterVersion::PARQUET_2_0,
- false,
- &[1, 2],
- None,
- &[Encoding::DELTA_BINARY_PACKED, Encoding::RLE],
- );
- }
-
- #[test]
- fn test_column_writer_default_encoding_support_int96() {
- check_encoding_write_support::<Int96Type>(
- WriterVersion::PARQUET_1_0,
- true,
- &[Int96::from(vec![1, 2, 3])],
- Some(0),
- &[Encoding::PLAIN, Encoding::RLE_DICTIONARY, Encoding::RLE],
- );
- check_encoding_write_support::<Int96Type>(
- WriterVersion::PARQUET_1_0,
- false,
- &[Int96::from(vec![1, 2, 3])],
- None,
- &[Encoding::PLAIN, Encoding::RLE],
- );
- check_encoding_write_support::<Int96Type>(
- WriterVersion::PARQUET_2_0,
- true,
- &[Int96::from(vec![1, 2, 3])],
- Some(0),
- &[Encoding::PLAIN, Encoding::RLE_DICTIONARY, Encoding::RLE],
- );
- check_encoding_write_support::<Int96Type>(
- WriterVersion::PARQUET_2_0,
- false,
- &[Int96::from(vec![1, 2, 3])],
- None,
- &[Encoding::PLAIN, Encoding::RLE],
- );
- }
-
- #[test]
- fn test_column_writer_default_encoding_support_float() {
- check_encoding_write_support::<FloatType>(
- WriterVersion::PARQUET_1_0,
- true,
- &[1.0, 2.0],
- Some(0),
- &[Encoding::PLAIN, Encoding::RLE_DICTIONARY, Encoding::RLE],
- );
- check_encoding_write_support::<FloatType>(
- WriterVersion::PARQUET_1_0,
- false,
- &[1.0, 2.0],
- None,
- &[Encoding::PLAIN, Encoding::RLE],
- );
- check_encoding_write_support::<FloatType>(
- WriterVersion::PARQUET_2_0,
- true,
- &[1.0, 2.0],
- Some(0),
- &[Encoding::PLAIN, Encoding::RLE_DICTIONARY, Encoding::RLE],
- );
- check_encoding_write_support::<FloatType>(
- WriterVersion::PARQUET_2_0,
- false,
- &[1.0, 2.0],
- None,
- &[Encoding::PLAIN, Encoding::RLE],
- );
- }
-
- #[test]
- fn test_column_writer_default_encoding_support_double() {
- check_encoding_write_support::<DoubleType>(
- WriterVersion::PARQUET_1_0,
- true,
- &[1.0, 2.0],
- Some(0),
- &[Encoding::PLAIN, Encoding::RLE_DICTIONARY, Encoding::RLE],
- );
- check_encoding_write_support::<DoubleType>(
- WriterVersion::PARQUET_1_0,
- false,
- &[1.0, 2.0],
- None,
- &[Encoding::PLAIN, Encoding::RLE],
- );
- check_encoding_write_support::<DoubleType>(
- WriterVersion::PARQUET_2_0,
- true,
- &[1.0, 2.0],
- Some(0),
- &[Encoding::PLAIN, Encoding::RLE_DICTIONARY, Encoding::RLE],
- );
- check_encoding_write_support::<DoubleType>(
- WriterVersion::PARQUET_2_0,
- false,
- &[1.0, 2.0],
- None,
- &[Encoding::PLAIN, Encoding::RLE],
- );
- }
-
- #[test]
- fn test_column_writer_default_encoding_support_byte_array() {
- check_encoding_write_support::<ByteArrayType>(
- WriterVersion::PARQUET_1_0,
- true,
- &[ByteArray::from(vec![1u8])],
- Some(0),
- &[Encoding::PLAIN, Encoding::RLE_DICTIONARY, Encoding::RLE],
- );
- check_encoding_write_support::<ByteArrayType>(
- WriterVersion::PARQUET_1_0,
- false,
- &[ByteArray::from(vec![1u8])],
- None,
- &[Encoding::PLAIN, Encoding::RLE],
- );
- check_encoding_write_support::<ByteArrayType>(
- WriterVersion::PARQUET_2_0,
- true,
- &[ByteArray::from(vec![1u8])],
- Some(0),
- &[Encoding::PLAIN, Encoding::RLE_DICTIONARY, Encoding::RLE],
- );
- check_encoding_write_support::<ByteArrayType>(
- WriterVersion::PARQUET_2_0,
- false,
- &[ByteArray::from(vec![1u8])],
- None,
- &[Encoding::DELTA_BYTE_ARRAY, Encoding::RLE],
- );
- }
-
- #[test]
- fn test_column_writer_default_encoding_support_fixed_len_byte_array() {
- check_encoding_write_support::<FixedLenByteArrayType>(
- WriterVersion::PARQUET_1_0,
- true,
- &[ByteArray::from(vec![1u8]).into()],
- None,
- &[Encoding::PLAIN, Encoding::RLE],
- );
- check_encoding_write_support::<FixedLenByteArrayType>(
- WriterVersion::PARQUET_1_0,
- false,
- &[ByteArray::from(vec![1u8]).into()],
- None,
- &[Encoding::PLAIN, Encoding::RLE],
- );
- check_encoding_write_support::<FixedLenByteArrayType>(
- WriterVersion::PARQUET_2_0,
- true,
- &[ByteArray::from(vec![1u8]).into()],
- Some(0),
- &[Encoding::PLAIN, Encoding::RLE_DICTIONARY, Encoding::RLE],
- );
- check_encoding_write_support::<FixedLenByteArrayType>(
- WriterVersion::PARQUET_2_0,
- false,
- &[ByteArray::from(vec![1u8]).into()],
- None,
- &[Encoding::DELTA_BYTE_ARRAY, Encoding::RLE],
- );
- }
-
- #[test]
- fn test_column_writer_check_metadata() {
- let page_writer = get_test_page_writer();
- let props = Arc::new(WriterProperties::builder().build());
- let mut writer = get_test_column_writer::<Int32Type>(page_writer, 0, 0, props);
- writer.write_batch(&[1, 2, 3, 4], None, None).unwrap();
-
- let (bytes_written, rows_written, metadata) = writer.close().unwrap();
- assert_eq!(bytes_written, 20);
- assert_eq!(rows_written, 4);
- assert_eq!(
- metadata.encodings(),
- &vec![Encoding::PLAIN, Encoding::RLE_DICTIONARY, Encoding::RLE]
- );
- assert_eq!(metadata.num_values(), 8); // dictionary + value indexes
- assert_eq!(metadata.compressed_size(), 20);
- assert_eq!(metadata.uncompressed_size(), 20);
- assert_eq!(metadata.data_page_offset(), 0);
- assert_eq!(metadata.dictionary_page_offset(), Some(0));
- if let Some(stats) = metadata.statistics() {
- assert!(stats.has_min_max_set());
- assert_eq!(stats.null_count(), 0);
- assert_eq!(stats.distinct_count(), None);
- if let Statistics::Int32(stats) = stats {
- assert_eq!(stats.min(), &1);
- assert_eq!(stats.max(), &4);
- } else {
- panic!("expecting Statistics::Int32");
- }
- } else {
- panic!("metadata missing statistics");
- }
- }
-
- #[test]
- fn test_column_writer_precalculated_statistics() {
- let page_writer = get_test_page_writer();
- let props = Arc::new(WriterProperties::builder().build());
- let mut writer = get_test_column_writer::<Int32Type>(page_writer, 0, 0, props);
- writer
- .write_batch_with_statistics(
- &[1, 2, 3, 4],
- None,
- None,
- &Some(-17),
- &Some(9000),
- Some(21),
- Some(55),
- )
- .unwrap();
-
- let (bytes_written, rows_written, metadata) = writer.close().unwrap();
- assert_eq!(bytes_written, 20);
- assert_eq!(rows_written, 4);
- assert_eq!(
- metadata.encodings(),
- &vec![Encoding::PLAIN, Encoding::RLE_DICTIONARY, Encoding::RLE]
- );
- assert_eq!(metadata.num_values(), 8); // dictionary + value indexes
- assert_eq!(metadata.compressed_size(), 20);
- assert_eq!(metadata.uncompressed_size(), 20);
- assert_eq!(metadata.data_page_offset(), 0);
- assert_eq!(metadata.dictionary_page_offset(), Some(0));
- if let Some(stats) = metadata.statistics() {
- assert!(stats.has_min_max_set());
- assert_eq!(stats.null_count(), 21);
- assert_eq!(stats.distinct_count().unwrap_or(0), 55);
- if let Statistics::Int32(stats) = stats {
- assert_eq!(stats.min(), &-17);
- assert_eq!(stats.max(), &9000);
- } else {
- panic!("expecting Statistics::Int32");
- }
- } else {
- panic!("metadata missing statistics");
- }
- }
-
- #[test]
- fn test_column_writer_empty_column_roundtrip() {
- let props = WriterProperties::builder().build();
- column_roundtrip::<Int32Type>("test_col_writer_rnd_1", props, &[], None, None);
- }
-
- #[test]
- fn test_column_writer_non_nullable_values_roundtrip() {
- let props = WriterProperties::builder().build();
- column_roundtrip_random::<Int32Type>(
- "test_col_writer_rnd_2",
- props,
- 1024,
- std::i32::MIN,
- std::i32::MAX,
- 0,
- 0,
- );
- }
-
- #[test]
- fn test_column_writer_nullable_non_repeated_values_roundtrip() {
- let props = WriterProperties::builder().build();
- column_roundtrip_random::<Int32Type>(
- "test_column_writer_nullable_non_repeated_values_roundtrip",
- props,
- 1024,
- std::i32::MIN,
- std::i32::MAX,
- 10,
- 0,
- );
- }
-
- #[test]
- fn test_column_writer_nullable_repeated_values_roundtrip() {
- let props = WriterProperties::builder().build();
- column_roundtrip_random::<Int32Type>(
- "test_col_writer_rnd_3",
- props,
- 1024,
- std::i32::MIN,
- std::i32::MAX,
- 10,
- 10,
- );
- }
-
- #[test]
- fn test_column_writer_dictionary_fallback_small_data_page() {
- let props = WriterProperties::builder()
- .set_dictionary_pagesize_limit(32)
- .set_data_pagesize_limit(32)
- .build();
- column_roundtrip_random::<Int32Type>(
- "test_col_writer_rnd_4",
- props,
- 1024,
- std::i32::MIN,
- std::i32::MAX,
- 10,
- 10,
- );
- }
-
- #[test]
- fn test_column_writer_small_write_batch_size() {
- for i in &[1usize, 2, 5, 10, 11, 1023] {
- let props = WriterProperties::builder().set_write_batch_size(*i).build();
-
- column_roundtrip_random::<Int32Type>(
- "test_col_writer_rnd_5",
- props,
- 1024,
- std::i32::MIN,
- std::i32::MAX,
- 10,
- 10,
- );
- }
- }
-
- #[test]
- fn test_column_writer_dictionary_disabled_v1() {
- let props = WriterProperties::builder()
- .set_writer_version(WriterVersion::PARQUET_1_0)
- .set_dictionary_enabled(false)
- .build();
- column_roundtrip_random::<Int32Type>(
- "test_col_writer_rnd_6",
- props,
- 1024,
- std::i32::MIN,
- std::i32::MAX,
- 10,
- 10,
- );
- }
-
- #[test]
- fn test_column_writer_dictionary_disabled_v2() {
- let props = WriterProperties::builder()
- .set_writer_version(WriterVersion::PARQUET_2_0)
- .set_dictionary_enabled(false)
- .build();
- column_roundtrip_random::<Int32Type>(
- "test_col_writer_rnd_7",
- props,
- 1024,
- std::i32::MIN,
- std::i32::MAX,
- 10,
- 10,
- );
- }
-
- #[test]
- fn test_column_writer_compression_v1() {
- let props = WriterProperties::builder()
- .set_writer_version(WriterVersion::PARQUET_1_0)
- .set_compression(Compression::SNAPPY)
- .build();
- column_roundtrip_random::<Int32Type>(
- "test_col_writer_rnd_8",
- props,
- 2048,
- std::i32::MIN,
- std::i32::MAX,
- 10,
- 10,
- );
- }
-
- #[test]
- fn test_column_writer_compression_v2() {
- let props = WriterProperties::builder()
- .set_writer_version(WriterVersion::PARQUET_2_0)
- .set_compression(Compression::SNAPPY)
- .build();
- column_roundtrip_random::<Int32Type>(
- "test_col_writer_rnd_9",
- props,
- 2048,
- std::i32::MIN,
- std::i32::MAX,
- 10,
- 10,
- );
- }
-
- #[test]
- fn test_column_writer_add_data_pages_with_dict() {
- // ARROW-5129: Test verifies that we add data page in case of dictionary encoding
- // and no fallback occurred so far.
- let file = get_temp_file("test_column_writer_add_data_pages_with_dict", &[]);
- let sink = FileSink::new(&file);
- let page_writer = Box::new(SerializedPageWriter::new(sink));
- let props = Arc::new(
- WriterProperties::builder()
- .set_data_pagesize_limit(15) // actually each page will have size 15-18 bytes
- .set_write_batch_size(3) // write 3 values at a time
- .build(),
- );
- let data = &[1, 2, 3, 4, 5, 6, 7, 8, 9, 10];
- let mut writer = get_test_column_writer::<Int32Type>(page_writer, 0, 0, props);
- writer.write_batch(data, None, None).unwrap();
- let (bytes_written, _, _) = writer.close().unwrap();
-
- // Read pages and check the sequence
- let source = FileSource::new(&file, 0, bytes_written as usize);
- let mut page_reader = Box::new(
- SerializedPageReader::new(
- source,
- data.len() as i64,
- Compression::UNCOMPRESSED,
- Int32Type::get_physical_type(),
- )
- .unwrap(),
- );
- let mut res = Vec::new();
- while let Some(page) = page_reader.get_next_page().unwrap() {
- res.push((page.page_type(), page.num_values()));
- }
- assert_eq!(
- res,
- vec![
- (PageType::DICTIONARY_PAGE, 10),
- (PageType::DATA_PAGE, 3),
- (PageType::DATA_PAGE, 3),
- (PageType::DATA_PAGE, 3),
- (PageType::DATA_PAGE, 1)
- ]
- );
- }
-
- #[test]
- fn test_float_statistics_nan_middle() {
- let stats = statistics_roundtrip::<FloatType>(&[1.0, f32::NAN, 2.0]);
- assert!(stats.has_min_max_set());
- if let Statistics::Float(stats) = stats {
- assert_eq!(stats.min(), &1.0);
- assert_eq!(stats.max(), &2.0);
- } else {
- panic!("expecting Statistics::Float");
- }
- }
-
- #[test]
- fn test_float_statistics_nan_start() {
- let stats = statistics_roundtrip::<FloatType>(&[f32::NAN, 1.0, 2.0]);
- assert!(stats.has_min_max_set());
- if let Statistics::Float(stats) = stats {
- assert_eq!(stats.min(), &1.0);
- assert_eq!(stats.max(), &2.0);
- } else {
- panic!("expecting Statistics::Float");
- }
- }
-
- #[test]
- fn test_float_statistics_nan_only() {
- let stats = statistics_roundtrip::<FloatType>(&[f32::NAN, f32::NAN]);
- assert!(!stats.has_min_max_set());
- assert!(matches!(stats, Statistics::Float(_)));
- }
-
- #[test]
- fn test_double_statistics_nan_middle() {
- let stats = statistics_roundtrip::<DoubleType>(&[1.0, f64::NAN, 2.0]);
- assert!(stats.has_min_max_set());
- if let Statistics::Double(stats) = stats {
- assert_eq!(stats.min(), &1.0);
- assert_eq!(stats.max(), &2.0);
- } else {
- panic!("expecting Statistics::Float");
- }
- }
-
- #[test]
- fn test_double_statistics_nan_start() {
- let stats = statistics_roundtrip::<DoubleType>(&[f64::NAN, 1.0, 2.0]);
- assert!(stats.has_min_max_set());
- if let Statistics::Double(stats) = stats {
- assert_eq!(stats.min(), &1.0);
- assert_eq!(stats.max(), &2.0);
- } else {
- panic!("expecting Statistics::Float");
- }
- }
-
- #[test]
- fn test_double_statistics_nan_only() {
- let stats = statistics_roundtrip::<DoubleType>(&[f64::NAN, f64::NAN]);
- assert!(!stats.has_min_max_set());
- assert!(matches!(stats, Statistics::Double(_)));
- }
-
- /// Performs write-read roundtrip with randomly generated values and levels.
- /// `max_size` is maximum number of values or levels (if `max_def_level` > 0) to write
- /// for a column.
- fn column_roundtrip_random<T: DataType>(
- file_name: &str,
- props: WriterProperties,
- max_size: usize,
- min_value: T::T,
- max_value: T::T,
- max_def_level: i16,
- max_rep_level: i16,
- ) where
- T::T: PartialOrd + SampleUniform + Copy,
- {
- let mut num_values: usize = 0;
-
- let mut buf: Vec<i16> = Vec::new();
- let def_levels = if max_def_level > 0 {
- random_numbers_range(max_size, 0, max_def_level + 1, &mut buf);
- for &dl in &buf[..] {
- if dl == max_def_level {
- num_values += 1;
- }
- }
- Some(&buf[..])
- } else {
- num_values = max_size;
- None
- };
-
- let mut buf: Vec<i16> = Vec::new();
- let rep_levels = if max_rep_level > 0 {
- random_numbers_range(max_size, 0, max_rep_level + 1, &mut buf);
- Some(&buf[..])
- } else {
- None
- };
-
- let mut values: Vec<T::T> = Vec::new();
- random_numbers_range(num_values, min_value, max_value, &mut values);
-
- column_roundtrip::<T>(file_name, props, &values[..], def_levels, rep_levels);
- }
-
- /// Performs write-read roundtrip and asserts written values and levels.
- fn column_roundtrip<'a, T: DataType>(
- file_name: &'a str,
- props: WriterProperties,
- values: &[T::T],
- def_levels: Option<&[i16]>,
- rep_levels: Option<&[i16]>,
- ) {
- let file = get_temp_file(file_name, &[]);
- let sink = FileSink::new(&file);
- let page_writer = Box::new(SerializedPageWriter::new(sink));
-
- let max_def_level = match def_levels {
- Some(buf) => *buf.iter().max().unwrap_or(&0i16),
- None => 0i16,
- };
-
- let max_rep_level = match rep_levels {
- Some(buf) => *buf.iter().max().unwrap_or(&0i16),
- None => 0i16,
- };
-
- let mut max_batch_size = values.len();
- if let Some(levels) = def_levels {
- max_batch_size = cmp::max(max_batch_size, levels.len());
- }
- if let Some(levels) = rep_levels {
- max_batch_size = cmp::max(max_batch_size, levels.len());
- }
-
- let mut writer = get_test_column_writer::<T>(
- page_writer,
- max_def_level,
- max_rep_level,
- Arc::new(props),
- );
-
- let values_written = writer.write_batch(values, def_levels, rep_levels).unwrap();
- assert_eq!(values_written, values.len());
- let (bytes_written, rows_written, column_metadata) = writer.close().unwrap();
-
- let source = FileSource::new(&file, 0, bytes_written as usize);
- let page_reader = Box::new(
- SerializedPageReader::new(
- source,
- column_metadata.num_values(),
- column_metadata.compression(),
- T::get_physical_type(),
- )
- .unwrap(),
- );
- let reader =
- get_test_column_reader::<T>(page_reader, max_def_level, max_rep_level);
-
- let mut actual_values = vec![T::T::default(); max_batch_size];
- let mut actual_def_levels = def_levels.map(|_| vec![0i16; max_batch_size]);
- let mut actual_rep_levels = rep_levels.map(|_| vec![0i16; max_batch_size]);
-
- let (values_read, levels_read) = read_fully(
- reader,
- max_batch_size,
- actual_def_levels.as_mut(),
- actual_rep_levels.as_mut(),
- actual_values.as_mut_slice(),
- );
-
- // Assert values, definition and repetition levels.
-
- assert_eq!(&actual_values[..values_read], values);
- match actual_def_levels {
- Some(ref vec) => assert_eq!(Some(&vec[..levels_read]), def_levels),
- None => assert_eq!(None, def_levels),
- }
- match actual_rep_levels {
- Some(ref vec) => assert_eq!(Some(&vec[..levels_read]), rep_levels),
- None => assert_eq!(None, rep_levels),
- }
-
- // Assert written rows.
-
- if let Some(levels) = actual_rep_levels {
- let mut actual_rows_written = 0;
- for l in levels {
- if l == 0 {
- actual_rows_written += 1;
- }
- }
- assert_eq!(actual_rows_written, rows_written);
- } else if actual_def_levels.is_some() {
- assert_eq!(levels_read as u64, rows_written);
- } else {
- assert_eq!(values_read as u64, rows_written);
- }
- }
-
- /// Performs write of provided values and returns column metadata of those values.
- /// Used to test encoding support for column writer.
- fn column_write_and_get_metadata<T: DataType>(
- props: WriterProperties,
- values: &[T::T],
- ) -> ColumnChunkMetaData {
- let page_writer = get_test_page_writer();
- let props = Arc::new(props);
- let mut writer = get_test_column_writer::<T>(page_writer, 0, 0, props);
- writer.write_batch(values, None, None).unwrap();
- let (_, _, metadata) = writer.close().unwrap();
- metadata
- }
-
- // Function to use in tests for EncodingWriteSupport. This checks that dictionary
- // offset and encodings to make sure that column writer uses provided by trait
- // encodings.
- fn check_encoding_write_support<T: DataType>(
- version: WriterVersion,
- dict_enabled: bool,
- data: &[T::T],
- dictionary_page_offset: Option<i64>,
- encodings: &[Encoding],
- ) {
- let props = WriterProperties::builder()
- .set_writer_version(version)
- .set_dictionary_enabled(dict_enabled)
- .build();
- let meta = column_write_and_get_metadata::<T>(props, data);
- assert_eq!(meta.dictionary_page_offset(), dictionary_page_offset);
- assert_eq!(meta.encodings(), &encodings);
- }
-
- /// Reads one batch of data, considering that batch is large enough to capture all of
- /// the values and levels.
- fn read_fully<T: DataType>(
- mut reader: ColumnReaderImpl<T>,
- batch_size: usize,
- mut def_levels: Option<&mut Vec<i16>>,
- mut rep_levels: Option<&mut Vec<i16>>,
- values: &mut [T::T],
- ) -> (usize, usize) {
- let actual_def_levels = def_levels.as_mut().map(|vec| &mut vec[..]);
- let actual_rep_levels = rep_levels.as_mut().map(|vec| &mut vec[..]);
- reader
- .read_batch(batch_size, actual_def_levels, actual_rep_levels, values)
- .unwrap()
- }
-
- /// Returns column writer.
- fn get_test_column_writer<T: DataType>(
- page_writer: Box<dyn PageWriter>,
- max_def_level: i16,
- max_rep_level: i16,
- props: WriterPropertiesPtr,
- ) -> ColumnWriterImpl<T> {
- let descr = Arc::new(get_test_column_descr::<T>(max_def_level, max_rep_level));
- let column_writer = get_column_writer(descr, props, page_writer);
- get_typed_column_writer::<T>(column_writer)
- }
-
- /// Returns column reader.
- fn get_test_column_reader<T: DataType>(
- page_reader: Box<dyn PageReader>,
- max_def_level: i16,
- max_rep_level: i16,
- ) -> ColumnReaderImpl<T> {
- let descr = Arc::new(get_test_column_descr::<T>(max_def_level, max_rep_level));
- let column_reader = get_column_reader(descr, page_reader);
- get_typed_column_reader::<T>(column_reader)
- }
-
- /// Returns descriptor for primitive column.
- fn get_test_column_descr<T: DataType>(
- max_def_level: i16,
- max_rep_level: i16,
- ) -> ColumnDescriptor {
- let path = ColumnPath::from("col");
- let tpe = SchemaType::primitive_type_builder("col", T::get_physical_type())
- // length is set for "encoding support" tests for FIXED_LEN_BYTE_ARRAY type,
- // it should be no-op for other types
- .with_length(1)
- .build()
- .unwrap();
- ColumnDescriptor::new(Arc::new(tpe), max_def_level, max_rep_level, path)
- }
-
- /// Returns page writer that collects pages without serializing them.
- fn get_test_page_writer() -> Box<dyn PageWriter> {
- Box::new(TestPageWriter {})
- }
-
- struct TestPageWriter {}
-
- impl PageWriter for TestPageWriter {
- fn write_page(&mut self, page: CompressedPage) -> Result<PageWriteSpec> {
- let mut res = PageWriteSpec::new();
- res.page_type = page.page_type();
- res.uncompressed_size = page.uncompressed_size();
- res.compressed_size = page.compressed_size();
- res.num_values = page.num_values();
- res.offset = 0;
- res.bytes_written = page.data().len() as u64;
- Ok(res)
- }
-
- fn write_metadata(&mut self, _metadata: &ColumnChunkMetaData) -> Result<()> {
- Ok(())
- }
-
- fn close(&mut self) -> Result<()> {
- Ok(())
- }
- }
-
- /// Write data into parquet using [`get_test_page_writer`] and [`get_test_column_writer`] and returns generated statistics.
- fn statistics_roundtrip<T: DataType>(values: &[<T as DataType>::T]) -> Statistics {
- let page_writer = get_test_page_writer();
- let props = Arc::new(WriterProperties::builder().build());
- let mut writer = get_test_column_writer::<T>(page_writer, 0, 0, props);
- writer.write_batch(values, None, None).unwrap();
-
- let (_bytes_written, _rows_written, metadata) = writer.close().unwrap();
- if let Some(stats) = metadata.statistics() {
- stats.clone()
- } else {
- panic!("metadata missing statistics");
- }
- }
-}
diff --git a/parquet/src/compression.rs b/parquet/src/compression.rs
deleted file mode 100644
index a115597..0000000
--- a/parquet/src/compression.rs
+++ /dev/null
@@ -1,393 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-//! Contains codec interface and supported codec implementations.
-//!
-//! See [`Compression`](crate::basic::Compression) enum for all available compression
-//! algorithms.
-//!
-//! # Example
-//!
-//! ```no_run
-//! use parquet::{basic::Compression, compression::create_codec};
-//!
-//! let mut codec = match create_codec(Compression::SNAPPY) {
-//! Ok(Some(codec)) => codec,
-//! _ => panic!(),
-//! };
-//!
-//! let data = vec![b'p', b'a', b'r', b'q', b'u', b'e', b't'];
-//! let mut compressed = vec![];
-//! codec.compress(&data[..], &mut compressed).unwrap();
-//!
-//! let mut output = vec![];
-//! codec.decompress(&compressed[..], &mut output).unwrap();
-//!
-//! assert_eq!(output, data);
-//! ```
-
-use crate::basic::Compression as CodecType;
-use crate::errors::{ParquetError, Result};
-
-/// Parquet compression codec interface.
-pub trait Codec {
- /// Compresses data stored in slice `input_buf` and writes the compressed result
- /// to `output_buf`.
- /// Note that you'll need to call `clear()` before reusing the same `output_buf`
- /// across different `compress` calls.
- fn compress(&mut self, input_buf: &[u8], output_buf: &mut Vec<u8>) -> Result<()>;
-
- /// Decompresses data stored in slice `input_buf` and writes output to `output_buf`.
- /// Returns the total number of bytes written.
- fn decompress(&mut self, input_buf: &[u8], output_buf: &mut Vec<u8>)
- -> Result<usize>;
-}
-
-/// Given the compression type `codec`, returns a codec used to compress and decompress
-/// bytes for the compression type.
-/// This returns `None` if the codec type is `UNCOMPRESSED`.
-pub fn create_codec(codec: CodecType) -> Result<Option<Box<dyn Codec>>> {
- match codec {
- #[cfg(any(feature = "brotli", test))]
- CodecType::BROTLI => Ok(Some(Box::new(BrotliCodec::new()))),
- #[cfg(any(feature = "flate2", test))]
- CodecType::GZIP => Ok(Some(Box::new(GZipCodec::new()))),
- #[cfg(any(feature = "snap", test))]
- CodecType::SNAPPY => Ok(Some(Box::new(SnappyCodec::new()))),
- #[cfg(any(feature = "lz4", test))]
- CodecType::LZ4 => Ok(Some(Box::new(LZ4Codec::new()))),
- #[cfg(any(feature = "zstd", test))]
- CodecType::ZSTD => Ok(Some(Box::new(ZSTDCodec::new()))),
- CodecType::UNCOMPRESSED => Ok(None),
- _ => Err(nyi_err!("The codec type {} is not supported yet", codec)),
- }
-}
-
-#[cfg(any(feature = "snap", test))]
-mod snappy_codec {
- use snap::raw::{decompress_len, max_compress_len, Decoder, Encoder};
-
- use crate::compression::Codec;
- use crate::errors::Result;
-
- /// Codec for Snappy compression format.
- pub struct SnappyCodec {
- decoder: Decoder,
- encoder: Encoder,
- }
-
- impl SnappyCodec {
- /// Creates new Snappy compression codec.
- pub(crate) fn new() -> Self {
- Self {
- decoder: Decoder::new(),
- encoder: Encoder::new(),
- }
- }
- }
-
- impl Codec for SnappyCodec {
- fn decompress(
- &mut self,
- input_buf: &[u8],
- output_buf: &mut Vec<u8>,
- ) -> Result<usize> {
- let len = decompress_len(input_buf)?;
- output_buf.resize(len, 0);
- self.decoder
- .decompress(input_buf, output_buf)
- .map_err(|e| e.into())
- }
-
- fn compress(&mut self, input_buf: &[u8], output_buf: &mut Vec<u8>) -> Result<()> {
- let output_buf_len = output_buf.len();
- let required_len = max_compress_len(input_buf.len());
- output_buf.resize(output_buf_len + required_len, 0);
- let n = self
- .encoder
- .compress(input_buf, &mut output_buf[output_buf_len..])?;
- output_buf.truncate(output_buf_len + n);
- Ok(())
- }
- }
-}
-#[cfg(any(feature = "snap", test))]
-pub use snappy_codec::*;
-
-#[cfg(any(feature = "flate2", test))]
-mod gzip_codec {
-
- use std::io::{Read, Write};
-
- use flate2::{read, write, Compression};
-
- use crate::compression::Codec;
- use crate::errors::Result;
-
- /// Codec for GZIP compression algorithm.
- pub struct GZipCodec {}
-
- impl GZipCodec {
- /// Creates new GZIP compression codec.
- pub(crate) fn new() -> Self {
- Self {}
- }
- }
-
- impl Codec for GZipCodec {
- fn decompress(
- &mut self,
- input_buf: &[u8],
- output_buf: &mut Vec<u8>,
- ) -> Result<usize> {
- let mut decoder = read::GzDecoder::new(input_buf);
- decoder.read_to_end(output_buf).map_err(|e| e.into())
- }
-
- fn compress(&mut self, input_buf: &[u8], output_buf: &mut Vec<u8>) -> Result<()> {
- let mut encoder = write::GzEncoder::new(output_buf, Compression::default());
- encoder.write_all(input_buf)?;
- encoder.try_finish().map_err(|e| e.into())
- }
- }
-}
-#[cfg(any(feature = "flate2", test))]
-pub use gzip_codec::*;
-
-#[cfg(any(feature = "brotli", test))]
-mod brotli_codec {
-
- use std::io::{Read, Write};
-
- use crate::compression::Codec;
- use crate::errors::Result;
-
- const BROTLI_DEFAULT_BUFFER_SIZE: usize = 4096;
- const BROTLI_DEFAULT_COMPRESSION_QUALITY: u32 = 1; // supported levels 0-9
- const BROTLI_DEFAULT_LG_WINDOW_SIZE: u32 = 22; // recommended between 20-22
-
- /// Codec for Brotli compression algorithm.
- pub struct BrotliCodec {}
-
- impl BrotliCodec {
- /// Creates new Brotli compression codec.
- pub(crate) fn new() -> Self {
- Self {}
- }
- }
-
- impl Codec for BrotliCodec {
- fn decompress(
- &mut self,
- input_buf: &[u8],
- output_buf: &mut Vec<u8>,
- ) -> Result<usize> {
- brotli::Decompressor::new(input_buf, BROTLI_DEFAULT_BUFFER_SIZE)
- .read_to_end(output_buf)
- .map_err(|e| e.into())
- }
-
- fn compress(&mut self, input_buf: &[u8], output_buf: &mut Vec<u8>) -> Result<()> {
- let mut encoder = brotli::CompressorWriter::new(
- output_buf,
- BROTLI_DEFAULT_BUFFER_SIZE,
- BROTLI_DEFAULT_COMPRESSION_QUALITY,
- BROTLI_DEFAULT_LG_WINDOW_SIZE,
- );
- encoder.write_all(input_buf)?;
- encoder.flush().map_err(|e| e.into())
- }
- }
-}
-#[cfg(any(feature = "brotli", test))]
-pub use brotli_codec::*;
-
-#[cfg(any(feature = "lz4", test))]
-mod lz4_codec {
- use std::io::{Read, Write};
-
- use crate::compression::Codec;
- use crate::errors::Result;
-
- const LZ4_BUFFER_SIZE: usize = 4096;
-
- /// Codec for LZ4 compression algorithm.
- pub struct LZ4Codec {}
-
- impl LZ4Codec {
- /// Creates new LZ4 compression codec.
- pub(crate) fn new() -> Self {
- Self {}
- }
- }
-
- impl Codec for LZ4Codec {
- fn decompress(
- &mut self,
- input_buf: &[u8],
- output_buf: &mut Vec<u8>,
- ) -> Result<usize> {
- let mut decoder = lz4::Decoder::new(input_buf)?;
- let mut buffer: [u8; LZ4_BUFFER_SIZE] = [0; LZ4_BUFFER_SIZE];
- let mut total_len = 0;
- loop {
- let len = decoder.read(&mut buffer)?;
- if len == 0 {
- break;
- }
- total_len += len;
- output_buf.write_all(&buffer[0..len])?;
- }
- Ok(total_len)
- }
-
- fn compress(&mut self, input_buf: &[u8], output_buf: &mut Vec<u8>) -> Result<()> {
- let mut encoder = lz4::EncoderBuilder::new().build(output_buf)?;
- let mut from = 0;
- loop {
- let to = std::cmp::min(from + LZ4_BUFFER_SIZE, input_buf.len());
- encoder.write_all(&input_buf[from..to])?;
- from += LZ4_BUFFER_SIZE;
- if from >= input_buf.len() {
- break;
- }
- }
- encoder.finish().1.map_err(|e| e.into())
- }
- }
-}
-#[cfg(any(feature = "lz4", test))]
-pub use lz4_codec::*;
-
-#[cfg(any(feature = "zstd", test))]
-mod zstd_codec {
- use std::io::{self, Write};
-
- use crate::compression::Codec;
- use crate::errors::Result;
-
- /// Codec for Zstandard compression algorithm.
- pub struct ZSTDCodec {}
-
- impl ZSTDCodec {
- /// Creates new Zstandard compression codec.
- pub(crate) fn new() -> Self {
- Self {}
- }
- }
-
- /// Compression level (1-21) for ZSTD. Choose 1 here for better compression speed.
- const ZSTD_COMPRESSION_LEVEL: i32 = 1;
-
- impl Codec for ZSTDCodec {
- fn decompress(
- &mut self,
- input_buf: &[u8],
- output_buf: &mut Vec<u8>,
- ) -> Result<usize> {
- let mut decoder = zstd::Decoder::new(input_buf)?;
- match io::copy(&mut decoder, output_buf) {
- Ok(n) => Ok(n as usize),
- Err(e) => Err(e.into()),
- }
- }
-
- fn compress(&mut self, input_buf: &[u8], output_buf: &mut Vec<u8>) -> Result<()> {
- let mut encoder = zstd::Encoder::new(output_buf, ZSTD_COMPRESSION_LEVEL)?;
- encoder.write_all(input_buf)?;
- match encoder.finish() {
- Ok(_) => Ok(()),
- Err(e) => Err(e.into()),
- }
- }
- }
-}
-#[cfg(any(feature = "zstd", test))]
-pub use zstd_codec::*;
-
-#[cfg(test)]
-mod tests {
- use super::*;
-
- use crate::util::test_common::*;
-
- fn test_roundtrip(c: CodecType, data: &[u8]) {
- let mut c1 = create_codec(c).unwrap().unwrap();
- let mut c2 = create_codec(c).unwrap().unwrap();
-
- // Compress with c1
- let mut compressed = Vec::new();
- let mut decompressed = Vec::new();
- c1.compress(data, &mut compressed)
- .expect("Error when compressing");
-
- // Decompress with c2
- let mut decompressed_size = c2
- .decompress(compressed.as_slice(), &mut decompressed)
- .expect("Error when decompressing");
- assert_eq!(data.len(), decompressed_size);
- decompressed.truncate(decompressed_size);
- assert_eq!(data, decompressed.as_slice());
-
- compressed.clear();
-
- // Compress with c2
- c2.compress(data, &mut compressed)
- .expect("Error when compressing");
-
- // Decompress with c1
- decompressed_size = c1
- .decompress(compressed.as_slice(), &mut decompressed)
- .expect("Error when decompressing");
- assert_eq!(data.len(), decompressed_size);
- decompressed.truncate(decompressed_size);
- assert_eq!(data, decompressed.as_slice());
- }
-
- fn test_codec(c: CodecType) {
- let sizes = vec![100, 10000, 100000];
- for size in sizes {
- let data = random_bytes(size);
- test_roundtrip(c, &data);
- }
- }
-
- #[test]
- fn test_codec_snappy() {
- test_codec(CodecType::SNAPPY);
- }
-
- #[test]
- fn test_codec_gzip() {
- test_codec(CodecType::GZIP);
- }
-
- #[test]
- fn test_codec_brotli() {
- test_codec(CodecType::BROTLI);
- }
-
- #[test]
- fn test_codec_lz4() {
- test_codec(CodecType::LZ4);
- }
-
- #[test]
- fn test_codec_zstd() {
- test_codec(CodecType::ZSTD);
- }
-}
diff --git a/parquet/src/data_type.rs b/parquet/src/data_type.rs
deleted file mode 100644
index f97df3c..0000000
--- a/parquet/src/data_type.rs
+++ /dev/null
@@ -1,1365 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-//! Data types that connect Parquet physical types with their Rust-specific
-//! representations.
-use std::cmp::Ordering;
-use std::fmt;
-use std::mem;
-use std::ops::{Deref, DerefMut};
-use std::str::from_utf8;
-
-use byteorder::{BigEndian, ByteOrder};
-
-use crate::basic::Type;
-use crate::column::reader::{ColumnReader, ColumnReaderImpl};
-use crate::column::writer::{ColumnWriter, ColumnWriterImpl};
-use crate::errors::{ParquetError, Result};
-use crate::util::{
- bit_util::{from_ne_slice, FromBytes},
- memory::{ByteBuffer, ByteBufferPtr},
-};
-
-/// Rust representation for logical type INT96, value is backed by an array of `u32`.
-/// The type only takes 12 bytes, without extra padding.
-#[derive(Clone, Debug, PartialOrd)]
-pub struct Int96 {
- value: Option<[u32; 3]>,
-}
-
-impl Int96 {
- /// Creates new INT96 type struct with no data set.
- pub fn new() -> Self {
- Self { value: None }
- }
-
- /// Returns underlying data as slice of [`u32`].
- #[inline]
- pub fn data(&self) -> &[u32] {
- self.value
- .as_ref()
- .expect("set_data should have been called")
- }
-
- /// Sets data for this INT96 type.
- #[inline]
- pub fn set_data(&mut self, elem0: u32, elem1: u32, elem2: u32) {
- self.value = Some([elem0, elem1, elem2]);
- }
-
- /// Converts this INT96 into an i64 representing the number of MILLISECONDS since Epoch
- pub fn to_i64(&self) -> i64 {
- const JULIAN_DAY_OF_EPOCH: i64 = 2_440_588;
- const SECONDS_PER_DAY: i64 = 86_400;
- const MILLIS_PER_SECOND: i64 = 1_000;
-
- let day = self.data()[2] as i64;
- let nanoseconds = ((self.data()[1] as i64) << 32) + self.data()[0] as i64;
- let seconds = (day - JULIAN_DAY_OF_EPOCH) * SECONDS_PER_DAY;
-
- seconds * MILLIS_PER_SECOND + nanoseconds / 1_000_000
- }
-}
-
-impl Default for Int96 {
- fn default() -> Self {
- Self { value: None }
- }
-}
-
-impl PartialEq for Int96 {
- fn eq(&self, other: &Int96) -> bool {
- match (&self.value, &other.value) {
- (Some(v1), Some(v2)) => v1 == v2,
- (None, None) => true,
- _ => false,
- }
- }
-}
-
-impl From<Vec<u32>> for Int96 {
- fn from(buf: Vec<u32>) -> Self {
- assert_eq!(buf.len(), 3);
- let mut result = Self::new();
- result.set_data(buf[0], buf[1], buf[2]);
- result
- }
-}
-
-impl fmt::Display for Int96 {
- #[cold]
- fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
- write!(f, "{:?}", self.data())
- }
-}
-
-/// Rust representation for BYTE_ARRAY and FIXED_LEN_BYTE_ARRAY Parquet physical types.
-/// Value is backed by a byte buffer.
-#[derive(Clone, Debug)]
-pub struct ByteArray {
- data: Option<ByteBufferPtr>,
-}
-
-impl PartialOrd for ByteArray {
- fn partial_cmp(&self, other: &ByteArray) -> Option<Ordering> {
- if self.data.is_some() && other.data.is_some() {
- match self.len().cmp(&other.len()) {
- Ordering::Greater => Some(Ordering::Greater),
- Ordering::Less => Some(Ordering::Less),
- Ordering::Equal => {
- for (v1, v2) in self.data().iter().zip(other.data().iter()) {
- match v1.cmp(v2) {
- Ordering::Greater => return Some(Ordering::Greater),
- Ordering::Less => return Some(Ordering::Less),
- _ => {}
- }
- }
- Some(Ordering::Equal)
- }
- }
- } else {
- None
- }
- }
-}
-
-impl ByteArray {
- /// Creates new byte array with no data set.
- #[inline]
- pub fn new() -> Self {
- ByteArray { data: None }
- }
-
- /// Gets length of the underlying byte buffer.
- #[inline]
- pub fn len(&self) -> usize {
- assert!(self.data.is_some());
- self.data.as_ref().unwrap().len()
- }
-
- /// Checks if the underlying buffer is empty.
- #[inline]
- pub fn is_empty(&self) -> bool {
- self.len() == 0
- }
-
- /// Returns slice of data.
- #[inline]
- pub fn data(&self) -> &[u8] {
- self.data
- .as_ref()
- .expect("set_data should have been called")
- .as_ref()
- }
-
- /// Set data from another byte buffer.
- #[inline]
- pub fn set_data(&mut self, data: ByteBufferPtr) {
- self.data = Some(data);
- }
-
- /// Returns `ByteArray` instance with slice of values for a data.
- #[inline]
- pub fn slice(&self, start: usize, len: usize) -> Self {
- Self::from(
- self.data
- .as_ref()
- .expect("set_data should have been called")
- .range(start, len),
- )
- }
-
- pub fn as_utf8(&self) -> Result<&str> {
- self.data
- .as_ref()
- .map(|ptr| ptr.as_ref())
- .ok_or_else(|| general_err!("Can't convert empty byte array to utf8"))
- .and_then(|bytes| from_utf8(bytes).map_err(|e| e.into()))
- }
-}
-
-impl From<Vec<u8>> for ByteArray {
- fn from(buf: Vec<u8>) -> ByteArray {
- Self {
- data: Some(ByteBufferPtr::new(buf)),
- }
- }
-}
-
-impl<'a> From<&'a str> for ByteArray {
- fn from(s: &'a str) -> ByteArray {
- let mut v = Vec::new();
- v.extend_from_slice(s.as_bytes());
- Self {
- data: Some(ByteBufferPtr::new(v)),
- }
- }
-}
-
-impl From<ByteBufferPtr> for ByteArray {
- fn from(ptr: ByteBufferPtr) -> ByteArray {
- Self { data: Some(ptr) }
- }
-}
-
-impl From<ByteBuffer> for ByteArray {
- fn from(mut buf: ByteBuffer) -> ByteArray {
- Self {
- data: Some(buf.consume()),
- }
- }
-}
-
-impl Default for ByteArray {
- fn default() -> Self {
- ByteArray { data: None }
- }
-}
-
-impl PartialEq for ByteArray {
- fn eq(&self, other: &ByteArray) -> bool {
- match (&self.data, &other.data) {
- (Some(d1), Some(d2)) => d1.as_ref() == d2.as_ref(),
- (None, None) => true,
- _ => false,
- }
- }
-}
-
-impl fmt::Display for ByteArray {
- fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
- write!(f, "{:?}", self.data())
- }
-}
-
-/// Wrapper type for performance reasons, this represents `FIXED_LEN_BYTE_ARRAY` but in all other
-/// considerations behaves the same as `ByteArray`
-///
-/// # Performance notes:
-/// This type is a little unfortunate, without it the compiler generates code that takes quite a
-/// big hit on the CPU pipeline. Essentially the previous version stalls awaiting the result of
-/// `T::get_physical_type() == Type::FIXED_LEN_BYTE_ARRAY`.
-///
-/// Its debatable if this is wanted, it is out of spec for what parquet documents as its base
-/// types, although there are code paths in the Rust (and potentially the C++) versions that
-/// warrant this.
-///
-/// With this wrapper type the compiler generates more targetted code paths matching the higher
-/// level logical types, removing the data-hazard from all decoding and encoding paths.
-#[repr(transparent)]
-#[derive(Clone, Debug, Default)]
-pub struct FixedLenByteArray(ByteArray);
-
-impl PartialEq for FixedLenByteArray {
- fn eq(&self, other: &FixedLenByteArray) -> bool {
- self.0.eq(&other.0)
- }
-}
-
-impl PartialEq<ByteArray> for FixedLenByteArray {
- fn eq(&self, other: &ByteArray) -> bool {
- self.0.eq(other)
- }
-}
-
-impl PartialEq<FixedLenByteArray> for ByteArray {
- fn eq(&self, other: &FixedLenByteArray) -> bool {
- self.eq(&other.0)
- }
-}
-
-impl fmt::Display for FixedLenByteArray {
- fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
- self.0.fmt(f)
- }
-}
-
-impl PartialOrd for FixedLenByteArray {
- fn partial_cmp(&self, other: &FixedLenByteArray) -> Option<Ordering> {
- self.0.partial_cmp(&other.0)
- }
-}
-
-impl PartialOrd<FixedLenByteArray> for ByteArray {
- fn partial_cmp(&self, other: &FixedLenByteArray) -> Option<Ordering> {
- self.partial_cmp(&other.0)
- }
-}
-
-impl PartialOrd<ByteArray> for FixedLenByteArray {
- fn partial_cmp(&self, other: &ByteArray) -> Option<Ordering> {
- self.0.partial_cmp(other)
- }
-}
-
-impl Deref for FixedLenByteArray {
- type Target = ByteArray;
-
- fn deref(&self) -> &Self::Target {
- &self.0
- }
-}
-
-impl DerefMut for FixedLenByteArray {
- fn deref_mut(&mut self) -> &mut Self::Target {
- &mut self.0
- }
-}
-
-impl From<ByteArray> for FixedLenByteArray {
- fn from(other: ByteArray) -> Self {
- Self(other)
- }
-}
-
-impl From<FixedLenByteArray> for ByteArray {
- fn from(other: FixedLenByteArray) -> Self {
- other.0
- }
-}
-
-/// Rust representation for Decimal values.
-///
-/// This is not a representation of Parquet physical type, but rather a wrapper for
-/// DECIMAL logical type, and serves as container for raw parts of decimal values:
-/// unscaled value in bytes, precision and scale.
-#[derive(Clone, Debug)]
-pub enum Decimal {
- /// Decimal backed by `i32`.
- Int32 {
- value: [u8; 4],
- precision: i32,
- scale: i32,
- },
- /// Decimal backed by `i64`.
- Int64 {
- value: [u8; 8],
- precision: i32,
- scale: i32,
- },
- /// Decimal backed by byte array.
- Bytes {
- value: ByteArray,
- precision: i32,
- scale: i32,
- },
-}
-
-impl Decimal {
- /// Creates new decimal value from `i32`.
- pub fn from_i32(value: i32, precision: i32, scale: i32) -> Self {
- let mut bytes = [0; 4];
- BigEndian::write_i32(&mut bytes, value);
- Decimal::Int32 {
- value: bytes,
- precision,
- scale,
- }
- }
-
- /// Creates new decimal value from `i64`.
- pub fn from_i64(value: i64, precision: i32, scale: i32) -> Self {
- let mut bytes = [0; 8];
- BigEndian::write_i64(&mut bytes, value);
- Decimal::Int64 {
- value: bytes,
- precision,
- scale,
- }
- }
-
- /// Creates new decimal value from `ByteArray`.
- pub fn from_bytes(value: ByteArray, precision: i32, scale: i32) -> Self {
- Decimal::Bytes {
- value,
- precision,
- scale,
- }
- }
-
- /// Returns bytes of unscaled value.
- pub fn data(&self) -> &[u8] {
- match *self {
- Decimal::Int32 { ref value, .. } => value,
- Decimal::Int64 { ref value, .. } => value,
- Decimal::Bytes { ref value, .. } => value.data(),
- }
- }
-
- /// Returns decimal precision.
- pub fn precision(&self) -> i32 {
- match *self {
- Decimal::Int32 { precision, .. } => precision,
- Decimal::Int64 { precision, .. } => precision,
- Decimal::Bytes { precision, .. } => precision,
- }
- }
-
- /// Returns decimal scale.
- pub fn scale(&self) -> i32 {
- match *self {
- Decimal::Int32 { scale, .. } => scale,
- Decimal::Int64 { scale, .. } => scale,
- Decimal::Bytes { scale, .. } => scale,
- }
- }
-}
-
-impl Default for Decimal {
- fn default() -> Self {
- Self::from_i32(0, 0, 0)
- }
-}
-
-impl PartialEq for Decimal {
- fn eq(&self, other: &Decimal) -> bool {
- self.precision() == other.precision()
- && self.scale() == other.scale()
- && self.data() == other.data()
- }
-}
-
-/// Converts an instance of data type to a slice of bytes as `u8`.
-pub trait AsBytes {
- /// Returns slice of bytes for this data type.
- fn as_bytes(&self) -> &[u8];
-}
-
-/// Converts an slice of a data type to a slice of bytes.
-pub trait SliceAsBytes: Sized {
- /// Returns slice of bytes for a slice of this data type.
- fn slice_as_bytes(self_: &[Self]) -> &[u8];
- /// Return the internal representation as a mutable slice
- ///
- /// # Safety
- /// If modified you are _required_ to ensure the internal representation
- /// is valid and correct for the actual raw data
- unsafe fn slice_as_bytes_mut(self_: &mut [Self]) -> &mut [u8];
-}
-
-impl AsBytes for [u8] {
- fn as_bytes(&self) -> &[u8] {
- self
- }
-}
-
-macro_rules! gen_as_bytes {
- ($source_ty:ident) => {
- impl AsBytes for $source_ty {
- #[allow(clippy::size_of_in_element_count)]
- fn as_bytes(&self) -> &[u8] {
- unsafe {
- std::slice::from_raw_parts(
- self as *const $source_ty as *const u8,
- std::mem::size_of::<$source_ty>(),
- )
- }
- }
- }
-
- impl SliceAsBytes for $source_ty {
- #[inline]
- #[allow(clippy::size_of_in_element_count)]
- fn slice_as_bytes(self_: &[Self]) -> &[u8] {
- unsafe {
- std::slice::from_raw_parts(
- self_.as_ptr() as *const u8,
- std::mem::size_of::<$source_ty>() * self_.len(),
- )
- }
- }
-
- #[inline]
- #[allow(clippy::size_of_in_element_count)]
- unsafe fn slice_as_bytes_mut(self_: &mut [Self]) -> &mut [u8] {
- std::slice::from_raw_parts_mut(
- self_.as_mut_ptr() as *mut u8,
- std::mem::size_of::<$source_ty>() * self_.len(),
- )
- }
- }
- };
-}
-
-gen_as_bytes!(i8);
-gen_as_bytes!(i16);
-gen_as_bytes!(i32);
-gen_as_bytes!(i64);
-gen_as_bytes!(u8);
-gen_as_bytes!(u16);
-gen_as_bytes!(u32);
-gen_as_bytes!(u64);
-gen_as_bytes!(f32);
-gen_as_bytes!(f64);
-
-macro_rules! unimplemented_slice_as_bytes {
- ($ty: ty) => {
- impl SliceAsBytes for $ty {
- fn slice_as_bytes(_self: &[Self]) -> &[u8] {
- unimplemented!()
- }
-
- unsafe fn slice_as_bytes_mut(_self: &mut [Self]) -> &mut [u8] {
- unimplemented!()
- }
- }
- };
-}
-
-// TODO - Can Int96 and bool be implemented in these terms?
-unimplemented_slice_as_bytes!(Int96);
-unimplemented_slice_as_bytes!(bool);
-unimplemented_slice_as_bytes!(ByteArray);
-unimplemented_slice_as_bytes!(FixedLenByteArray);
-
-impl AsBytes for bool {
- fn as_bytes(&self) -> &[u8] {
- unsafe { std::slice::from_raw_parts(self as *const bool as *const u8, 1) }
- }
-}
-
-impl AsBytes for Int96 {
- fn as_bytes(&self) -> &[u8] {
- unsafe {
- std::slice::from_raw_parts(self.data() as *const [u32] as *const u8, 12)
- }
- }
-}
-
-impl AsBytes for ByteArray {
- fn as_bytes(&self) -> &[u8] {
- self.data()
- }
-}
-
-impl AsBytes for FixedLenByteArray {
- fn as_bytes(&self) -> &[u8] {
- self.data()
- }
-}
-
-impl AsBytes for Decimal {
- fn as_bytes(&self) -> &[u8] {
- self.data()
- }
-}
-
-impl AsBytes for Vec<u8> {
- fn as_bytes(&self) -> &[u8] {
- self.as_slice()
- }
-}
-
-impl<'a> AsBytes for &'a str {
- fn as_bytes(&self) -> &[u8] {
- (self as &str).as_bytes()
- }
-}
-
-impl AsBytes for str {
- fn as_bytes(&self) -> &[u8] {
- (self as &str).as_bytes()
- }
-}
-
-pub(crate) mod private {
- use crate::encodings::decoding::PlainDecoderDetails;
- use crate::util::bit_util::{BitReader, BitWriter};
- use crate::util::memory::ByteBufferPtr;
-
- use byteorder::ByteOrder;
- use std::convert::TryInto;
-
- use super::{ParquetError, Result, SliceAsBytes};
-
- pub type BitIndex = u64;
-
- /// Sealed trait to start to remove specialisation from implementations
- ///
- /// This is done to force the associated value type to be unimplementable outside of this
- /// crate, and thus hint to the type system (and end user) traits are public for the contract
- /// and not for extension.
- pub trait ParquetValueType:
- std::cmp::PartialEq
- + std::fmt::Debug
- + std::fmt::Display
- + std::default::Default
- + std::clone::Clone
- + super::AsBytes
- + super::FromBytes
- + super::SliceAsBytes
- + PartialOrd
- {
- /// Encode the value directly from a higher level encoder
- fn encode<W: std::io::Write>(
- values: &[Self],
- writer: &mut W,
- bit_writer: &mut BitWriter,
- ) -> Result<()>;
-
- /// Establish the data that will be decoded in a buffer
- fn set_data(
- decoder: &mut PlainDecoderDetails,
- data: ByteBufferPtr,
- num_values: usize,
- );
-
- /// Decode the value from a given buffer for a higher level decoder
- fn decode(
- buffer: &mut [Self],
- decoder: &mut PlainDecoderDetails,
- ) -> Result<usize>;
-
- /// Return the encoded size for a type
- fn dict_encoding_size(&self) -> (usize, usize) {
- (std::mem::size_of::<Self>(), 1)
- }
-
- /// Return the value as i64 if possible
- ///
- /// This is essentially the same as `std::convert::TryInto<i64>` but can
- /// implemented for `f32` and `f64`, types that would fail orphan rules
- fn as_i64(&self) -> Result<i64> {
- Err(general_err!("Type cannot be converted to i64"))
- }
-
- /// Return the value as u64 if possible
- ///
- /// This is essentially the same as `std::convert::TryInto<u64>` but can
- /// implemented for `f32` and `f64`, types that would fail orphan rules
- fn as_u64(&self) -> Result<u64> {
- self.as_i64()
- .map_err(|_| general_err!("Type cannot be converted to u64"))
- .map(|x| x as u64)
- }
-
- /// Return the value as an Any to allow for downcasts without transmutation
- fn as_any(&self) -> &dyn std::any::Any;
-
- /// Return the value as an mutable Any to allow for downcasts without transmutation
- fn as_mut_any(&mut self) -> &mut dyn std::any::Any;
- }
-
- impl ParquetValueType for bool {
- #[inline]
- fn encode<W: std::io::Write>(
- values: &[Self],
- _: &mut W,
- bit_writer: &mut BitWriter,
- ) -> Result<()> {
- if bit_writer.bytes_written() + values.len() / 8 >= bit_writer.capacity() {
- bit_writer.extend(256);
- }
- for value in values {
- if !bit_writer.put_value(*value as u64, 1) {
- return Err(ParquetError::EOF(
- "unable to put boolean value".to_string(),
- ));
- }
- }
- Ok(())
- }
-
- #[inline]
- fn set_data(
- decoder: &mut PlainDecoderDetails,
- data: ByteBufferPtr,
- num_values: usize,
- ) {
- decoder.bit_reader.replace(BitReader::new(data));
- decoder.num_values = num_values;
- }
-
- #[inline]
- fn decode(
- buffer: &mut [Self],
- decoder: &mut PlainDecoderDetails,
- ) -> Result<usize> {
- let bit_reader = decoder.bit_reader.as_mut().unwrap();
- let num_values = std::cmp::min(buffer.len(), decoder.num_values);
- let values_read = bit_reader.get_batch(&mut buffer[..num_values], 1);
- decoder.num_values -= values_read;
- Ok(values_read)
- }
-
- #[inline]
- fn as_i64(&self) -> Result<i64> {
- Ok(*self as i64)
- }
-
- #[inline]
- fn as_any(&self) -> &dyn std::any::Any {
- self
- }
-
- #[inline]
- fn as_mut_any(&mut self) -> &mut dyn std::any::Any {
- self
- }
- }
-
- /// Hopelessly unsafe function that emulates `num::as_ne_bytes`
- ///
- /// It is not recommended to use this outside of this private module as, while it
- /// _should_ work for primitive values, it is little better than a transmutation
- /// and can act as a backdoor into mis-interpreting types as arbitary byte slices
- #[inline]
- fn as_raw<'a, T>(value: *const T) -> &'a [u8] {
- unsafe {
- let value = value as *const u8;
- std::slice::from_raw_parts(value, std::mem::size_of::<T>())
- }
- }
-
- macro_rules! impl_from_raw {
- ($ty: ty, $self: ident => $as_i64: block) => {
- impl ParquetValueType for $ty {
- #[inline]
- fn encode<W: std::io::Write>(values: &[Self], writer: &mut W, _: &mut BitWriter) -> Result<()> {
- let raw = unsafe {
- std::slice::from_raw_parts(
- values.as_ptr() as *const u8,
- std::mem::size_of::<$ty>() * values.len(),
- )
- };
- writer.write_all(raw)?;
-
- Ok(())
- }
-
- #[inline]
- fn set_data(decoder: &mut PlainDecoderDetails, data: ByteBufferPtr, num_values: usize) {
- decoder.data.replace(data);
- decoder.start = 0;
- decoder.num_values = num_values;
- }
-
- #[inline]
- fn decode(buffer: &mut [Self], decoder: &mut PlainDecoderDetails) -> Result<usize> {
- let data = decoder.data.as_ref().expect("set_data should have been called");
- let num_values = std::cmp::min(buffer.len(), decoder.num_values);
- let bytes_left = data.len() - decoder.start;
- let bytes_to_decode = std::mem::size_of::<Self>() * num_values;
-
- if bytes_left < bytes_to_decode {
- return Err(eof_err!("Not enough bytes to decode"));
- }
-
- // SAFETY: Raw types should be as per the standard rust bit-vectors
- unsafe {
- let raw_buffer = &mut Self::slice_as_bytes_mut(buffer)[..bytes_to_decode];
- raw_buffer.copy_from_slice(data.range(decoder.start, bytes_to_decode).as_ref());
- };
- decoder.start += bytes_to_decode;
- decoder.num_values -= num_values;
-
- Ok(num_values)
- }
-
- #[inline]
- fn as_i64(&$self) -> Result<i64> {
- $as_i64
- }
-
- #[inline]
- fn as_any(&self) -> &dyn std::any::Any {
- self
- }
-
- #[inline]
- fn as_mut_any(&mut self) -> &mut dyn std::any::Any {
- self
- }
- }
- }
- }
-
- impl_from_raw!(i32, self => { Ok(*self as i64) });
- impl_from_raw!(i64, self => { Ok(*self) });
- impl_from_raw!(f32, self => { Err(general_err!("Type cannot be converted to i64")) });
- impl_from_raw!(f64, self => { Err(general_err!("Type cannot be converted to i64")) });
-
- impl ParquetValueType for super::Int96 {
- #[inline]
- fn encode<W: std::io::Write>(
- values: &[Self],
- writer: &mut W,
- _: &mut BitWriter,
- ) -> Result<()> {
- for value in values {
- let raw = unsafe {
- std::slice::from_raw_parts(
- value.data() as *const [u32] as *const u8,
- 12,
- )
- };
- writer.write_all(raw)?;
- }
- Ok(())
- }
-
- #[inline]
- fn set_data(
- decoder: &mut PlainDecoderDetails,
- data: ByteBufferPtr,
- num_values: usize,
- ) {
- decoder.data.replace(data);
- decoder.start = 0;
- decoder.num_values = num_values;
- }
-
- #[inline]
- fn decode(
- buffer: &mut [Self],
- decoder: &mut PlainDecoderDetails,
- ) -> Result<usize> {
- // TODO - Remove the duplication between this and the general slice method
- let data = decoder
- .data
- .as_ref()
- .expect("set_data should have been called");
- let num_values = std::cmp::min(buffer.len(), decoder.num_values);
- let bytes_left = data.len() - decoder.start;
- let bytes_to_decode = 12 * num_values;
-
- if bytes_left < bytes_to_decode {
- return Err(eof_err!("Not enough bytes to decode"));
- }
-
- let data_range = data.range(decoder.start, bytes_to_decode);
- let bytes: &[u8] = data_range.data();
- decoder.start += bytes_to_decode;
-
- let mut pos = 0; // position in byte array
- for i in 0..num_values {
- let elem0 = byteorder::LittleEndian::read_u32(&bytes[pos..pos + 4]);
- let elem1 = byteorder::LittleEndian::read_u32(&bytes[pos + 4..pos + 8]);
- let elem2 = byteorder::LittleEndian::read_u32(&bytes[pos + 8..pos + 12]);
-
- buffer[i]
- .as_mut_any()
- .downcast_mut::<Self>()
- .unwrap()
- .set_data(elem0, elem1, elem2);
-
- pos += 12;
- }
- decoder.num_values -= num_values;
-
- Ok(num_values)
- }
-
- #[inline]
- fn as_any(&self) -> &dyn std::any::Any {
- self
- }
-
- #[inline]
- fn as_mut_any(&mut self) -> &mut dyn std::any::Any {
- self
- }
- }
-
- // TODO - Why does macro importing fail?
- /// Reads `$size` of bytes from `$src`, and reinterprets them as type `$ty`, in
- /// little-endian order. `$ty` must implement the `Default` trait. Otherwise this won't
- /// compile.
- /// This is copied and modified from byteorder crate.
- macro_rules! read_num_bytes {
- ($ty:ty, $size:expr, $src:expr) => {{
- assert!($size <= $src.len());
- let mut buffer =
- <$ty as $crate::util::bit_util::FromBytes>::Buffer::default();
- buffer.as_mut()[..$size].copy_from_slice(&$src[..$size]);
- <$ty>::from_ne_bytes(buffer)
- }};
- }
-
- impl ParquetValueType for super::ByteArray {
- #[inline]
- fn encode<W: std::io::Write>(
- values: &[Self],
- writer: &mut W,
- _: &mut BitWriter,
- ) -> Result<()> {
- for value in values {
- let len: u32 = value.len().try_into().unwrap();
- writer.write_all(&len.to_ne_bytes())?;
- let raw = value.data();
- writer.write_all(raw)?;
- }
- Ok(())
- }
-
- #[inline]
- fn set_data(
- decoder: &mut PlainDecoderDetails,
- data: ByteBufferPtr,
- num_values: usize,
- ) {
- decoder.data.replace(data);
- decoder.start = 0;
- decoder.num_values = num_values;
- }
-
- #[inline]
- fn decode(
- buffer: &mut [Self],
- decoder: &mut PlainDecoderDetails,
- ) -> Result<usize> {
- let data = decoder
- .data
- .as_mut()
- .expect("set_data should have been called");
- let num_values = std::cmp::min(buffer.len(), decoder.num_values);
- for i in 0..num_values {
- let len: usize =
- read_num_bytes!(u32, 4, data.start_from(decoder.start).as_ref())
- as usize;
- decoder.start += std::mem::size_of::<u32>();
-
- if data.len() < decoder.start + len {
- return Err(eof_err!("Not enough bytes to decode"));
- }
-
- let val: &mut Self = buffer[i].as_mut_any().downcast_mut().unwrap();
-
- val.set_data(data.range(decoder.start, len));
- decoder.start += len;
- }
- decoder.num_values -= num_values;
-
- Ok(num_values)
- }
-
- #[inline]
- fn dict_encoding_size(&self) -> (usize, usize) {
- (std::mem::size_of::<u32>(), self.len())
- }
-
- #[inline]
- fn as_any(&self) -> &dyn std::any::Any {
- self
- }
-
- #[inline]
- fn as_mut_any(&mut self) -> &mut dyn std::any::Any {
- self
- }
- }
-
- impl ParquetValueType for super::FixedLenByteArray {
- #[inline]
- fn encode<W: std::io::Write>(
- values: &[Self],
- writer: &mut W,
- _: &mut BitWriter,
- ) -> Result<()> {
- for value in values {
- let raw = value.data();
- writer.write_all(raw)?;
- }
- Ok(())
- }
-
- #[inline]
- fn set_data(
- decoder: &mut PlainDecoderDetails,
- data: ByteBufferPtr,
- num_values: usize,
- ) {
- decoder.data.replace(data);
- decoder.start = 0;
- decoder.num_values = num_values;
- }
-
- #[inline]
- fn decode(
- buffer: &mut [Self],
- decoder: &mut PlainDecoderDetails,
- ) -> Result<usize> {
- assert!(decoder.type_length > 0);
-
- let data = decoder
- .data
- .as_mut()
- .expect("set_data should have been called");
- let num_values = std::cmp::min(buffer.len(), decoder.num_values);
- for i in 0..num_values {
- let len = decoder.type_length as usize;
-
- if data.len() < decoder.start + len {
- return Err(eof_err!("Not enough bytes to decode"));
- }
-
- let val: &mut Self = buffer[i].as_mut_any().downcast_mut().unwrap();
-
- val.set_data(data.range(decoder.start, len));
- decoder.start += len;
- }
- decoder.num_values -= num_values;
-
- Ok(num_values)
- }
-
- #[inline]
- fn dict_encoding_size(&self) -> (usize, usize) {
- (std::mem::size_of::<u32>(), self.len())
- }
-
- #[inline]
- fn as_any(&self) -> &dyn std::any::Any {
- self
- }
-
- #[inline]
- fn as_mut_any(&mut self) -> &mut dyn std::any::Any {
- self
- }
- }
-}
-
-/// Contains the Parquet physical type information as well as the Rust primitive type
-/// presentation.
-pub trait DataType: 'static {
- type T: private::ParquetValueType;
-
- /// Returns Parquet physical type.
- fn get_physical_type() -> Type;
-
- /// Returns size in bytes for Rust representation of the physical type.
- fn get_type_size() -> usize;
-
- fn get_column_reader(column_writer: ColumnReader) -> Option<ColumnReaderImpl<Self>>
- where
- Self: Sized;
-
- fn get_column_writer(column_writer: ColumnWriter) -> Option<ColumnWriterImpl<Self>>
- where
- Self: Sized;
-
- fn get_column_writer_ref(
- column_writer: &ColumnWriter,
- ) -> Option<&ColumnWriterImpl<Self>>
- where
- Self: Sized;
-
- fn get_column_writer_mut(
- column_writer: &mut ColumnWriter,
- ) -> Option<&mut ColumnWriterImpl<Self>>
- where
- Self: Sized;
-}
-
-// Workaround bug in specialization
-pub trait SliceAsBytesDataType: DataType
-where
- Self::T: SliceAsBytes,
-{
-}
-
-impl<T> SliceAsBytesDataType for T
-where
- T: DataType,
- <T as DataType>::T: SliceAsBytes,
-{
-}
-
-macro_rules! make_type {
- ($name:ident, $physical_ty:path, $reader_ident: ident, $writer_ident: ident, $native_ty:ty, $size:expr) => {
- #[derive(Clone)]
- pub struct $name {}
-
- impl DataType for $name {
- type T = $native_ty;
-
- fn get_physical_type() -> Type {
- $physical_ty
- }
-
- fn get_type_size() -> usize {
- $size
- }
-
- fn get_column_reader(
- column_writer: ColumnReader,
- ) -> Option<ColumnReaderImpl<Self>> {
- match column_writer {
- ColumnReader::$reader_ident(w) => Some(w),
- _ => None,
- }
- }
-
- fn get_column_writer(
- column_writer: ColumnWriter,
- ) -> Option<ColumnWriterImpl<Self>> {
- match column_writer {
- ColumnWriter::$writer_ident(w) => Some(w),
- _ => None,
- }
- }
-
- fn get_column_writer_ref(
- column_writer: &ColumnWriter,
- ) -> Option<&ColumnWriterImpl<Self>> {
- match column_writer {
- ColumnWriter::$writer_ident(w) => Some(w),
- _ => None,
- }
- }
-
- fn get_column_writer_mut(
- column_writer: &mut ColumnWriter,
- ) -> Option<&mut ColumnWriterImpl<Self>> {
- match column_writer {
- ColumnWriter::$writer_ident(w) => Some(w),
- _ => None,
- }
- }
- }
- };
-}
-
-// Generate struct definitions for all physical types
-
-make_type!(
- BoolType,
- Type::BOOLEAN,
- BoolColumnReader,
- BoolColumnWriter,
- bool,
- 1
-);
-make_type!(
- Int32Type,
- Type::INT32,
- Int32ColumnReader,
- Int32ColumnWriter,
- i32,
- 4
-);
-make_type!(
- Int64Type,
- Type::INT64,
- Int64ColumnReader,
- Int64ColumnWriter,
- i64,
- 8
-);
-make_type!(
- Int96Type,
- Type::INT96,
- Int96ColumnReader,
- Int96ColumnWriter,
- Int96,
- mem::size_of::<Int96>()
-);
-make_type!(
- FloatType,
- Type::FLOAT,
- FloatColumnReader,
- FloatColumnWriter,
- f32,
- 4
-);
-make_type!(
- DoubleType,
- Type::DOUBLE,
- DoubleColumnReader,
- DoubleColumnWriter,
- f64,
- 8
-);
-make_type!(
- ByteArrayType,
- Type::BYTE_ARRAY,
- ByteArrayColumnReader,
- ByteArrayColumnWriter,
- ByteArray,
- mem::size_of::<ByteArray>()
-);
-make_type!(
- FixedLenByteArrayType,
- Type::FIXED_LEN_BYTE_ARRAY,
- FixedLenByteArrayColumnReader,
- FixedLenByteArrayColumnWriter,
- FixedLenByteArray,
- mem::size_of::<FixedLenByteArray>()
-);
-
-impl FromBytes for Int96 {
- type Buffer = [u8; 12];
- fn from_le_bytes(_bs: Self::Buffer) -> Self {
- unimplemented!()
- }
- fn from_be_bytes(_bs: Self::Buffer) -> Self {
- unimplemented!()
- }
- fn from_ne_bytes(bs: Self::Buffer) -> Self {
- let mut i = Int96::new();
- i.set_data(
- from_ne_slice(&bs[0..4]),
- from_ne_slice(&bs[4..8]),
- from_ne_slice(&bs[8..12]),
- );
- i
- }
-}
-
-// FIXME Needed to satisfy the constraint of many decoding functions but ByteArray does not
-// appear to actual be converted directly from bytes
-impl FromBytes for ByteArray {
- type Buffer = [u8; 8];
- fn from_le_bytes(_bs: Self::Buffer) -> Self {
- unreachable!()
- }
- fn from_be_bytes(_bs: Self::Buffer) -> Self {
- unreachable!()
- }
- fn from_ne_bytes(bs: Self::Buffer) -> Self {
- ByteArray::from(bs.to_vec())
- }
-}
-
-impl FromBytes for FixedLenByteArray {
- type Buffer = [u8; 8];
-
- fn from_le_bytes(_bs: Self::Buffer) -> Self {
- unreachable!()
- }
- fn from_be_bytes(_bs: Self::Buffer) -> Self {
- unreachable!()
- }
- fn from_ne_bytes(bs: Self::Buffer) -> Self {
- Self(ByteArray::from(bs.to_vec()))
- }
-}
-
-/// Macro to reduce repetition in making type assertions on the physical type against `T`
-macro_rules! ensure_phys_ty {
- ($($ty: pat)|+ , $err: literal) => {
- match T::get_physical_type() {
- $($ty => (),)*
- _ => panic!($err),
- };
- }
-}
-
-#[cfg(test)]
-#[allow(clippy::float_cmp, clippy::approx_constant)]
-mod tests {
- use super::*;
-
- #[test]
- #[allow(clippy::string_lit_as_bytes)]
- fn test_as_bytes() {
- assert_eq!(false.as_bytes(), &[0]);
- assert_eq!(true.as_bytes(), &[1]);
- assert_eq!(7_i32.as_bytes(), &[7, 0, 0, 0]);
- assert_eq!(555_i32.as_bytes(), &[43, 2, 0, 0]);
- assert_eq!(555_u32.as_bytes(), &[43, 2, 0, 0]);
- assert_eq!(i32::max_value().as_bytes(), &[255, 255, 255, 127]);
- assert_eq!(i32::min_value().as_bytes(), &[0, 0, 0, 128]);
- assert_eq!(7_i64.as_bytes(), &[7, 0, 0, 0, 0, 0, 0, 0]);
- assert_eq!(555_i64.as_bytes(), &[43, 2, 0, 0, 0, 0, 0, 0]);
- assert_eq!(
- (i64::max_value()).as_bytes(),
- &[255, 255, 255, 255, 255, 255, 255, 127]
- );
- assert_eq!((i64::min_value()).as_bytes(), &[0, 0, 0, 0, 0, 0, 0, 128]);
- assert_eq!(3.14_f32.as_bytes(), &[195, 245, 72, 64]);
- assert_eq!(3.14_f64.as_bytes(), &[31, 133, 235, 81, 184, 30, 9, 64]);
- assert_eq!("hello".as_bytes(), &[b'h', b'e', b'l', b'l', b'o']);
- assert_eq!(
- Vec::from("hello".as_bytes()).as_bytes(),
- &[b'h', b'e', b'l', b'l', b'o']
- );
-
- // Test Int96
- let i96 = Int96::from(vec![1, 2, 3]);
- assert_eq!(i96.as_bytes(), &[1, 0, 0, 0, 2, 0, 0, 0, 3, 0, 0, 0]);
-
- // Test ByteArray
- let ba = ByteArray::from(vec![1, 2, 3]);
- assert_eq!(ba.as_bytes(), &[1, 2, 3]);
-
- // Test Decimal
- let decimal = Decimal::from_i32(123, 5, 2);
- assert_eq!(decimal.as_bytes(), &[0, 0, 0, 123]);
- let decimal = Decimal::from_i64(123, 5, 2);
- assert_eq!(decimal.as_bytes(), &[0, 0, 0, 0, 0, 0, 0, 123]);
- let decimal = Decimal::from_bytes(ByteArray::from(vec![1, 2, 3]), 5, 2);
- assert_eq!(decimal.as_bytes(), &[1, 2, 3]);
- }
-
- #[test]
- fn test_int96_from() {
- assert_eq!(
- Int96::from(vec![1, 12345, 1234567890]).data(),
- &[1, 12345, 1234567890]
- );
- }
-
- #[test]
- fn test_byte_array_from() {
- assert_eq!(
- ByteArray::from(vec![b'A', b'B', b'C']).data(),
- &[b'A', b'B', b'C']
- );
- assert_eq!(ByteArray::from("ABC").data(), &[b'A', b'B', b'C']);
- assert_eq!(
- ByteArray::from(ByteBufferPtr::new(vec![1u8, 2u8, 3u8, 4u8, 5u8])).data(),
- &[1u8, 2u8, 3u8, 4u8, 5u8]
- );
- let mut buf = ByteBuffer::new();
- buf.set_data(vec![6u8, 7u8, 8u8, 9u8, 10u8]);
- assert_eq!(ByteArray::from(buf).data(), &[6u8, 7u8, 8u8, 9u8, 10u8]);
- }
-
- #[test]
- fn test_decimal_partial_eq() {
- assert_eq!(Decimal::default(), Decimal::from_i32(0, 0, 0));
- assert_eq!(Decimal::from_i32(222, 5, 2), Decimal::from_i32(222, 5, 2));
- assert_eq!(
- Decimal::from_bytes(ByteArray::from(vec![0, 0, 0, 3]), 5, 2),
- Decimal::from_i32(3, 5, 2)
- );
-
- assert!(Decimal::from_i32(222, 5, 2) != Decimal::from_i32(111, 5, 2));
- assert!(Decimal::from_i32(222, 5, 2) != Decimal::from_i32(222, 6, 2));
- assert!(Decimal::from_i32(222, 5, 2) != Decimal::from_i32(222, 5, 3));
-
- assert!(Decimal::from_i64(222, 5, 2) != Decimal::from_i32(222, 5, 2));
- }
-
- #[test]
- fn test_byte_array_ord() {
- let ba1 = ByteArray::from(vec![1, 2, 3]);
- let ba11 = ByteArray::from(vec![1, 2, 3]);
- let ba2 = ByteArray::from(vec![3, 4]);
- let ba3 = ByteArray::from(vec![1, 2, 4]);
- let ba4 = ByteArray::from(vec![]);
- let ba5 = ByteArray::from(vec![2, 2, 3]);
-
- assert!(ba1 > ba2);
- assert!(ba3 > ba1);
- assert!(ba1 > ba4);
- assert_eq!(ba1, ba11);
- assert!(ba5 > ba1);
- }
-}
diff --git a/parquet/src/encodings/decoding.rs b/parquet/src/encodings/decoding.rs
deleted file mode 100644
index e83e277..0000000
--- a/parquet/src/encodings/decoding.rs
+++ /dev/null
@@ -1,1386 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-//! Contains all supported decoders for Parquet.
-
-use std::{cmp, marker::PhantomData, mem};
-
-use super::rle::RleDecoder;
-
-use crate::basic::*;
-use crate::data_type::private::*;
-use crate::data_type::*;
-use crate::errors::{ParquetError, Result};
-use crate::schema::types::ColumnDescPtr;
-use crate::util::{
- bit_util::{self, BitReader, FromBytes},
- memory::{ByteBuffer, ByteBufferPtr},
-};
-
-// ----------------------------------------------------------------------
-// Decoders
-
-/// A Parquet decoder for the data type `T`.
-pub trait Decoder<T: DataType> {
- /// Sets the data to decode to be `data`, which should contain `num_values` of values
- /// to decode.
- fn set_data(&mut self, data: ByteBufferPtr, num_values: usize) -> Result<()>;
-
- /// Consumes values from this decoder and write the results to `buffer`. This will try
- /// to fill up `buffer`.
- ///
- /// Returns the actual number of values decoded, which should be equal to
- /// `buffer.len()` unless the remaining number of values is less than
- /// `buffer.len()`.
- fn get(&mut self, buffer: &mut [T::T]) -> Result<usize>;
-
- /// Consume values from this decoder and write the results to `buffer`, leaving
- /// "spaces" for null values.
- ///
- /// `null_count` is the number of nulls we expect to see in `buffer`, after reading.
- /// `valid_bits` stores the valid bit for each value in the buffer. It should contain
- /// at least number of bits that equal to `buffer.len()`.
- ///
- /// Returns the actual number of values decoded.
- ///
- /// # Panics
- ///
- /// Panics if `null_count` is greater than `buffer.len()`.
- fn get_spaced(
- &mut self,
- buffer: &mut [T::T],
- null_count: usize,
- valid_bits: &[u8],
- ) -> Result<usize> {
- assert!(buffer.len() >= null_count);
-
- // TODO: check validity of the input arguments?
- if null_count == 0 {
- return self.get(buffer);
- }
-
- let num_values = buffer.len();
- let values_to_read = num_values - null_count;
- let values_read = self.get(buffer)?;
- if values_read != values_to_read {
- return Err(general_err!(
- "Number of values read: {}, doesn't match expected: {}",
- values_read,
- values_to_read
- ));
- }
- let mut values_to_move = values_read;
- for i in (0..num_values).rev() {
- if bit_util::get_bit(valid_bits, i) {
- values_to_move -= 1;
- buffer.swap(i, values_to_move);
- }
- }
-
- Ok(num_values)
- }
-
- /// Returns the number of values left in this decoder stream.
- fn values_left(&self) -> usize;
-
- /// Returns the encoding for this decoder.
- fn encoding(&self) -> Encoding;
-}
-
-/// Gets a decoder for the column descriptor `descr` and encoding type `encoding`.
-///
-/// NOTE: the primitive type in `descr` MUST match the data type `T`, otherwise
-/// disastrous consequence could occur.
-pub fn get_decoder<T: DataType>(
- descr: ColumnDescPtr,
- encoding: Encoding,
-) -> Result<Box<dyn Decoder<T>>> {
- let decoder: Box<dyn Decoder<T>> = match encoding {
- Encoding::PLAIN => Box::new(PlainDecoder::new(descr.type_length())),
- Encoding::RLE_DICTIONARY | Encoding::PLAIN_DICTIONARY => {
- return Err(general_err!(
- "Cannot initialize this encoding through this function"
- ));
- }
- Encoding::RLE => Box::new(RleValueDecoder::new()),
- Encoding::DELTA_BINARY_PACKED => Box::new(DeltaBitPackDecoder::new()),
- Encoding::DELTA_LENGTH_BYTE_ARRAY => Box::new(DeltaLengthByteArrayDecoder::new()),
- Encoding::DELTA_BYTE_ARRAY => Box::new(DeltaByteArrayDecoder::new()),
- e => return Err(nyi_err!("Encoding {} is not supported", e)),
- };
- Ok(decoder)
-}
-
-// ----------------------------------------------------------------------
-// PLAIN Decoding
-
-#[derive(Default)]
-pub struct PlainDecoderDetails {
- // The remaining number of values in the byte array
- pub(crate) num_values: usize,
-
- // The current starting index in the byte array. Not used when `T` is bool.
- pub(crate) start: usize,
-
- // The length for the type `T`. Only used when `T` is `FixedLenByteArrayType`
- pub(crate) type_length: i32,
-
- // The byte array to decode from. Not set if `T` is bool.
- pub(crate) data: Option<ByteBufferPtr>,
-
- // Read `data` bit by bit. Only set if `T` is bool.
- pub(crate) bit_reader: Option<BitReader>,
-}
-
-/// Plain decoding that supports all types.
-/// Values are encoded back to back. For native types, data is encoded as little endian.
-/// Floating point types are encoded in IEEE.
-/// See [`PlainEncoder`](crate::encoding::PlainEncoder) for more information.
-pub struct PlainDecoder<T: DataType> {
- // The binary details needed for decoding
- inner: PlainDecoderDetails,
-
- // To allow `T` in the generic parameter for this struct. This doesn't take any
- // space.
- _phantom: PhantomData<T>,
-}
-
-impl<T: DataType> PlainDecoder<T> {
- /// Creates new plain decoder.
- pub fn new(type_length: i32) -> Self {
- PlainDecoder {
- inner: PlainDecoderDetails {
- type_length,
- num_values: 0,
- start: 0,
- data: None,
- bit_reader: None,
- },
- _phantom: PhantomData,
- }
- }
-}
-
-impl<T: DataType> Decoder<T> for PlainDecoder<T> {
- #[inline]
- fn set_data(&mut self, data: ByteBufferPtr, num_values: usize) -> Result<()> {
- T::T::set_data(&mut self.inner, data, num_values);
- Ok(())
- }
-
- #[inline]
- fn values_left(&self) -> usize {
- self.inner.num_values
- }
-
- #[inline]
- fn encoding(&self) -> Encoding {
- Encoding::PLAIN
- }
-
- #[inline]
- fn get(&mut self, buffer: &mut [T::T]) -> Result<usize> {
- T::T::decode(buffer, &mut self.inner)
- }
-}
-
-// ----------------------------------------------------------------------
-// RLE_DICTIONARY/PLAIN_DICTIONARY Decoding
-
-/// Dictionary decoder.
-/// The dictionary encoding builds a dictionary of values encountered in a given column.
-/// The dictionary is be stored in a dictionary page per column chunk.
-/// See [`DictEncoder`](crate::encoding::DictEncoder) for more information.
-pub struct DictDecoder<T: DataType> {
- // The dictionary, which maps ids to the values
- dictionary: Vec<T::T>,
-
- // Whether `dictionary` has been initialized
- has_dictionary: bool,
-
- // The decoder for the value ids
- rle_decoder: Option<RleDecoder>,
-
- // Number of values left in the data stream
- num_values: usize,
-}
-
-impl<T: DataType> DictDecoder<T> {
- /// Creates new dictionary decoder.
- pub fn new() -> Self {
- Self {
- dictionary: vec![],
- has_dictionary: false,
- rle_decoder: None,
- num_values: 0,
- }
- }
-
- /// Decodes and sets values for dictionary using `decoder` decoder.
- pub fn set_dict(&mut self, mut decoder: Box<dyn Decoder<T>>) -> Result<()> {
- let num_values = decoder.values_left();
- self.dictionary.resize(num_values, T::T::default());
- let _ = decoder.get(&mut self.dictionary)?;
- self.has_dictionary = true;
- Ok(())
- }
-}
-
-impl<T: DataType> Decoder<T> for DictDecoder<T> {
- fn set_data(&mut self, data: ByteBufferPtr, num_values: usize) -> Result<()> {
- // First byte in `data` is bit width
- let bit_width = data.as_ref()[0];
- let mut rle_decoder = RleDecoder::new(bit_width);
- rle_decoder.set_data(data.start_from(1));
- self.num_values = num_values;
- self.rle_decoder = Some(rle_decoder);
- Ok(())
- }
-
- fn get(&mut self, buffer: &mut [T::T]) -> Result<usize> {
- assert!(self.rle_decoder.is_some());
- assert!(self.has_dictionary, "Must call set_dict() first!");
-
- let rle = self.rle_decoder.as_mut().unwrap();
- let num_values = cmp::min(buffer.len(), self.num_values);
- rle.get_batch_with_dict(&self.dictionary[..], buffer, num_values)
- }
-
- /// Number of values left in this decoder stream
- fn values_left(&self) -> usize {
- self.num_values
- }
-
- fn encoding(&self) -> Encoding {
- Encoding::RLE_DICTIONARY
- }
-}
-
-// ----------------------------------------------------------------------
-// RLE Decoding
-
-/// RLE/Bit-Packing hybrid decoding for values.
-/// Currently is used only for data pages v2 and supports boolean types.
-/// See [`RleValueEncoder`](crate::encoding::RleValueEncoder) for more information.
-pub struct RleValueDecoder<T: DataType> {
- values_left: usize,
- decoder: RleDecoder,
- _phantom: PhantomData<T>,
-}
-
-impl<T: DataType> RleValueDecoder<T> {
- pub fn new() -> Self {
- Self {
- values_left: 0,
- decoder: RleDecoder::new(1),
- _phantom: PhantomData,
- }
- }
-}
-
-impl<T: DataType> Decoder<T> for RleValueDecoder<T> {
- #[inline]
- fn set_data(&mut self, data: ByteBufferPtr, num_values: usize) -> Result<()> {
- // Only support RLE value reader for boolean values with bit width of 1.
- ensure_phys_ty!(Type::BOOLEAN, "RleValueDecoder only supports BoolType");
-
- // We still need to remove prefix of i32 from the stream.
- const I32_SIZE: usize = mem::size_of::<i32>();
- let data_size = read_num_bytes!(i32, I32_SIZE, data.as_ref()) as usize;
- self.decoder = RleDecoder::new(1);
- self.decoder.set_data(data.range(I32_SIZE, data_size));
- self.values_left = num_values;
- Ok(())
- }
-
- #[inline]
- fn values_left(&self) -> usize {
- self.values_left
- }
-
- #[inline]
- fn encoding(&self) -> Encoding {
- Encoding::RLE
- }
-
- #[inline]
- fn get(&mut self, buffer: &mut [T::T]) -> Result<usize> {
- let num_values = cmp::min(buffer.len(), self.values_left);
- let values_read = self.decoder.get_batch(&mut buffer[..num_values])?;
- self.values_left -= values_read;
- Ok(values_read)
- }
-}
-
-// ----------------------------------------------------------------------
-// DELTA_BINARY_PACKED Decoding
-
-/// Delta binary packed decoder.
-/// Supports INT32 and INT64 types.
-/// See [`DeltaBitPackEncoder`](crate::encoding::DeltaBitPackEncoder) for more
-/// information.
-pub struct DeltaBitPackDecoder<T: DataType> {
- bit_reader: BitReader,
- initialized: bool,
-
- // Header info
- num_values: usize,
- num_mini_blocks: i64,
- values_per_mini_block: usize,
- values_current_mini_block: usize,
- first_value: i64,
- first_value_read: bool,
-
- // Per block info
- min_delta: i64,
- mini_block_idx: usize,
- delta_bit_width: u8,
- delta_bit_widths: ByteBuffer,
- deltas_in_mini_block: Vec<T::T>, // eagerly loaded deltas for a mini block
- use_batch: bool,
-
- current_value: i64,
-
- _phantom: PhantomData<T>,
-}
-
-impl<T: DataType> DeltaBitPackDecoder<T> {
- /// Creates new delta bit packed decoder.
- pub fn new() -> Self {
- Self {
- bit_reader: BitReader::from(vec![]),
- initialized: false,
- num_values: 0,
- num_mini_blocks: 0,
- values_per_mini_block: 0,
- values_current_mini_block: 0,
- first_value: 0,
- first_value_read: false,
- min_delta: 0,
- mini_block_idx: 0,
- delta_bit_width: 0,
- delta_bit_widths: ByteBuffer::new(),
- deltas_in_mini_block: vec![],
- use_batch: mem::size_of::<T::T>() == 4,
- current_value: 0,
- _phantom: PhantomData,
- }
- }
-
- /// Returns underlying bit reader offset.
- pub fn get_offset(&self) -> usize {
- assert!(self.initialized, "Bit reader is not initialized");
- self.bit_reader.get_byte_offset()
- }
-
- /// Initializes new mini block.
- #[inline]
- fn init_block(&mut self) -> Result<()> {
- self.min_delta = self
- .bit_reader
- .get_zigzag_vlq_int()
- .ok_or_else(|| eof_err!("Not enough data to decode 'min_delta'"))?;
-
- self.delta_bit_widths.clear();
- for _ in 0..self.num_mini_blocks {
- let w = self
- .bit_reader
- .get_aligned::<u8>(1)
- .ok_or_else(|| eof_err!("Not enough data to decode 'width'"))?;
- self.delta_bit_widths.push(w);
- }
-
- self.mini_block_idx = 0;
- self.delta_bit_width = self.delta_bit_widths.data()[0];
- self.values_current_mini_block = self.values_per_mini_block;
- Ok(())
- }
-
- /// Loads delta into mini block.
- #[inline]
- fn load_deltas_in_mini_block(&mut self) -> Result<()>
- where
- T::T: FromBytes,
- {
- if self.use_batch {
- self.deltas_in_mini_block
- .resize(self.values_current_mini_block, T::T::default());
- let loaded = self.bit_reader.get_batch::<T::T>(
- &mut self.deltas_in_mini_block[..],
- self.delta_bit_width as usize,
- );
- assert!(loaded == self.values_current_mini_block);
- } else {
- self.deltas_in_mini_block.clear();
- for _ in 0..self.values_current_mini_block {
- // TODO: load one batch at a time similar to int32
- let delta = self
- .bit_reader
- .get_value::<T::T>(self.delta_bit_width as usize)
- .ok_or_else(|| eof_err!("Not enough data to decode 'delta'"))?;
- self.deltas_in_mini_block.push(delta);
- }
- }
-
- Ok(())
- }
-}
-
-impl<T: DataType> Decoder<T> for DeltaBitPackDecoder<T> {
- // # of total values is derived from encoding
- #[inline]
- fn set_data(&mut self, data: ByteBufferPtr, _index: usize) -> Result<()> {
- self.bit_reader = BitReader::new(data);
- self.initialized = true;
-
- let block_size = self
- .bit_reader
- .get_vlq_int()
- .ok_or_else(|| eof_err!("Not enough data to decode 'block_size'"))?;
- self.num_mini_blocks = self
- .bit_reader
- .get_vlq_int()
- .ok_or_else(|| eof_err!("Not enough data to decode 'num_mini_blocks'"))?;
- self.num_values = self
- .bit_reader
- .get_vlq_int()
- .ok_or_else(|| eof_err!("Not enough data to decode 'num_values'"))?
- as usize;
- self.first_value = self
- .bit_reader
- .get_zigzag_vlq_int()
- .ok_or_else(|| eof_err!("Not enough data to decode 'first_value'"))?;
-
- // Reset decoding state
- self.first_value_read = false;
- self.mini_block_idx = 0;
- self.delta_bit_widths.clear();
- self.values_current_mini_block = 0;
-
- self.values_per_mini_block = (block_size / self.num_mini_blocks) as usize;
- assert!(self.values_per_mini_block % 8 == 0);
-
- Ok(())
- }
-
- fn get(&mut self, buffer: &mut [T::T]) -> Result<usize> {
- assert!(self.initialized, "Bit reader is not initialized");
-
- let num_values = cmp::min(buffer.len(), self.num_values);
- for i in 0..num_values {
- if !self.first_value_read {
- self.set_decoded_value(buffer, i, self.first_value);
- self.current_value = self.first_value;
- self.first_value_read = true;
- continue;
- }
-
- if self.values_current_mini_block == 0 {
- self.mini_block_idx += 1;
- if self.mini_block_idx < self.delta_bit_widths.size() {
- self.delta_bit_width =
- self.delta_bit_widths.data()[self.mini_block_idx];
- self.values_current_mini_block = self.values_per_mini_block;
- } else {
- self.init_block()?;
- }
- self.load_deltas_in_mini_block()?;
- }
-
- // we decrement values in current mini block, so we need to invert index for
- // delta
- let delta = self.get_delta(
- self.deltas_in_mini_block.len() - self.values_current_mini_block,
- );
- // It is OK for deltas to contain "overflowed" values after encoding,
- // e.g. i64::MAX - i64::MIN, so we use `wrapping_add` to "overflow" again and
- // restore original value.
- self.current_value = self.current_value.wrapping_add(self.min_delta);
- self.current_value = self.current_value.wrapping_add(delta as i64);
- self.set_decoded_value(buffer, i, self.current_value);
- self.values_current_mini_block -= 1;
- }
-
- self.num_values -= num_values;
- Ok(num_values)
- }
-
- fn values_left(&self) -> usize {
- self.num_values
- }
-
- fn encoding(&self) -> Encoding {
- Encoding::DELTA_BINARY_PACKED
- }
-}
-
-/// Helper trait to define specific conversions when decoding values
-trait DeltaBitPackDecoderConversion<T: DataType> {
- /// Sets decoded value based on type `T`.
- fn get_delta(&self, index: usize) -> i64;
-
- fn set_decoded_value(&self, buffer: &mut [T::T], index: usize, value: i64);
-}
-
-impl<T: DataType> DeltaBitPackDecoderConversion<T> for DeltaBitPackDecoder<T> {
- #[inline]
- fn get_delta(&self, index: usize) -> i64 {
- ensure_phys_ty!(
- Type::INT32 | Type::INT64,
- "DeltaBitPackDecoder only supports Int32Type and Int64Type"
- );
- self.deltas_in_mini_block[index].as_i64().unwrap()
- }
-
- #[inline]
- fn set_decoded_value(&self, buffer: &mut [T::T], index: usize, value: i64) {
- match T::get_physical_type() {
- Type::INT32 => {
- let val = buffer[index].as_mut_any().downcast_mut::<i32>().unwrap();
-
- *val = value as i32;
- }
- Type::INT64 => {
- let val = buffer[index].as_mut_any().downcast_mut::<i64>().unwrap();
-
- *val = value;
- }
- _ => panic!("DeltaBitPackDecoder only supports Int32Type and Int64Type"),
- };
- }
-}
-
-// ----------------------------------------------------------------------
-// DELTA_LENGTH_BYTE_ARRAY Decoding
-
-/// Delta length byte array decoder.
-/// Only applied to byte arrays to separate the length values and the data, the lengths
-/// are encoded using DELTA_BINARY_PACKED encoding.
-/// See [`DeltaLengthByteArrayEncoder`](crate::encoding::DeltaLengthByteArrayEncoder)
-/// for more information.
-pub struct DeltaLengthByteArrayDecoder<T: DataType> {
- // Lengths for each byte array in `data`
- // TODO: add memory tracker to this
- lengths: Vec<i32>,
-
- // Current index into `lengths`
- current_idx: usize,
-
- // Concatenated byte array data
- data: Option<ByteBufferPtr>,
-
- // Offset into `data`, always point to the beginning of next byte array.
- offset: usize,
-
- // Number of values left in this decoder stream
- num_values: usize,
-
- // Placeholder to allow `T` as generic parameter
- _phantom: PhantomData<T>,
-}
-
-impl<T: DataType> DeltaLengthByteArrayDecoder<T> {
- /// Creates new delta length byte array decoder.
- pub fn new() -> Self {
- Self {
- lengths: vec![],
- current_idx: 0,
- data: None,
- offset: 0,
- num_values: 0,
- _phantom: PhantomData,
- }
- }
-}
-
-impl<T: DataType> Decoder<T> for DeltaLengthByteArrayDecoder<T> {
- fn set_data(&mut self, data: ByteBufferPtr, num_values: usize) -> Result<()> {
- match T::get_physical_type() {
- Type::BYTE_ARRAY => {
- let mut len_decoder = DeltaBitPackDecoder::<Int32Type>::new();
- len_decoder.set_data(data.all(), num_values)?;
- let num_lengths = len_decoder.values_left();
- self.lengths.resize(num_lengths, 0);
- len_decoder.get(&mut self.lengths[..])?;
-
- self.data = Some(data.start_from(len_decoder.get_offset()));
- self.offset = 0;
- self.current_idx = 0;
- self.num_values = num_lengths;
- Ok(())
- }
- _ => Err(general_err!(
- "DeltaLengthByteArrayDecoder only support ByteArrayType"
- )),
- }
- }
-
- fn get(&mut self, buffer: &mut [T::T]) -> Result<usize> {
- match T::get_physical_type() {
- Type::BYTE_ARRAY => {
- assert!(self.data.is_some());
-
- let data = self.data.as_ref().unwrap();
- let num_values = cmp::min(buffer.len(), self.num_values);
- for i in 0..num_values {
- let len = self.lengths[self.current_idx] as usize;
-
- buffer[i]
- .as_mut_any()
- .downcast_mut::<ByteArray>()
- .unwrap()
- .set_data(data.range(self.offset, len));
-
- self.offset += len;
- self.current_idx += 1;
- }
-
- self.num_values -= num_values;
- Ok(num_values)
- }
- _ => Err(general_err!(
- "DeltaLengthByteArrayDecoder only support ByteArrayType"
- )),
- }
- }
-
- fn values_left(&self) -> usize {
- self.num_values
- }
-
- fn encoding(&self) -> Encoding {
- Encoding::DELTA_LENGTH_BYTE_ARRAY
- }
-}
-
-// ----------------------------------------------------------------------
-// DELTA_BYTE_ARRAY Decoding
-
-/// Delta byte array decoder.
-/// Prefix lengths are encoded using `DELTA_BINARY_PACKED` encoding, Suffixes are stored
-/// using `DELTA_LENGTH_BYTE_ARRAY` encoding.
-/// See [`DeltaByteArrayEncoder`](crate::encoding::DeltaByteArrayEncoder) for more
-/// information.
-pub struct DeltaByteArrayDecoder<T: DataType> {
- // Prefix lengths for each byte array
- // TODO: add memory tracker to this
- prefix_lengths: Vec<i32>,
-
- // The current index into `prefix_lengths`,
- current_idx: usize,
-
- // Decoder for all suffixes, the # of which should be the same as
- // `prefix_lengths.len()`
- suffix_decoder: Option<DeltaLengthByteArrayDecoder<ByteArrayType>>,
-
- // The last byte array, used to derive the current prefix
- previous_value: Vec<u8>,
-
- // Number of values left
- num_values: usize,
-
- // Placeholder to allow `T` as generic parameter
- _phantom: PhantomData<T>,
-}
-
-impl<T: DataType> DeltaByteArrayDecoder<T> {
- /// Creates new delta byte array decoder.
- pub fn new() -> Self {
- Self {
- prefix_lengths: vec![],
- current_idx: 0,
- suffix_decoder: None,
- previous_value: vec![],
- num_values: 0,
- _phantom: PhantomData,
- }
- }
-}
-
-impl<'m, T: DataType> Decoder<T> for DeltaByteArrayDecoder<T> {
- fn set_data(&mut self, data: ByteBufferPtr, num_values: usize) -> Result<()> {
- match T::get_physical_type() {
- Type::BYTE_ARRAY | Type::FIXED_LEN_BYTE_ARRAY => {
- let mut prefix_len_decoder = DeltaBitPackDecoder::<Int32Type>::new();
- prefix_len_decoder.set_data(data.all(), num_values)?;
- let num_prefixes = prefix_len_decoder.values_left();
- self.prefix_lengths.resize(num_prefixes, 0);
- prefix_len_decoder.get(&mut self.prefix_lengths[..])?;
-
- let mut suffix_decoder = DeltaLengthByteArrayDecoder::new();
- suffix_decoder
- .set_data(data.start_from(prefix_len_decoder.get_offset()), num_values)?;
- self.suffix_decoder = Some(suffix_decoder);
- self.num_values = num_prefixes;
- self.current_idx = 0;
- self.previous_value.clear();
- Ok(())
- }
- _ => {
- Err(general_err!(
- "DeltaByteArrayDecoder only supports ByteArrayType and FixedLenByteArrayType"
- ))
- }
- }
- }
-
- fn get(&mut self, buffer: &mut [T::T]) -> Result<usize> {
- match T::get_physical_type() {
- ty @ Type::BYTE_ARRAY | ty @ Type::FIXED_LEN_BYTE_ARRAY => {
- let num_values = cmp::min(buffer.len(), self.num_values);
- let mut v: [ByteArray; 1] = [ByteArray::new(); 1];
- for i in 0..num_values {
- // Process suffix
- // TODO: this is awkward - maybe we should add a non-vectorized API?
- let suffix_decoder = self.suffix_decoder.as_mut().expect("decoder not initialized");
- suffix_decoder.get(&mut v[..])?;
- let suffix = v[0].data();
-
- // Extract current prefix length, can be 0
- let prefix_len = self.prefix_lengths[self.current_idx] as usize;
-
- // Concatenate prefix with suffix
- let mut result = Vec::new();
- result.extend_from_slice(&self.previous_value[0..prefix_len]);
- result.extend_from_slice(suffix);
-
- let data = ByteBufferPtr::new(result.clone());
-
- match ty {
- Type::BYTE_ARRAY => buffer[i]
- .as_mut_any()
- .downcast_mut::<ByteArray>()
- .unwrap()
- .set_data(data),
- Type::FIXED_LEN_BYTE_ARRAY => buffer[i]
- .as_mut_any()
- .downcast_mut::<FixedLenByteArray>()
- .unwrap()
- .set_data(data),
- _ => unreachable!(),
- };
-
- self.previous_value = result;
- self.current_idx += 1;
- }
-
- self.num_values -= num_values;
- Ok(num_values)
- }
- _ => {
- Err(general_err!(
- "DeltaByteArrayDecoder only supports ByteArrayType and FixedLenByteArrayType"
- ))
- }
- }
- }
-
- fn values_left(&self) -> usize {
- self.num_values
- }
-
- fn encoding(&self) -> Encoding {
- Encoding::DELTA_BYTE_ARRAY
- }
-}
-
-#[cfg(test)]
-#[allow(clippy::approx_constant)]
-mod tests {
- use super::{super::encoding::*, *};
-
- use std::sync::Arc;
-
- use crate::schema::types::{
- ColumnDescPtr, ColumnDescriptor, ColumnPath, Type as SchemaType,
- };
- use crate::util::{
- bit_util::set_array_bit, memory::MemTracker, test_common::RandGen,
- };
-
- #[test]
- fn test_get_decoders() {
- // supported encodings
- create_and_check_decoder::<Int32Type>(Encoding::PLAIN, None);
- create_and_check_decoder::<Int32Type>(Encoding::DELTA_BINARY_PACKED, None);
- create_and_check_decoder::<Int32Type>(Encoding::DELTA_LENGTH_BYTE_ARRAY, None);
- create_and_check_decoder::<Int32Type>(Encoding::DELTA_BYTE_ARRAY, None);
- create_and_check_decoder::<BoolType>(Encoding::RLE, None);
-
- // error when initializing
- create_and_check_decoder::<Int32Type>(
- Encoding::RLE_DICTIONARY,
- Some(general_err!(
- "Cannot initialize this encoding through this function"
- )),
- );
- create_and_check_decoder::<Int32Type>(
- Encoding::PLAIN_DICTIONARY,
- Some(general_err!(
- "Cannot initialize this encoding through this function"
- )),
- );
-
- // unsupported
- create_and_check_decoder::<Int32Type>(
- Encoding::BIT_PACKED,
- Some(nyi_err!("Encoding BIT_PACKED is not supported")),
- );
- }
-
- #[test]
- fn test_plain_decode_int32() {
- let data = vec![42, 18, 52];
- let data_bytes = Int32Type::to_byte_array(&data[..]);
- let mut buffer = vec![0; 3];
- test_plain_decode::<Int32Type>(
- ByteBufferPtr::new(data_bytes),
- 3,
- -1,
- &mut buffer[..],
- &data[..],
- );
- }
-
- #[test]
- fn test_plain_decode_int32_spaced() {
- let data = [42, 18, 52];
- let expected_data = [0, 42, 0, 18, 0, 0, 52, 0];
- let data_bytes = Int32Type::to_byte_array(&data[..]);
- let mut buffer = vec![0; 8];
- let num_nulls = 5;
- let valid_bits = [0b01001010];
- test_plain_decode_spaced::<Int32Type>(
- ByteBufferPtr::new(data_bytes),
- 3,
- -1,
- &mut buffer[..],
- num_nulls,
- &valid_bits,
- &expected_data[..],
- );
- }
-
- #[test]
- fn test_plain_decode_int64() {
- let data = vec![42, 18, 52];
- let data_bytes = Int64Type::to_byte_array(&data[..]);
- let mut buffer = vec![0; 3];
- test_plain_decode::<Int64Type>(
- ByteBufferPtr::new(data_bytes),
- 3,
- -1,
- &mut buffer[..],
- &data[..],
- );
- }
-
- #[test]
- fn test_plain_decode_float() {
- let data = vec![3.14, 2.414, 12.51];
- let data_bytes = FloatType::to_byte_array(&data[..]);
- let mut buffer = vec![0.0; 3];
- test_plain_decode::<FloatType>(
- ByteBufferPtr::new(data_bytes),
- 3,
- -1,
- &mut buffer[..],
- &data[..],
- );
- }
-
- #[test]
- fn test_plain_decode_double() {
- let data = vec![3.14f64, 2.414f64, 12.51f64];
- let data_bytes = DoubleType::to_byte_array(&data[..]);
- let mut buffer = vec![0.0f64; 3];
- test_plain_decode::<DoubleType>(
- ByteBufferPtr::new(data_bytes),
- 3,
- -1,
- &mut buffer[..],
- &data[..],
- );
- }
-
- #[test]
- fn test_plain_decode_int96() {
- let mut data = vec![Int96::new(); 4];
- data[0].set_data(11, 22, 33);
- data[1].set_data(44, 55, 66);
- data[2].set_data(10, 20, 30);
- data[3].set_data(40, 50, 60);
- let data_bytes = Int96Type::to_byte_array(&data[..]);
- let mut buffer = vec![Int96::new(); 4];
- test_plain_decode::<Int96Type>(
- ByteBufferPtr::new(data_bytes),
- 4,
- -1,
- &mut buffer[..],
- &data[..],
- );
- }
-
- #[test]
- fn test_plain_decode_bool() {
- let data = vec![
- false, true, false, false, true, false, true, true, false, true,
- ];
- let data_bytes = BoolType::to_byte_array(&data[..]);
- let mut buffer = vec![false; 10];
- test_plain_decode::<BoolType>(
- ByteBufferPtr::new(data_bytes),
- 10,
- -1,
- &mut buffer[..],
- &data[..],
- );
- }
-
- #[test]
- fn test_plain_decode_byte_array() {
- let mut data = vec![ByteArray::new(); 2];
- data[0].set_data(ByteBufferPtr::new(String::from("hello").into_bytes()));
- data[1].set_data(ByteBufferPtr::new(String::from("parquet").into_bytes()));
- let data_bytes = ByteArrayType::to_byte_array(&data[..]);
- let mut buffer = vec![ByteArray::new(); 2];
- test_plain_decode::<ByteArrayType>(
- ByteBufferPtr::new(data_bytes),
- 2,
- -1,
- &mut buffer[..],
- &data[..],
- );
- }
-
- #[test]
- fn test_plain_decode_fixed_len_byte_array() {
- let mut data = vec![FixedLenByteArray::default(); 3];
- data[0].set_data(ByteBufferPtr::new(String::from("bird").into_bytes()));
- data[1].set_data(ByteBufferPtr::new(String::from("come").into_bytes()));
- data[2].set_data(ByteBufferPtr::new(String::from("flow").into_bytes()));
- let data_bytes = FixedLenByteArrayType::to_byte_array(&data[..]);
- let mut buffer = vec![FixedLenByteArray::default(); 3];
- test_plain_decode::<FixedLenByteArrayType>(
- ByteBufferPtr::new(data_bytes),
- 3,
- 4,
- &mut buffer[..],
- &data[..],
- );
- }
-
- fn test_plain_decode<T: DataType>(
- data: ByteBufferPtr,
- num_values: usize,
- type_length: i32,
- buffer: &mut [T::T],
- expected: &[T::T],
- ) {
- let mut decoder: PlainDecoder<T> = PlainDecoder::new(type_length);
- let result = decoder.set_data(data, num_values);
- assert!(result.is_ok());
- let result = decoder.get(buffer);
- assert!(result.is_ok());
- assert_eq!(decoder.values_left(), 0);
- assert_eq!(buffer, expected);
- }
-
- fn test_plain_decode_spaced<T: DataType>(
- data: ByteBufferPtr,
- num_values: usize,
- type_length: i32,
- buffer: &mut [T::T],
- num_nulls: usize,
- valid_bits: &[u8],
- expected: &[T::T],
- ) {
- let mut decoder: PlainDecoder<T> = PlainDecoder::new(type_length);
- let result = decoder.set_data(data, num_values);
- assert!(result.is_ok());
- let result = decoder.get_spaced(buffer, num_nulls, valid_bits);
- assert!(result.is_ok());
- assert_eq!(num_values + num_nulls, result.unwrap());
- assert_eq!(decoder.values_left(), 0);
- assert_eq!(buffer, expected);
- }
-
- #[test]
- #[should_panic(expected = "RleValueEncoder only supports BoolType")]
- fn test_rle_value_encode_int32_not_supported() {
- let mut encoder = RleValueEncoder::<Int32Type>::new();
- encoder.put(&[1, 2, 3, 4]).unwrap();
- }
-
- #[test]
- #[should_panic(expected = "RleValueDecoder only supports BoolType")]
- fn test_rle_value_decode_int32_not_supported() {
- let mut decoder = RleValueDecoder::<Int32Type>::new();
- decoder
- .set_data(ByteBufferPtr::new(vec![5, 0, 0, 0]), 1)
- .unwrap();
- }
-
- #[test]
- fn test_rle_value_decode_bool_decode() {
- // Test multiple 'put' calls on the same encoder
- let data = vec![
- BoolType::gen_vec(-1, 256),
- BoolType::gen_vec(-1, 257),
- BoolType::gen_vec(-1, 126),
- ];
- test_rle_value_decode::<BoolType>(data);
- }
-
- #[test]
- #[should_panic(expected = "Bit reader is not initialized")]
- fn test_delta_bit_packed_not_initialized_offset() {
- // Fail if set_data() is not called before get_offset()
- let decoder = DeltaBitPackDecoder::<Int32Type>::new();
- decoder.get_offset();
- }
-
- #[test]
- #[should_panic(expected = "Bit reader is not initialized")]
- fn test_delta_bit_packed_not_initialized_get() {
- // Fail if set_data() is not called before get()
- let mut decoder = DeltaBitPackDecoder::<Int32Type>::new();
- let mut buffer = vec![];
- decoder.get(&mut buffer).unwrap();
- }
-
- #[test]
- fn test_delta_bit_packed_int32_empty() {
- let data = vec![vec![0; 0]];
- test_delta_bit_packed_decode::<Int32Type>(data);
- }
-
- #[test]
- fn test_delta_bit_packed_int32_repeat() {
- let block_data = vec![
- 1, 2, 3, 4, 5, 6, 7, 8, 1, 2, 3, 4, 5, 6, 7, 8, 1, 2, 3, 4, 5, 6, 7, 8, 1, 2,
- 3, 4, 5, 6, 7, 8,
- ];
- test_delta_bit_packed_decode::<Int32Type>(vec![block_data]);
- }
-
- #[test]
- fn test_delta_bit_packed_int32_uneven() {
- let block_data = vec![1, -2, 3, -4, 5, 6, 7, 8, 9, 10, 11];
- test_delta_bit_packed_decode::<Int32Type>(vec![block_data]);
- }
-
- #[test]
- fn test_delta_bit_packed_int32_same_values() {
- let block_data = vec![
- 127, 127, 127, 127, 127, 127, 127, 127, 127, 127, 127, 127, 127, 127, 127,
- 127,
- ];
- test_delta_bit_packed_decode::<Int32Type>(vec![block_data]);
-
- let block_data = vec![
- -127, -127, -127, -127, -127, -127, -127, -127, -127, -127, -127, -127, -127,
- -127, -127, -127,
- ];
- test_delta_bit_packed_decode::<Int32Type>(vec![block_data]);
- }
-
- #[test]
- fn test_delta_bit_packed_int32_min_max() {
- let block_data = vec![
- i32::min_value(),
- i32::max_value(),
- i32::min_value(),
- i32::max_value(),
- i32::min_value(),
- i32::max_value(),
- i32::min_value(),
- i32::max_value(),
- ];
- test_delta_bit_packed_decode::<Int32Type>(vec![block_data]);
- }
-
- #[test]
- fn test_delta_bit_packed_int32_multiple_blocks() {
- // Test multiple 'put' calls on the same encoder
- let data = vec![
- Int32Type::gen_vec(-1, 64),
- Int32Type::gen_vec(-1, 128),
- Int32Type::gen_vec(-1, 64),
- ];
- test_delta_bit_packed_decode::<Int32Type>(data);
- }
-
- #[test]
- fn test_delta_bit_packed_int32_data_across_blocks() {
- // Test multiple 'put' calls on the same encoder
- let data = vec![Int32Type::gen_vec(-1, 256), Int32Type::gen_vec(-1, 257)];
- test_delta_bit_packed_decode::<Int32Type>(data);
- }
-
- #[test]
- fn test_delta_bit_packed_int32_with_empty_blocks() {
- let data = vec![
- Int32Type::gen_vec(-1, 128),
- vec![0; 0],
- Int32Type::gen_vec(-1, 64),
- ];
- test_delta_bit_packed_decode::<Int32Type>(data);
- }
-
- #[test]
- fn test_delta_bit_packed_int64_empty() {
- let data = vec![vec![0; 0]];
- test_delta_bit_packed_decode::<Int64Type>(data);
- }
-
- #[test]
- fn test_delta_bit_packed_int64_min_max() {
- let block_data = vec![
- i64::min_value(),
- i64::max_value(),
- i64::min_value(),
- i64::max_value(),
- i64::min_value(),
- i64::max_value(),
- i64::min_value(),
- i64::max_value(),
- ];
- test_delta_bit_packed_decode::<Int64Type>(vec![block_data]);
- }
-
- #[test]
- fn test_delta_bit_packed_int64_multiple_blocks() {
- // Test multiple 'put' calls on the same encoder
- let data = vec![
- Int64Type::gen_vec(-1, 64),
- Int64Type::gen_vec(-1, 128),
- Int64Type::gen_vec(-1, 64),
- ];
- test_delta_bit_packed_decode::<Int64Type>(data);
- }
-
- #[test]
- fn test_delta_bit_packed_decoder_sample() {
- let data_bytes = vec![
- 128, 1, 4, 3, 58, 28, 6, 0, 0, 0, 0, 8, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
- 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
- ];
- let buffer = ByteBufferPtr::new(data_bytes);
- let mut decoder: DeltaBitPackDecoder<Int32Type> = DeltaBitPackDecoder::new();
- decoder.set_data(buffer, 3).unwrap();
- // check exact offsets, because when reading partial values we end up with
- // some data not being read from bit reader
- assert_eq!(decoder.get_offset(), 5);
- let mut result = vec![0, 0, 0];
- decoder.get(&mut result).unwrap();
- assert_eq!(decoder.get_offset(), 34);
- assert_eq!(result, vec![29, 43, 89]);
- }
-
- #[test]
- fn test_delta_byte_array_same_arrays() {
- let data = vec![
- vec![ByteArray::from(vec![1, 2, 3, 4, 5, 6])],
- vec![
- ByteArray::from(vec![1, 2, 3, 4, 5, 6]),
- ByteArray::from(vec![1, 2, 3, 4, 5, 6]),
- ],
- vec![
- ByteArray::from(vec![1, 2, 3, 4, 5, 6]),
- ByteArray::from(vec![1, 2, 3, 4, 5, 6]),
- ],
- ];
- test_delta_byte_array_decode(data);
- }
-
- #[test]
- fn test_delta_byte_array_unique_arrays() {
- let data = vec![
- vec![ByteArray::from(vec![1])],
- vec![ByteArray::from(vec![2, 3]), ByteArray::from(vec![4, 5, 6])],
- vec![
- ByteArray::from(vec![7, 8]),
- ByteArray::from(vec![9, 0, 1, 2]),
- ],
- ];
- test_delta_byte_array_decode(data);
- }
-
- #[test]
- fn test_delta_byte_array_single_array() {
- let data = vec![vec![ByteArray::from(vec![1, 2, 3, 4, 5, 6])]];
- test_delta_byte_array_decode(data);
- }
-
- fn test_rle_value_decode<T: DataType>(data: Vec<Vec<T::T>>) {
- test_encode_decode::<T>(data, Encoding::RLE);
- }
-
- fn test_delta_bit_packed_decode<T: DataType>(data: Vec<Vec<T::T>>) {
- test_encode_decode::<T>(data, Encoding::DELTA_BINARY_PACKED);
- }
-
- fn test_delta_byte_array_decode(data: Vec<Vec<ByteArray>>) {
- test_encode_decode::<ByteArrayType>(data, Encoding::DELTA_BYTE_ARRAY);
- }
-
- // Input data represents vector of data slices to write (test multiple `put()` calls)
- // For example,
- // vec![vec![1, 2, 3]] invokes `put()` once and writes {1, 2, 3}
- // vec![vec![1, 2], vec![3]] invokes `put()` twice and writes {1, 2, 3}
- fn test_encode_decode<T: DataType>(data: Vec<Vec<T::T>>, encoding: Encoding) {
- // Type length should not really matter for encode/decode test,
- // otherwise change it based on type
- let col_descr = create_test_col_desc_ptr(-1, T::get_physical_type());
-
- // Encode data
- let mut encoder =
- get_encoder::<T>(col_descr.clone(), encoding, Arc::new(MemTracker::new()))
- .expect("get encoder");
-
- for v in &data[..] {
- encoder.put(&v[..]).expect("ok to encode");
- }
- let bytes = encoder.flush_buffer().expect("ok to flush buffer");
-
- // Flatten expected data as contiguous array of values
- let expected: Vec<T::T> = data.iter().flat_map(|s| s.clone()).collect();
-
- // Decode data and compare with original
- let mut decoder = get_decoder::<T>(col_descr, encoding).expect("get decoder");
-
- let mut result = vec![T::T::default(); expected.len()];
- decoder
- .set_data(bytes, expected.len())
- .expect("ok to set data");
- let mut result_num_values = 0;
- while decoder.values_left() > 0 {
- result_num_values += decoder
- .get(&mut result[result_num_values..])
- .expect("ok to decode");
- }
- assert_eq!(result_num_values, expected.len());
- assert_eq!(result, expected);
- }
-
- fn create_and_check_decoder<T: DataType>(
- encoding: Encoding,
- err: Option<ParquetError>,
- ) {
- let descr = create_test_col_desc_ptr(-1, T::get_physical_type());
- let decoder = get_decoder::<T>(descr, encoding);
- match err {
- Some(parquet_error) => {
- assert!(decoder.is_err());
- assert_eq!(decoder.err().unwrap(), parquet_error);
- }
- None => {
- assert!(decoder.is_ok());
- assert_eq!(decoder.unwrap().encoding(), encoding);
- }
- }
- }
-
- // Creates test column descriptor.
- fn create_test_col_desc_ptr(type_len: i32, t: Type) -> ColumnDescPtr {
- let ty = SchemaType::primitive_type_builder("t", t)
- .with_length(type_len)
- .build()
- .unwrap();
- Arc::new(ColumnDescriptor::new(
- Arc::new(ty),
- 0,
- 0,
- ColumnPath::new(vec![]),
- ))
- }
-
- fn usize_to_bytes(v: usize) -> [u8; 4] {
- (v as u32).to_ne_bytes()
- }
-
- /// A util trait to convert slices of different types to byte arrays
- trait ToByteArray<T: DataType> {
- #[allow(clippy::wrong_self_convention)]
- fn to_byte_array(data: &[T::T]) -> Vec<u8>;
- }
-
- macro_rules! to_byte_array_impl {
- ($ty: ty) => {
- impl ToByteArray<$ty> for $ty {
- fn to_byte_array(data: &[<$ty as DataType>::T]) -> Vec<u8> {
- <$ty as DataType>::T::slice_as_bytes(data).to_vec()
- }
- }
- };
- }
-
- to_byte_array_impl!(Int32Type);
- to_byte_array_impl!(Int64Type);
- to_byte_array_impl!(FloatType);
- to_byte_array_impl!(DoubleType);
-
- impl ToByteArray<BoolType> for BoolType {
- fn to_byte_array(data: &[bool]) -> Vec<u8> {
- let mut v = vec![];
- for i in 0..data.len() {
- if i % 8 == 0 {
- v.push(0);
- }
- if data[i] {
- set_array_bit(&mut v[..], i);
- }
- }
- v
- }
- }
-
- impl ToByteArray<Int96Type> for Int96Type {
- fn to_byte_array(data: &[Int96]) -> Vec<u8> {
- let mut v = vec![];
- for d in data {
- v.extend_from_slice(d.as_bytes());
- }
- v
- }
- }
-
- impl ToByteArray<ByteArrayType> for ByteArrayType {
- fn to_byte_array(data: &[ByteArray]) -> Vec<u8> {
- let mut v = vec![];
- for d in data {
- let buf = d.data();
- let len = &usize_to_bytes(buf.len());
- v.extend_from_slice(len);
- v.extend(buf);
- }
- v
- }
- }
-
- impl ToByteArray<FixedLenByteArrayType> for FixedLenByteArrayType {
- fn to_byte_array(data: &[FixedLenByteArray]) -> Vec<u8> {
- let mut v = vec![];
- for d in data {
- let buf = d.data();
- v.extend(buf);
- }
- v
- }
- }
-}
diff --git a/parquet/src/encodings/encoding.rs b/parquet/src/encodings/encoding.rs
deleted file mode 100644
index d042738..0000000
--- a/parquet/src/encodings/encoding.rs
+++ /dev/null
@@ -1,1334 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-//! Contains all supported encoders for Parquet.
-
-use std::{cmp, io::Write, marker::PhantomData};
-
-use crate::basic::*;
-use crate::data_type::private::ParquetValueType;
-use crate::data_type::*;
-use crate::encodings::rle::RleEncoder;
-use crate::errors::{ParquetError, Result};
-use crate::schema::types::ColumnDescPtr;
-use crate::util::{
- bit_util::{self, log2, num_required_bits, BitWriter},
- hash_util,
- memory::{Buffer, ByteBuffer, ByteBufferPtr, MemTrackerPtr},
-};
-
-// ----------------------------------------------------------------------
-// Encoders
-
-/// An Parquet encoder for the data type `T`.
-///
-/// Currently this allocates internal buffers for the encoded values. After done putting
-/// values, caller should call `flush_buffer()` to get an immutable buffer pointer.
-pub trait Encoder<T: DataType> {
- /// Encodes data from `values`.
- fn put(&mut self, values: &[T::T]) -> Result<()>;
-
- /// Encodes data from `values`, which contains spaces for null values, that is
- /// identified by `valid_bits`.
- ///
- /// Returns the number of non-null values encoded.
- fn put_spaced(&mut self, values: &[T::T], valid_bits: &[u8]) -> Result<usize> {
- let num_values = values.len();
- let mut buffer = Vec::with_capacity(num_values);
- // TODO: this is pretty inefficient. Revisit in future.
- for i in 0..num_values {
- if bit_util::get_bit(valid_bits, i) {
- buffer.push(values[i].clone());
- }
- }
- self.put(&buffer[..])?;
- Ok(buffer.len())
- }
-
- /// Returns the encoding type of this encoder.
- fn encoding(&self) -> Encoding;
-
- /// Returns an estimate of the encoded data, in bytes.
- /// Method call must be O(1).
- fn estimated_data_encoded_size(&self) -> usize;
-
- /// Flushes the underlying byte buffer that's being processed by this encoder, and
- /// return the immutable copy of it. This will also reset the internal state.
- fn flush_buffer(&mut self) -> Result<ByteBufferPtr>;
-}
-
-/// Gets a encoder for the particular data type `T` and encoding `encoding`. Memory usage
-/// for the encoder instance is tracked by `mem_tracker`.
-pub fn get_encoder<T: DataType>(
- desc: ColumnDescPtr,
- encoding: Encoding,
- mem_tracker: MemTrackerPtr,
-) -> Result<Box<dyn Encoder<T>>> {
- let encoder: Box<dyn Encoder<T>> = match encoding {
- Encoding::PLAIN => Box::new(PlainEncoder::new(desc, mem_tracker, vec![])),
- Encoding::RLE_DICTIONARY | Encoding::PLAIN_DICTIONARY => {
- return Err(general_err!(
- "Cannot initialize this encoding through this function"
- ));
- }
- Encoding::RLE => Box::new(RleValueEncoder::new()),
- Encoding::DELTA_BINARY_PACKED => Box::new(DeltaBitPackEncoder::new()),
- Encoding::DELTA_LENGTH_BYTE_ARRAY => Box::new(DeltaLengthByteArrayEncoder::new()),
- Encoding::DELTA_BYTE_ARRAY => Box::new(DeltaByteArrayEncoder::new()),
- e => return Err(nyi_err!("Encoding {} is not supported", e)),
- };
- Ok(encoder)
-}
-
-// ----------------------------------------------------------------------
-// Plain encoding
-
-/// Plain encoding that supports all types.
-/// Values are encoded back to back.
-/// The plain encoding is used whenever a more efficient encoding can not be used.
-/// It stores the data in the following format:
-/// - BOOLEAN - 1 bit per value, 0 is false; 1 is true.
-/// - INT32 - 4 bytes per value, stored as little-endian.
-/// - INT64 - 8 bytes per value, stored as little-endian.
-/// - FLOAT - 4 bytes per value, stored as IEEE little-endian.
-/// - DOUBLE - 8 bytes per value, stored as IEEE little-endian.
-/// - BYTE_ARRAY - 4 byte length stored as little endian, followed by bytes.
-/// - FIXED_LEN_BYTE_ARRAY - just the bytes are stored.
-pub struct PlainEncoder<T: DataType> {
- buffer: ByteBuffer,
- bit_writer: BitWriter,
- desc: ColumnDescPtr,
- _phantom: PhantomData<T>,
-}
-
-impl<T: DataType> PlainEncoder<T> {
- /// Creates new plain encoder.
- pub fn new(desc: ColumnDescPtr, mem_tracker: MemTrackerPtr, vec: Vec<u8>) -> Self {
- let mut byte_buffer = ByteBuffer::new().with_mem_tracker(mem_tracker);
- byte_buffer.set_data(vec);
- Self {
- buffer: byte_buffer,
- bit_writer: BitWriter::new(256),
- desc,
- _phantom: PhantomData,
- }
- }
-}
-
-impl<T: DataType> Encoder<T> for PlainEncoder<T> {
- // Performance Note:
- // As far as can be seen these functions are rarely called and as such we can hint to the
- // compiler that they dont need to be folded into hot locations in the final output.
- #[cold]
- fn encoding(&self) -> Encoding {
- Encoding::PLAIN
- }
-
- fn estimated_data_encoded_size(&self) -> usize {
- self.buffer.size() + self.bit_writer.bytes_written()
- }
-
- #[inline]
- fn flush_buffer(&mut self) -> Result<ByteBufferPtr> {
- self.buffer.write_all(self.bit_writer.flush_buffer())?;
- self.buffer.flush()?;
- self.bit_writer.clear();
-
- Ok(self.buffer.consume())
- }
-
- #[inline]
- fn put(&mut self, values: &[T::T]) -> Result<()> {
- T::T::encode(values, &mut self.buffer, &mut self.bit_writer)?;
- Ok(())
- }
-}
-
-// ----------------------------------------------------------------------
-// Dictionary encoding
-
-const INITIAL_HASH_TABLE_SIZE: usize = 1024;
-const MAX_HASH_LOAD: f32 = 0.7;
-const HASH_SLOT_EMPTY: i32 = -1;
-
-/// Dictionary encoder.
-/// The dictionary encoding builds a dictionary of values encountered in a given column.
-/// The dictionary page is written first, before the data pages of the column chunk.
-///
-/// Dictionary page format: the entries in the dictionary - in dictionary order -
-/// using the plain encoding.
-///
-/// Data page format: the bit width used to encode the entry ids stored as 1 byte
-/// (max bit width = 32), followed by the values encoded using RLE/Bit packed described
-/// above (with the given bit width).
-pub struct DictEncoder<T: DataType> {
- // Descriptor for the column to be encoded.
- desc: ColumnDescPtr,
-
- // Size of the table. **Must be** a power of 2.
- hash_table_size: usize,
-
- // Store `hash_table_size` - 1, so that `j & mod_bitmask` is equivalent to
- // `j % hash_table_size`, but uses far fewer CPU cycles.
- mod_bitmask: u32,
-
- // Stores indices which map (many-to-one) to the values in the `uniques` array.
- // Here we are using fix-sized array with linear probing.
- // A slot with `HASH_SLOT_EMPTY` indicates the slot is not currently occupied.
- hash_slots: Buffer<i32>,
-
- // Indices that have not yet be written out by `write_indices()`.
- buffered_indices: Buffer<i32>,
-
- // The unique observed values.
- uniques: Buffer<T::T>,
-
- // Size in bytes needed to encode this dictionary.
- uniques_size_in_bytes: usize,
-
- // Tracking memory usage for the various data structures in this struct.
- mem_tracker: MemTrackerPtr,
-}
-
-impl<T: DataType> DictEncoder<T> {
- /// Creates new dictionary encoder.
- pub fn new(desc: ColumnDescPtr, mem_tracker: MemTrackerPtr) -> Self {
- let mut slots = Buffer::new().with_mem_tracker(mem_tracker.clone());
- slots.resize(INITIAL_HASH_TABLE_SIZE, -1);
- Self {
- desc,
- hash_table_size: INITIAL_HASH_TABLE_SIZE,
- mod_bitmask: (INITIAL_HASH_TABLE_SIZE - 1) as u32,
- hash_slots: slots,
- buffered_indices: Buffer::new().with_mem_tracker(mem_tracker.clone()),
- uniques: Buffer::new().with_mem_tracker(mem_tracker.clone()),
- uniques_size_in_bytes: 0,
- mem_tracker,
- }
- }
-
- /// Returns true if dictionary entries are sorted, false otherwise.
- #[inline]
- pub fn is_sorted(&self) -> bool {
- // Sorting is not supported currently.
- false
- }
-
- /// Returns number of unique values (keys) in the dictionary.
- pub fn num_entries(&self) -> usize {
- self.uniques.size()
- }
-
- /// Returns size of unique values (keys) in the dictionary, in bytes.
- pub fn dict_encoded_size(&self) -> usize {
- self.uniques_size_in_bytes
- }
-
- /// Writes out the dictionary values with PLAIN encoding in a byte buffer, and return
- /// the result.
- #[inline]
- pub fn write_dict(&self) -> Result<ByteBufferPtr> {
- let mut plain_encoder =
- PlainEncoder::<T>::new(self.desc.clone(), self.mem_tracker.clone(), vec![]);
- plain_encoder.put(self.uniques.data())?;
- plain_encoder.flush_buffer()
- }
-
- /// Writes out the dictionary values with RLE encoding in a byte buffer, and return
- /// the result.
- pub fn write_indices(&mut self) -> Result<ByteBufferPtr> {
- // TODO: the caller should allocate the buffer
- let buffer_len = self.estimated_data_encoded_size();
- let mut buffer: Vec<u8> = vec![0; buffer_len as usize];
- buffer[0] = self.bit_width() as u8;
- self.mem_tracker.alloc(buffer.capacity() as i64);
-
- // Write bit width in the first byte
- buffer.write_all((self.bit_width() as u8).as_bytes())?;
- let mut encoder = RleEncoder::new_from_buf(self.bit_width(), buffer, 1);
- for index in self.buffered_indices.data() {
- if !encoder.put(*index as u64)? {
- return Err(general_err!("Encoder doesn't have enough space"));
- }
- }
- self.buffered_indices.clear();
- Ok(ByteBufferPtr::new(encoder.consume()?))
- }
-
- #[inline]
- #[allow(clippy::unnecessary_wraps)]
- fn put_one(&mut self, value: &T::T) -> Result<()> {
- let mut j = (hash_util::hash(value, 0) & self.mod_bitmask) as usize;
- let mut index = self.hash_slots[j];
-
- while index != HASH_SLOT_EMPTY && self.uniques[index as usize] != *value {
- j += 1;
- if j == self.hash_table_size {
- j = 0;
- }
- index = self.hash_slots[j];
- }
-
- if index == HASH_SLOT_EMPTY {
- index = self.insert_fresh_slot(j, value.clone());
- }
-
- self.buffered_indices.push(index);
- Ok(())
- }
-
- #[inline(never)]
- fn insert_fresh_slot(&mut self, slot: usize, value: T::T) -> i32 {
- let index = self.uniques.size() as i32;
- self.hash_slots[slot] = index;
-
- let (base_size, num_elements) = value.dict_encoding_size();
-
- let unique_size = match T::get_physical_type() {
- Type::BYTE_ARRAY => base_size + num_elements,
- Type::FIXED_LEN_BYTE_ARRAY => self.desc.type_length() as usize,
- _ => base_size,
- };
-
- self.uniques_size_in_bytes += unique_size;
- self.uniques.push(value);
-
- if self.uniques.size() > (self.hash_table_size as f32 * MAX_HASH_LOAD) as usize {
- self.double_table_size();
- }
-
- index
- }
-
- #[inline]
- fn bit_width(&self) -> u8 {
- let num_entries = self.uniques.size();
- if num_entries == 0 {
- 0
- } else if num_entries == 1 {
- 1
- } else {
- log2(num_entries as u64) as u8
- }
- }
-
- fn double_table_size(&mut self) {
- let new_size = self.hash_table_size * 2;
- let mut new_hash_slots = Buffer::new().with_mem_tracker(self.mem_tracker.clone());
- new_hash_slots.resize(new_size, HASH_SLOT_EMPTY);
- for i in 0..self.hash_table_size {
- let index = self.hash_slots[i];
- if index == HASH_SLOT_EMPTY {
- continue;
- }
- let value = &self.uniques[index as usize];
- let mut j = (hash_util::hash(value, 0) & ((new_size - 1) as u32)) as usize;
- let mut slot = new_hash_slots[j];
- while slot != HASH_SLOT_EMPTY && self.uniques[slot as usize] != *value {
- j += 1;
- if j == new_size {
- j = 0;
- }
- slot = new_hash_slots[j];
- }
-
- new_hash_slots[j] = index;
- }
-
- self.hash_table_size = new_size;
- self.mod_bitmask = (new_size - 1) as u32;
- self.hash_slots = new_hash_slots;
- }
-}
-
-impl<T: DataType> Encoder<T> for DictEncoder<T> {
- #[inline]
- fn put(&mut self, values: &[T::T]) -> Result<()> {
- for i in values {
- self.put_one(&i)?
- }
- Ok(())
- }
-
- // Performance Note:
- // As far as can be seen these functions are rarely called and as such we can hint to the
- // compiler that they dont need to be folded into hot locations in the final output.
- #[cold]
- fn encoding(&self) -> Encoding {
- Encoding::PLAIN_DICTIONARY
- }
-
- #[inline]
- fn estimated_data_encoded_size(&self) -> usize {
- let bit_width = self.bit_width();
- 1 + RleEncoder::min_buffer_size(bit_width)
- + RleEncoder::max_buffer_size(bit_width, self.buffered_indices.size())
- }
-
- #[inline]
- fn flush_buffer(&mut self) -> Result<ByteBufferPtr> {
- self.write_indices()
- }
-}
-
-// ----------------------------------------------------------------------
-// RLE encoding
-
-const DEFAULT_RLE_BUFFER_LEN: usize = 1024;
-
-/// RLE/Bit-Packing hybrid encoding for values.
-/// Currently is used only for data pages v2 and supports boolean types.
-pub struct RleValueEncoder<T: DataType> {
- // Buffer with raw values that we collect,
- // when flushing buffer they are encoded using RLE encoder
- encoder: Option<RleEncoder>,
- _phantom: PhantomData<T>,
-}
-
-impl<T: DataType> RleValueEncoder<T> {
- /// Creates new rle value encoder.
- pub fn new() -> Self {
- Self {
- encoder: None,
- _phantom: PhantomData,
- }
- }
-}
-
-impl<T: DataType> Encoder<T> for RleValueEncoder<T> {
- #[inline]
- fn put(&mut self, values: &[T::T]) -> Result<()> {
- ensure_phys_ty!(Type::BOOLEAN, "RleValueEncoder only supports BoolType");
-
- if self.encoder.is_none() {
- self.encoder = Some(RleEncoder::new(1, DEFAULT_RLE_BUFFER_LEN));
- }
- let rle_encoder = self.encoder.as_mut().unwrap();
- for value in values {
- let value = value.as_u64()?;
- if !rle_encoder.put(value)? {
- return Err(general_err!("RLE buffer is full"));
- }
- }
- Ok(())
- }
-
- // Performance Note:
- // As far as can be seen these functions are rarely called and as such we can hint to the
- // compiler that they dont need to be folded into hot locations in the final output.
- #[cold]
- fn encoding(&self) -> Encoding {
- Encoding::RLE
- }
-
- #[inline]
- fn estimated_data_encoded_size(&self) -> usize {
- match self.encoder {
- Some(ref enc) => enc.len(),
- None => 0,
- }
- }
-
- #[inline]
- fn flush_buffer(&mut self) -> Result<ByteBufferPtr> {
- ensure_phys_ty!(Type::BOOLEAN, "RleValueEncoder only supports BoolType");
- let rle_encoder = self
- .encoder
- .as_mut()
- .expect("RLE value encoder is not initialized");
-
- // Flush all encoder buffers and raw values
- let encoded_data = {
- let buf = rle_encoder.flush_buffer()?;
-
- // Note that buf does not have any offset, all data is encoded bytes
- let len = (buf.len() as i32).to_le();
- let len_bytes = len.as_bytes();
- let mut encoded_data = vec![];
- encoded_data.extend_from_slice(len_bytes);
- encoded_data.extend_from_slice(buf);
- encoded_data
- };
- // Reset rle encoder for the next batch
- rle_encoder.clear();
-
- Ok(ByteBufferPtr::new(encoded_data))
- }
-}
-
-// ----------------------------------------------------------------------
-// DELTA_BINARY_PACKED encoding
-
-const MAX_PAGE_HEADER_WRITER_SIZE: usize = 32;
-const MAX_BIT_WRITER_SIZE: usize = 10 * 1024 * 1024;
-const DEFAULT_BLOCK_SIZE: usize = 128;
-const DEFAULT_NUM_MINI_BLOCKS: usize = 4;
-
-/// Delta bit packed encoder.
-/// Consists of a header followed by blocks of delta encoded values binary packed.
-///
-/// Delta-binary-packing:
-/// ```shell
-/// [page-header] [block 1], [block 2], ... [block N]
-/// ```
-///
-/// Each page header consists of:
-/// ```shell
-/// [block size] [number of miniblocks in a block] [total value count] [first value]
-/// ```
-///
-/// Each block consists of:
-/// ```shell
-/// [min delta] [list of bitwidths of miniblocks] [miniblocks]
-/// ```
-///
-/// Current implementation writes values in `put` method, multiple calls to `put` to
-/// existing block or start new block if block size is exceeded. Calling `flush_buffer`
-/// writes out all data and resets internal state, including page header.
-///
-/// Supports only INT32 and INT64.
-pub struct DeltaBitPackEncoder<T: DataType> {
- page_header_writer: BitWriter,
- bit_writer: BitWriter,
- total_values: usize,
- first_value: i64,
- current_value: i64,
- block_size: usize,
- mini_block_size: usize,
- num_mini_blocks: usize,
- values_in_block: usize,
- deltas: Vec<i64>,
- _phantom: PhantomData<T>,
-}
-
-impl<T: DataType> DeltaBitPackEncoder<T> {
- /// Creates new delta bit packed encoder.
- pub fn new() -> Self {
- let block_size = DEFAULT_BLOCK_SIZE;
- let num_mini_blocks = DEFAULT_NUM_MINI_BLOCKS;
- let mini_block_size = block_size / num_mini_blocks;
- assert!(mini_block_size % 8 == 0);
- Self::assert_supported_type();
-
- DeltaBitPackEncoder {
- page_header_writer: BitWriter::new(MAX_PAGE_HEADER_WRITER_SIZE),
- bit_writer: BitWriter::new(MAX_BIT_WRITER_SIZE),
- total_values: 0,
- first_value: 0,
- current_value: 0, // current value to keep adding deltas
- block_size, // can write fewer values than block size for last block
- mini_block_size,
- num_mini_blocks,
- values_in_block: 0, // will be at most block_size
- deltas: vec![0; block_size],
- _phantom: PhantomData,
- }
- }
-
- /// Writes page header for blocks, this method is invoked when we are done encoding
- /// values. It is also okay to encode when no values have been provided
- fn write_page_header(&mut self) {
- // We ignore the result of each 'put' operation, because
- // MAX_PAGE_HEADER_WRITER_SIZE is chosen to fit all header values and
- // guarantees that writes will not fail.
-
- // Write the size of each block
- self.page_header_writer.put_vlq_int(self.block_size as u64);
- // Write the number of mini blocks
- self.page_header_writer
- .put_vlq_int(self.num_mini_blocks as u64);
- // Write the number of all values (including non-encoded first value)
- self.page_header_writer
- .put_vlq_int(self.total_values as u64);
- // Write first value
- self.page_header_writer.put_zigzag_vlq_int(self.first_value);
- }
-
- // Write current delta buffer (<= 'block size' values) into bit writer
- #[inline(never)]
- fn flush_block_values(&mut self) -> Result<()> {
- if self.values_in_block == 0 {
- return Ok(());
- }
-
- let mut min_delta = i64::max_value();
- for i in 0..self.values_in_block {
- min_delta = cmp::min(min_delta, self.deltas[i]);
- }
-
- // Write min delta
- self.bit_writer.put_zigzag_vlq_int(min_delta);
-
- // Slice to store bit width for each mini block
- let offset = self.bit_writer.skip(self.num_mini_blocks)?;
-
- for i in 0..self.num_mini_blocks {
- // Find how many values we need to encode - either block size or whatever
- // values left
- let n = cmp::min(self.mini_block_size, self.values_in_block);
- if n == 0 {
- break;
- }
-
- // Compute the max delta in current mini block
- let mut max_delta = i64::min_value();
- for j in 0..n {
- max_delta =
- cmp::max(max_delta, self.deltas[i * self.mini_block_size + j]);
- }
-
- // Compute bit width to store (max_delta - min_delta)
- let bit_width = num_required_bits(self.subtract_u64(max_delta, min_delta));
- self.bit_writer.write_at(offset + i, bit_width as u8);
-
- // Encode values in current mini block using min_delta and bit_width
- for j in 0..n {
- let packed_value = self
- .subtract_u64(self.deltas[i * self.mini_block_size + j], min_delta);
- self.bit_writer.put_value(packed_value, bit_width);
- }
-
- // Pad the last block (n < mini_block_size)
- for _ in n..self.mini_block_size {
- self.bit_writer.put_value(0, bit_width);
- }
-
- self.values_in_block -= n;
- }
-
- assert!(
- self.values_in_block == 0,
- "Expected 0 values in block, found {}",
- self.values_in_block
- );
- Ok(())
- }
-}
-
-// Implementation is shared between Int32Type and Int64Type,
-// see `DeltaBitPackEncoderConversion` below for specifics.
-impl<T: DataType> Encoder<T> for DeltaBitPackEncoder<T> {
- fn put(&mut self, values: &[T::T]) -> Result<()> {
- if values.is_empty() {
- return Ok(());
- }
-
- // Define values to encode, initialize state
- let mut idx = if self.total_values == 0 {
- self.first_value = self.as_i64(values, 0);
- self.current_value = self.first_value;
- 1
- } else {
- 0
- };
- // Add all values (including first value)
- self.total_values += values.len();
-
- // Write block
- while idx < values.len() {
- let value = self.as_i64(values, idx);
- self.deltas[self.values_in_block] = self.subtract(value, self.current_value);
- self.current_value = value;
- idx += 1;
- self.values_in_block += 1;
- if self.values_in_block == self.block_size {
- self.flush_block_values()?;
- }
- }
- Ok(())
- }
-
- // Performance Note:
- // As far as can be seen these functions are rarely called and as such we can hint to the
- // compiler that they dont need to be folded into hot locations in the final output.
- #[cold]
- fn encoding(&self) -> Encoding {
- Encoding::DELTA_BINARY_PACKED
- }
-
- fn estimated_data_encoded_size(&self) -> usize {
- self.bit_writer.bytes_written()
- }
-
- fn flush_buffer(&mut self) -> Result<ByteBufferPtr> {
- // Write remaining values
- self.flush_block_values()?;
- // Write page header with total values
- self.write_page_header();
-
- let mut buffer = ByteBuffer::new();
- buffer.write_all(self.page_header_writer.flush_buffer())?;
- buffer.write_all(self.bit_writer.flush_buffer())?;
- buffer.flush()?;
-
- // Reset state
- self.page_header_writer.clear();
- self.bit_writer.clear();
- self.total_values = 0;
- self.first_value = 0;
- self.current_value = 0;
- self.values_in_block = 0;
-
- Ok(buffer.consume())
- }
-}
-
-/// Helper trait to define specific conversions and subtractions when computing deltas
-trait DeltaBitPackEncoderConversion<T: DataType> {
- // Method should panic if type is not supported, otherwise no-op
- fn assert_supported_type();
-
- fn as_i64(&self, values: &[T::T], index: usize) -> i64;
-
- fn subtract(&self, left: i64, right: i64) -> i64;
-
- fn subtract_u64(&self, left: i64, right: i64) -> u64;
-}
-
-impl<T: DataType> DeltaBitPackEncoderConversion<T> for DeltaBitPackEncoder<T> {
- #[inline]
- fn assert_supported_type() {
- ensure_phys_ty!(
- Type::INT32 | Type::INT64,
- "DeltaBitPackDecoder only supports Int32Type and Int64Type"
- );
- }
-
- #[inline]
- fn as_i64(&self, values: &[T::T], index: usize) -> i64 {
- values[index]
- .as_i64()
- .expect("DeltaBitPackDecoder only supports Int32Type and Int64Type")
- }
-
- #[inline]
- fn subtract(&self, left: i64, right: i64) -> i64 {
- // It is okay for values to overflow, wrapping_sub wrapping around at the boundary
- match T::get_physical_type() {
- Type::INT32 => (left as i32).wrapping_sub(right as i32) as i64,
- Type::INT64 => left.wrapping_sub(right),
- _ => panic!("DeltaBitPackDecoder only supports Int32Type and Int64Type"),
- }
- }
-
- #[inline]
- fn subtract_u64(&self, left: i64, right: i64) -> u64 {
- match T::get_physical_type() {
- // Conversion of i32 -> u32 -> u64 is to avoid non-zero left most bytes in int repr
- Type::INT32 => (left as i32).wrapping_sub(right as i32) as u32 as u64,
- Type::INT64 => left.wrapping_sub(right) as u64,
- _ => panic!("DeltaBitPackDecoder only supports Int32Type and Int64Type"),
- }
- }
-}
-
-// ----------------------------------------------------------------------
-// DELTA_LENGTH_BYTE_ARRAY encoding
-
-/// Encoding for byte arrays to separate the length values and the data.
-/// The lengths are encoded using DELTA_BINARY_PACKED encoding, data is
-/// stored as raw bytes.
-pub struct DeltaLengthByteArrayEncoder<T: DataType> {
- // length encoder
- len_encoder: DeltaBitPackEncoder<Int32Type>,
- // byte array data
- data: Vec<ByteArray>,
- // data size in bytes of encoded values
- encoded_size: usize,
- _phantom: PhantomData<T>,
-}
-
-impl<T: DataType> DeltaLengthByteArrayEncoder<T> {
- /// Creates new delta length byte array encoder.
- pub fn new() -> Self {
- Self {
- len_encoder: DeltaBitPackEncoder::new(),
- data: vec![],
- encoded_size: 0,
- _phantom: PhantomData,
- }
- }
-}
-
-impl<T: DataType> Encoder<T> for DeltaLengthByteArrayEncoder<T> {
- fn put(&mut self, values: &[T::T]) -> Result<()> {
- ensure_phys_ty!(
- Type::BYTE_ARRAY | Type::FIXED_LEN_BYTE_ARRAY,
- "DeltaLengthByteArrayEncoder only supports ByteArrayType"
- );
-
- let val_it = || {
- values
- .iter()
- .map(|x| x.as_any().downcast_ref::<ByteArray>().unwrap())
- };
-
- let lengths: Vec<i32> =
- val_it().map(|byte_array| byte_array.len() as i32).collect();
- self.len_encoder.put(&lengths)?;
- for byte_array in val_it() {
- self.encoded_size += byte_array.len();
- self.data.push(byte_array.clone());
- }
-
- Ok(())
- }
-
- // Performance Note:
- // As far as can be seen these functions are rarely called and as such we can hint to the
- // compiler that they dont need to be folded into hot locations in the final output.
- #[cold]
- fn encoding(&self) -> Encoding {
- Encoding::DELTA_LENGTH_BYTE_ARRAY
- }
-
- fn estimated_data_encoded_size(&self) -> usize {
- self.len_encoder.estimated_data_encoded_size() + self.encoded_size
- }
-
- fn flush_buffer(&mut self) -> Result<ByteBufferPtr> {
- ensure_phys_ty!(
- Type::BYTE_ARRAY | Type::FIXED_LEN_BYTE_ARRAY,
- "DeltaLengthByteArrayEncoder only supports ByteArrayType"
- );
-
- let mut total_bytes = vec![];
- let lengths = self.len_encoder.flush_buffer()?;
- total_bytes.extend_from_slice(lengths.data());
- self.data.iter().for_each(|byte_array| {
- total_bytes.extend_from_slice(byte_array.data());
- });
- self.data.clear();
- self.encoded_size = 0;
-
- Ok(ByteBufferPtr::new(total_bytes))
- }
-}
-
-// ----------------------------------------------------------------------
-// DELTA_BYTE_ARRAY encoding
-
-/// Encoding for byte arrays, prefix lengths are encoded using DELTA_BINARY_PACKED
-/// encoding, followed by suffixes with DELTA_LENGTH_BYTE_ARRAY encoding.
-pub struct DeltaByteArrayEncoder<T: DataType> {
- prefix_len_encoder: DeltaBitPackEncoder<Int32Type>,
- suffix_writer: DeltaLengthByteArrayEncoder<ByteArrayType>,
- previous: Vec<u8>,
- _phantom: PhantomData<T>,
-}
-
-impl<T: DataType> DeltaByteArrayEncoder<T> {
- /// Creates new delta byte array encoder.
- pub fn new() -> Self {
- Self {
- prefix_len_encoder: DeltaBitPackEncoder::new(),
- suffix_writer: DeltaLengthByteArrayEncoder::new(),
- previous: vec![],
- _phantom: PhantomData,
- }
- }
-}
-
-impl<T: DataType> Encoder<T> for DeltaByteArrayEncoder<T> {
- fn put(&mut self, values: &[T::T]) -> Result<()> {
- let mut prefix_lengths: Vec<i32> = vec![];
- let mut suffixes: Vec<ByteArray> = vec![];
-
- let values = values.iter()
- .map(|x| x.as_any())
- .map(|x| match T::get_physical_type() {
- Type::BYTE_ARRAY => x.downcast_ref::<ByteArray>().unwrap(),
- Type::FIXED_LEN_BYTE_ARRAY => x.downcast_ref::<FixedLenByteArray>().unwrap(),
- _ => panic!(
- "DeltaByteArrayEncoder only supports ByteArrayType and FixedLenByteArrayType"
- )
- });
-
- for byte_array in values {
- let current = byte_array.data();
- // Maximum prefix length that is shared between previous value and current
- // value
- let prefix_len = cmp::min(self.previous.len(), current.len());
- let mut match_len = 0;
- while match_len < prefix_len && self.previous[match_len] == current[match_len]
- {
- match_len += 1;
- }
- prefix_lengths.push(match_len as i32);
- suffixes.push(byte_array.slice(match_len, byte_array.len() - match_len));
- // Update previous for the next prefix
- self.previous.clear();
- self.previous.extend_from_slice(current);
- }
- self.prefix_len_encoder.put(&prefix_lengths)?;
- self.suffix_writer.put(&suffixes)?;
-
- Ok(())
- }
-
- // Performance Note:
- // As far as can be seen these functions are rarely called and as such we can hint to the
- // compiler that they dont need to be folded into hot locations in the final output.
- #[cold]
- fn encoding(&self) -> Encoding {
- Encoding::DELTA_BYTE_ARRAY
- }
-
- fn estimated_data_encoded_size(&self) -> usize {
- self.prefix_len_encoder.estimated_data_encoded_size()
- + self.suffix_writer.estimated_data_encoded_size()
- }
-
- fn flush_buffer(&mut self) -> Result<ByteBufferPtr> {
- match T::get_physical_type() {
- Type::BYTE_ARRAY | Type::FIXED_LEN_BYTE_ARRAY => {
- // TODO: investigate if we can merge lengths and suffixes
- // without copying data into new vector.
- let mut total_bytes = vec![];
- // Insert lengths ...
- let lengths = self.prefix_len_encoder.flush_buffer()?;
- total_bytes.extend_from_slice(lengths.data());
- // ... followed by suffixes
- let suffixes = self.suffix_writer.flush_buffer()?;
- total_bytes.extend_from_slice(suffixes.data());
-
- self.previous.clear();
- Ok(ByteBufferPtr::new(total_bytes))
- }
- _ => panic!(
- "DeltaByteArrayEncoder only supports ByteArrayType and FixedLenByteArrayType"
- )
- }
- }
-}
-
-#[cfg(test)]
-mod tests {
- use super::*;
-
- use std::sync::Arc;
-
- use crate::decoding::{get_decoder, Decoder, DictDecoder, PlainDecoder};
- use crate::schema::types::{
- ColumnDescPtr, ColumnDescriptor, ColumnPath, Type as SchemaType,
- };
- use crate::util::{
- memory::MemTracker,
- test_common::{random_bytes, RandGen},
- };
-
- const TEST_SET_SIZE: usize = 1024;
-
- #[test]
- fn test_get_encoders() {
- // supported encodings
- create_and_check_encoder::<Int32Type>(Encoding::PLAIN, None);
- create_and_check_encoder::<Int32Type>(Encoding::DELTA_BINARY_PACKED, None);
- create_and_check_encoder::<Int32Type>(Encoding::DELTA_LENGTH_BYTE_ARRAY, None);
- create_and_check_encoder::<Int32Type>(Encoding::DELTA_BYTE_ARRAY, None);
- create_and_check_encoder::<BoolType>(Encoding::RLE, None);
-
- // error when initializing
- create_and_check_encoder::<Int32Type>(
- Encoding::RLE_DICTIONARY,
- Some(general_err!(
- "Cannot initialize this encoding through this function"
- )),
- );
- create_and_check_encoder::<Int32Type>(
- Encoding::PLAIN_DICTIONARY,
- Some(general_err!(
- "Cannot initialize this encoding through this function"
- )),
- );
-
- // unsupported
- create_and_check_encoder::<Int32Type>(
- Encoding::BIT_PACKED,
- Some(nyi_err!("Encoding BIT_PACKED is not supported")),
- );
- }
-
- #[test]
- fn test_bool() {
- BoolType::test(Encoding::PLAIN, TEST_SET_SIZE, -1);
- BoolType::test(Encoding::PLAIN_DICTIONARY, TEST_SET_SIZE, -1);
- BoolType::test(Encoding::RLE, TEST_SET_SIZE, -1);
- }
-
- #[test]
- fn test_i32() {
- Int32Type::test(Encoding::PLAIN, TEST_SET_SIZE, -1);
- Int32Type::test(Encoding::PLAIN_DICTIONARY, TEST_SET_SIZE, -1);
- Int32Type::test(Encoding::DELTA_BINARY_PACKED, TEST_SET_SIZE, -1);
- }
-
- #[test]
- fn test_i64() {
- Int64Type::test(Encoding::PLAIN, TEST_SET_SIZE, -1);
- Int64Type::test(Encoding::PLAIN_DICTIONARY, TEST_SET_SIZE, -1);
- Int64Type::test(Encoding::DELTA_BINARY_PACKED, TEST_SET_SIZE, -1);
- }
-
- #[test]
- fn test_i96() {
- Int96Type::test(Encoding::PLAIN, TEST_SET_SIZE, -1);
- Int96Type::test(Encoding::PLAIN_DICTIONARY, TEST_SET_SIZE, -1);
- }
-
- #[test]
- fn test_float() {
- FloatType::test(Encoding::PLAIN, TEST_SET_SIZE, -1);
- FloatType::test(Encoding::PLAIN_DICTIONARY, TEST_SET_SIZE, -1);
- }
-
- #[test]
- fn test_double() {
- DoubleType::test(Encoding::PLAIN, TEST_SET_SIZE, -1);
- DoubleType::test(Encoding::PLAIN_DICTIONARY, TEST_SET_SIZE, -1);
- }
-
- #[test]
- fn test_byte_array() {
- ByteArrayType::test(Encoding::PLAIN, TEST_SET_SIZE, -1);
- ByteArrayType::test(Encoding::PLAIN_DICTIONARY, TEST_SET_SIZE, -1);
- ByteArrayType::test(Encoding::DELTA_LENGTH_BYTE_ARRAY, TEST_SET_SIZE, -1);
- ByteArrayType::test(Encoding::DELTA_BYTE_ARRAY, TEST_SET_SIZE, -1);
- }
-
- #[test]
- fn test_fixed_lenbyte_array() {
- FixedLenByteArrayType::test(Encoding::PLAIN, TEST_SET_SIZE, 100);
- FixedLenByteArrayType::test(Encoding::PLAIN_DICTIONARY, TEST_SET_SIZE, 100);
- FixedLenByteArrayType::test(Encoding::DELTA_BYTE_ARRAY, TEST_SET_SIZE, 100);
- }
-
- #[test]
- fn test_dict_encoded_size() {
- fn run_test<T: DataType>(
- type_length: i32,
- values: &[T::T],
- expected_size: usize,
- ) {
- let mut encoder = create_test_dict_encoder::<T>(type_length);
- assert_eq!(encoder.dict_encoded_size(), 0);
- encoder.put(values).unwrap();
- assert_eq!(encoder.dict_encoded_size(), expected_size);
- // We do not reset encoded size of the dictionary keys after flush_buffer
- encoder.flush_buffer().unwrap();
- assert_eq!(encoder.dict_encoded_size(), expected_size);
- }
-
- // Only 2 variations of values 1 byte each
- run_test::<BoolType>(-1, &[true, false, true, false, true], 2);
- run_test::<Int32Type>(-1, &[1i32, 2i32, 3i32, 4i32, 5i32], 20);
- run_test::<Int64Type>(-1, &[1i64, 2i64, 3i64, 4i64, 5i64], 40);
- run_test::<FloatType>(-1, &[1f32, 2f32, 3f32, 4f32, 5f32], 20);
- run_test::<DoubleType>(-1, &[1f64, 2f64, 3f64, 4f64, 5f64], 40);
- // Int96: len + reference
- run_test::<Int96Type>(
- -1,
- &[Int96::from(vec![1, 2, 3]), Int96::from(vec![2, 3, 4])],
- 32,
- );
- run_test::<ByteArrayType>(
- -1,
- &[ByteArray::from("abcd"), ByteArray::from("efj")],
- 15,
- );
- run_test::<FixedLenByteArrayType>(
- 2,
- &[ByteArray::from("ab").into(), ByteArray::from("bc").into()],
- 4,
- );
- }
-
- #[test]
- fn test_estimated_data_encoded_size() {
- fn run_test<T: DataType>(
- encoding: Encoding,
- type_length: i32,
- values: &[T::T],
- initial_size: usize,
- max_size: usize,
- flush_size: usize,
- ) {
- let mut encoder = match encoding {
- Encoding::PLAIN_DICTIONARY | Encoding::RLE_DICTIONARY => {
- Box::new(create_test_dict_encoder::<T>(type_length))
- }
- _ => create_test_encoder::<T>(type_length, encoding),
- };
- assert_eq!(encoder.estimated_data_encoded_size(), initial_size);
-
- encoder.put(values).unwrap();
- assert_eq!(encoder.estimated_data_encoded_size(), max_size);
-
- encoder.flush_buffer().unwrap();
- assert_eq!(encoder.estimated_data_encoded_size(), flush_size);
- }
-
- // PLAIN
- run_test::<Int32Type>(Encoding::PLAIN, -1, &[123; 1024], 0, 4096, 0);
-
- // DICTIONARY
- // NOTE: The final size is almost the same because the dictionary entries are
- // preserved after encoded values have been written.
- run_test::<Int32Type>(Encoding::RLE_DICTIONARY, -1, &[123, 1024], 11, 68, 66);
-
- // DELTA_BINARY_PACKED
- run_test::<Int32Type>(Encoding::DELTA_BINARY_PACKED, -1, &[123; 1024], 0, 35, 0);
-
- // RLE
- let mut values = vec![];
- values.extend_from_slice(&[true; 16]);
- values.extend_from_slice(&[false; 16]);
- run_test::<BoolType>(Encoding::RLE, -1, &values, 0, 2, 0);
-
- // DELTA_LENGTH_BYTE_ARRAY
- run_test::<ByteArrayType>(
- Encoding::DELTA_LENGTH_BYTE_ARRAY,
- -1,
- &[ByteArray::from("ab"), ByteArray::from("abc")],
- 0,
- 5, // only value bytes, length encoder is not flushed yet
- 0,
- );
-
- // DELTA_BYTE_ARRAY
- run_test::<ByteArrayType>(
- Encoding::DELTA_BYTE_ARRAY,
- -1,
- &[ByteArray::from("ab"), ByteArray::from("abc")],
- 0,
- 3, // only suffix bytes, length encoder is not flushed yet
- 0,
- );
- }
-
- // See: https://github.com/sunchao/parquet-rs/issues/47
- #[test]
- fn test_issue_47() {
- let mut encoder =
- create_test_encoder::<ByteArrayType>(0, Encoding::DELTA_BYTE_ARRAY);
- let mut decoder =
- create_test_decoder::<ByteArrayType>(0, Encoding::DELTA_BYTE_ARRAY);
-
- let mut input = vec![];
- input.push(ByteArray::from("aa"));
- input.push(ByteArray::from("aaa"));
- input.push(ByteArray::from("aa"));
- input.push(ByteArray::from("aaa"));
- let mut output = vec![ByteArray::default(); input.len()];
-
- let mut result =
- put_and_get(&mut encoder, &mut decoder, &input[..2], &mut output[..2]);
- assert!(
- result.is_ok(),
- "first put_and_get() failed with: {}",
- result.unwrap_err()
- );
- result = put_and_get(&mut encoder, &mut decoder, &input[2..], &mut output[2..]);
- assert!(
- result.is_ok(),
- "second put_and_get() failed with: {}",
- result.unwrap_err()
- );
- assert_eq!(output, input);
- }
-
- trait EncodingTester<T: DataType> {
- fn test(enc: Encoding, total: usize, type_length: i32) {
- let result = match enc {
- Encoding::PLAIN_DICTIONARY | Encoding::RLE_DICTIONARY => {
- Self::test_dict_internal(total, type_length)
- }
- enc => Self::test_internal(enc, total, type_length),
- };
-
- assert!(
- result.is_ok(),
- "Expected result to be OK but got err:\n {}",
- result.unwrap_err()
- );
- }
-
- fn test_internal(enc: Encoding, total: usize, type_length: i32) -> Result<()>;
-
- fn test_dict_internal(total: usize, type_length: i32) -> Result<()>;
- }
-
- impl<T: DataType + RandGen<T>> EncodingTester<T> for T {
- fn test_internal(enc: Encoding, total: usize, type_length: i32) -> Result<()> {
- let mut encoder = create_test_encoder::<T>(type_length, enc);
- let mut decoder = create_test_decoder::<T>(type_length, enc);
- let mut values = <T as RandGen<T>>::gen_vec(type_length, total);
- let mut result_data = vec![T::T::default(); total];
-
- // Test put/get spaced.
- let num_bytes = bit_util::ceil(total as i64, 8);
- let valid_bits = random_bytes(num_bytes as usize);
- let values_written = encoder.put_spaced(&values[..], &valid_bits[..])?;
- let data = encoder.flush_buffer()?;
- decoder.set_data(data, values_written)?;
- let _ = decoder.get_spaced(
- &mut result_data[..],
- values.len() - values_written,
- &valid_bits[..],
- )?;
-
- // Check equality
- for i in 0..total {
- if bit_util::get_bit(&valid_bits[..], i) {
- assert_eq!(result_data[i], values[i]);
- } else {
- assert_eq!(result_data[i], T::T::default());
- }
- }
-
- let mut actual_total = put_and_get(
- &mut encoder,
- &mut decoder,
- &values[..],
- &mut result_data[..],
- )?;
- assert_eq!(actual_total, total);
- assert_eq!(result_data, values);
-
- // Encode more data after flush and test with decoder
-
- values = <T as RandGen<T>>::gen_vec(type_length, total);
- actual_total = put_and_get(
- &mut encoder,
- &mut decoder,
- &values[..],
- &mut result_data[..],
- )?;
- assert_eq!(actual_total, total);
- assert_eq!(result_data, values);
-
- Ok(())
- }
-
- fn test_dict_internal(total: usize, type_length: i32) -> Result<()> {
- let mut encoder = create_test_dict_encoder::<T>(type_length);
- let mut values = <T as RandGen<T>>::gen_vec(type_length, total);
- encoder.put(&values[..])?;
-
- let mut data = encoder.flush_buffer()?;
- let mut decoder = create_test_dict_decoder::<T>();
- let mut dict_decoder = PlainDecoder::<T>::new(type_length);
- dict_decoder.set_data(encoder.write_dict()?, encoder.num_entries())?;
- decoder.set_dict(Box::new(dict_decoder))?;
- let mut result_data = vec![T::T::default(); total];
- decoder.set_data(data, total)?;
- let mut actual_total = decoder.get(&mut result_data)?;
-
- assert_eq!(actual_total, total);
- assert_eq!(result_data, values);
-
- // Encode more data after flush and test with decoder
-
- values = <T as RandGen<T>>::gen_vec(type_length, total);
- encoder.put(&values[..])?;
- data = encoder.flush_buffer()?;
-
- let mut dict_decoder = PlainDecoder::<T>::new(type_length);
- dict_decoder.set_data(encoder.write_dict()?, encoder.num_entries())?;
- decoder.set_dict(Box::new(dict_decoder))?;
- decoder.set_data(data, total)?;
- actual_total = decoder.get(&mut result_data)?;
-
- assert_eq!(actual_total, total);
- assert_eq!(result_data, values);
-
- Ok(())
- }
- }
-
- fn put_and_get<T: DataType>(
- encoder: &mut Box<dyn Encoder<T>>,
- decoder: &mut Box<dyn Decoder<T>>,
- input: &[T::T],
- output: &mut [T::T],
- ) -> Result<usize> {
- encoder.put(input)?;
- let data = encoder.flush_buffer()?;
- decoder.set_data(data, input.len())?;
- decoder.get(output)
- }
-
- fn create_and_check_encoder<T: DataType>(
- encoding: Encoding,
- err: Option<ParquetError>,
- ) {
- let descr = create_test_col_desc_ptr(-1, T::get_physical_type());
- let mem_tracker = Arc::new(MemTracker::new());
- let encoder = get_encoder::<T>(descr, encoding, mem_tracker);
- match err {
- Some(parquet_error) => {
- assert!(encoder.is_err());
- assert_eq!(encoder.err().unwrap(), parquet_error);
- }
- None => {
- assert!(encoder.is_ok());
- assert_eq!(encoder.unwrap().encoding(), encoding);
- }
- }
- }
-
- // Creates test column descriptor.
- fn create_test_col_desc_ptr(type_len: i32, t: Type) -> ColumnDescPtr {
- let ty = SchemaType::primitive_type_builder("t", t)
- .with_length(type_len)
- .build()
- .unwrap();
- Arc::new(ColumnDescriptor::new(
- Arc::new(ty),
- 0,
- 0,
- ColumnPath::new(vec![]),
- ))
- }
-
- fn create_test_encoder<T: DataType>(
- type_len: i32,
- enc: Encoding,
- ) -> Box<dyn Encoder<T>> {
- let desc = create_test_col_desc_ptr(type_len, T::get_physical_type());
- let mem_tracker = Arc::new(MemTracker::new());
- get_encoder(desc, enc, mem_tracker).unwrap()
- }
-
- fn create_test_decoder<T: DataType>(
- type_len: i32,
- enc: Encoding,
- ) -> Box<dyn Decoder<T>> {
- let desc = create_test_col_desc_ptr(type_len, T::get_physical_type());
- get_decoder(desc, enc).unwrap()
- }
-
- fn create_test_dict_encoder<T: DataType>(type_len: i32) -> DictEncoder<T> {
- let desc = create_test_col_desc_ptr(type_len, T::get_physical_type());
- let mem_tracker = Arc::new(MemTracker::new());
- DictEncoder::<T>::new(desc, mem_tracker)
- }
-
- fn create_test_dict_decoder<T: DataType>() -> DictDecoder<T> {
- DictDecoder::<T>::new()
- }
-}
diff --git a/parquet/src/encodings/levels.rs b/parquet/src/encodings/levels.rs
deleted file mode 100644
index 6727589..0000000
--- a/parquet/src/encodings/levels.rs
+++ /dev/null
@@ -1,563 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-use std::{cmp, mem};
-
-use super::rle::{RleDecoder, RleEncoder};
-
-use crate::basic::Encoding;
-use crate::data_type::AsBytes;
-use crate::errors::{ParquetError, Result};
-use crate::util::{
- bit_util::{ceil, log2, BitReader, BitWriter},
- memory::ByteBufferPtr,
-};
-
-/// Computes max buffer size for level encoder/decoder based on encoding, max
-/// repetition/definition level and number of total buffered values (includes null
-/// values).
-#[inline]
-pub fn max_buffer_size(
- encoding: Encoding,
- max_level: i16,
- num_buffered_values: usize,
-) -> usize {
- let bit_width = log2(max_level as u64 + 1) as u8;
- match encoding {
- Encoding::RLE => {
- RleEncoder::max_buffer_size(bit_width, num_buffered_values)
- + RleEncoder::min_buffer_size(bit_width)
- }
- Encoding::BIT_PACKED => {
- ceil((num_buffered_values * bit_width as usize) as i64, 8) as usize
- }
- _ => panic!("Unsupported encoding type {}", encoding),
- }
-}
-
-/// Encoder for definition/repetition levels.
-/// Currently only supports RLE and BIT_PACKED (dev/null) encoding, including v2.
-pub enum LevelEncoder {
- RLE(RleEncoder),
- RLE_V2(RleEncoder),
- BIT_PACKED(u8, BitWriter),
-}
-
-impl LevelEncoder {
- /// Creates new level encoder based on encoding, max level and underlying byte buffer.
- /// For bit packed encoding it is assumed that buffer is already allocated with
- /// `levels::max_buffer_size` method.
- ///
- /// Used to encode levels for Data Page v1.
- ///
- /// Panics, if encoding is not supported.
- pub fn v1(encoding: Encoding, max_level: i16, byte_buffer: Vec<u8>) -> Self {
- let bit_width = log2(max_level as u64 + 1) as u8;
- match encoding {
- Encoding::RLE => LevelEncoder::RLE(RleEncoder::new_from_buf(
- bit_width,
- byte_buffer,
- mem::size_of::<i32>(),
- )),
- Encoding::BIT_PACKED => {
- // Here we set full byte buffer without adjusting for num_buffered_values,
- // because byte buffer will already be allocated with size from
- // `max_buffer_size()` method.
- LevelEncoder::BIT_PACKED(
- bit_width,
- BitWriter::new_from_buf(byte_buffer, 0),
- )
- }
- _ => panic!("Unsupported encoding type {}", encoding),
- }
- }
-
- /// Creates new level encoder based on RLE encoding. Used to encode Data Page v2
- /// repetition and definition levels.
- pub fn v2(max_level: i16, byte_buffer: Vec<u8>) -> Self {
- let bit_width = log2(max_level as u64 + 1) as u8;
- LevelEncoder::RLE_V2(RleEncoder::new_from_buf(bit_width, byte_buffer, 0))
- }
-
- /// Put/encode levels vector into this level encoder.
- /// Returns number of encoded values that are less than or equal to length of the
- /// input buffer.
- ///
- /// RLE and BIT_PACKED level encoders return Err() when internal buffer overflows or
- /// flush fails.
- #[inline]
- pub fn put(&mut self, buffer: &[i16]) -> Result<usize> {
- let mut num_encoded = 0;
- match *self {
- LevelEncoder::RLE(ref mut encoder)
- | LevelEncoder::RLE_V2(ref mut encoder) => {
- for value in buffer {
- if !encoder.put(*value as u64)? {
- return Err(general_err!("RLE buffer is full"));
- }
- num_encoded += 1;
- }
- encoder.flush()?;
- }
- LevelEncoder::BIT_PACKED(bit_width, ref mut encoder) => {
- for value in buffer {
- if !encoder.put_value(*value as u64, bit_width as usize) {
- return Err(general_err!("Not enough bytes left"));
- }
- num_encoded += 1;
- }
- encoder.flush();
- }
- }
- Ok(num_encoded)
- }
-
- /// Finalizes level encoder, flush all intermediate buffers and return resulting
- /// encoded buffer. Returned buffer is already truncated to encoded bytes only.
- #[inline]
- pub fn consume(self) -> Result<Vec<u8>> {
- match self {
- LevelEncoder::RLE(encoder) => {
- let mut encoded_data = encoder.consume()?;
- // Account for the buffer offset
- let encoded_len = encoded_data.len() - mem::size_of::<i32>();
- let len = (encoded_len as i32).to_le();
- let len_bytes = len.as_bytes();
- encoded_data[0..len_bytes.len()].copy_from_slice(len_bytes);
- Ok(encoded_data)
- }
- LevelEncoder::RLE_V2(encoder) => encoder.consume(),
- LevelEncoder::BIT_PACKED(_, encoder) => Ok(encoder.consume()),
- }
- }
-}
-
-/// Decoder for definition/repetition levels.
-/// Currently only supports RLE and BIT_PACKED encoding for Data Page v1 and
-/// RLE for Data Page v2.
-pub enum LevelDecoder {
- RLE(Option<usize>, RleDecoder),
- RLE_V2(Option<usize>, RleDecoder),
- BIT_PACKED(Option<usize>, u8, BitReader),
-}
-
-impl LevelDecoder {
- /// Creates new level decoder based on encoding and max definition/repetition level.
- /// This method only initializes level decoder, `set_data` method must be called
- /// before reading any value.
- ///
- /// Used to encode levels for Data Page v1.
- ///
- /// Panics if encoding is not supported
- pub fn v1(encoding: Encoding, max_level: i16) -> Self {
- let bit_width = log2(max_level as u64 + 1) as u8;
- match encoding {
- Encoding::RLE => LevelDecoder::RLE(None, RleDecoder::new(bit_width)),
- Encoding::BIT_PACKED => {
- LevelDecoder::BIT_PACKED(None, bit_width, BitReader::from(Vec::new()))
- }
- _ => panic!("Unsupported encoding type {}", encoding),
- }
- }
-
- /// Creates new level decoder based on RLE encoding.
- /// Used to decode Data Page v2 repetition and definition levels.
- ///
- /// To set data for this decoder, use `set_data_range` method.
- pub fn v2(max_level: i16) -> Self {
- let bit_width = log2(max_level as u64 + 1) as u8;
- LevelDecoder::RLE_V2(None, RleDecoder::new(bit_width))
- }
-
- /// Sets data for this level decoder, and returns total number of bytes set.
- /// This is used for Data Page v1 levels.
- ///
- /// `data` is encoded data as byte buffer, `num_buffered_values` represents total
- /// number of values that is expected.
- ///
- /// Both RLE and BIT_PACKED level decoders set `num_buffered_values` as total number
- /// of values that they can return and track num values.
- #[inline]
- pub fn set_data(&mut self, num_buffered_values: usize, data: ByteBufferPtr) -> usize {
- match *self {
- LevelDecoder::RLE(ref mut num_values, ref mut decoder) => {
- *num_values = Some(num_buffered_values);
- let i32_size = mem::size_of::<i32>();
- let data_size = read_num_bytes!(i32, i32_size, data.as_ref()) as usize;
- decoder.set_data(data.range(i32_size, data_size));
- i32_size + data_size
- }
- LevelDecoder::BIT_PACKED(ref mut num_values, bit_width, ref mut decoder) => {
- *num_values = Some(num_buffered_values);
- // Set appropriate number of bytes: if max size is larger than buffer -
- // set full buffer
- let num_bytes =
- ceil((num_buffered_values * bit_width as usize) as i64, 8);
- let data_size = cmp::min(num_bytes as usize, data.len());
- decoder.reset(data.range(data.start(), data_size));
- data_size
- }
- _ => panic!(),
- }
- }
-
- /// Sets byte array explicitly when start position `start` and length `len` are known
- /// in advance. Only supported by RLE level decoder and used for Data Page v2 levels.
- /// Returns number of total bytes set for this decoder (len).
- #[inline]
- pub fn set_data_range(
- &mut self,
- num_buffered_values: usize,
- data: &ByteBufferPtr,
- start: usize,
- len: usize,
- ) -> usize {
- match *self {
- LevelDecoder::RLE_V2(ref mut num_values, ref mut decoder) => {
- decoder.set_data(data.range(start, len));
- *num_values = Some(num_buffered_values);
- len
- }
- _ => panic!(
- "set_data_range() method is only supported by RLE v2 encoding type"
- ),
- }
- }
-
- /// Returns true if data is set for decoder, false otherwise.
- #[inline]
- pub fn is_data_set(&self) -> bool {
- match self {
- LevelDecoder::RLE(ref num_values, _) => num_values.is_some(),
- LevelDecoder::RLE_V2(ref num_values, _) => num_values.is_some(),
- LevelDecoder::BIT_PACKED(ref num_values, ..) => num_values.is_some(),
- }
- }
-
- /// Decodes values and puts them into `buffer`.
- /// Returns number of values that were successfully decoded (less than or equal to
- /// buffer length).
- #[inline]
- pub fn get(&mut self, buffer: &mut [i16]) -> Result<usize> {
- assert!(self.is_data_set(), "No data set for decoding");
- match *self {
- LevelDecoder::RLE(ref mut num_values, ref mut decoder)
- | LevelDecoder::RLE_V2(ref mut num_values, ref mut decoder) => {
- // Max length we can read
- let len = cmp::min(num_values.unwrap(), buffer.len());
- let values_read = decoder.get_batch::<i16>(&mut buffer[0..len])?;
- *num_values = num_values.map(|len| len - values_read);
- Ok(values_read)
- }
- LevelDecoder::BIT_PACKED(ref mut num_values, bit_width, ref mut decoder) => {
- // When extracting values from bit reader, it might return more values
- // than left because of padding to a full byte, we use
- // num_values to track precise number of values.
- let len = cmp::min(num_values.unwrap(), buffer.len());
- let values_read =
- decoder.get_batch::<i16>(&mut buffer[..len], bit_width as usize);
- *num_values = num_values.map(|len| len - values_read);
- Ok(values_read)
- }
- }
- }
-}
-
-#[cfg(test)]
-mod tests {
- use super::*;
-
- use crate::util::test_common::random_numbers_range;
-
- fn test_internal_roundtrip(enc: Encoding, levels: &[i16], max_level: i16, v2: bool) {
- let size = max_buffer_size(enc, max_level, levels.len());
- let mut encoder = if v2 {
- LevelEncoder::v2(max_level, vec![0; size])
- } else {
- LevelEncoder::v1(enc, max_level, vec![0; size])
- };
- encoder.put(&levels).expect("put() should be OK");
- let encoded_levels = encoder.consume().expect("consume() should be OK");
-
- let byte_buf = ByteBufferPtr::new(encoded_levels);
- let mut decoder;
- if v2 {
- decoder = LevelDecoder::v2(max_level);
- decoder.set_data_range(levels.len(), &byte_buf, 0, byte_buf.len());
- } else {
- decoder = LevelDecoder::v1(enc, max_level);
- decoder.set_data(levels.len(), byte_buf);
- };
-
- let mut buffer = vec![0; levels.len()];
- let num_decoded = decoder.get(&mut buffer).expect("get() should be OK");
- assert_eq!(num_decoded, levels.len());
- assert_eq!(buffer, levels);
- }
-
- // Performs incremental read until all bytes are read
- fn test_internal_roundtrip_incremental(
- enc: Encoding,
- levels: &[i16],
- max_level: i16,
- v2: bool,
- ) {
- let size = max_buffer_size(enc, max_level, levels.len());
- let mut encoder = if v2 {
- LevelEncoder::v2(max_level, vec![0; size])
- } else {
- LevelEncoder::v1(enc, max_level, vec![0; size])
- };
- encoder.put(&levels).expect("put() should be OK");
- let encoded_levels = encoder.consume().expect("consume() should be OK");
-
- let byte_buf = ByteBufferPtr::new(encoded_levels);
- let mut decoder;
- if v2 {
- decoder = LevelDecoder::v2(max_level);
- decoder.set_data_range(levels.len(), &byte_buf, 0, byte_buf.len());
- } else {
- decoder = LevelDecoder::v1(enc, max_level);
- decoder.set_data(levels.len(), byte_buf);
- }
-
- let mut buffer = vec![0; levels.len() * 2];
- let mut total_decoded = 0;
- let mut safe_stop = levels.len() * 2; // still terminate in case of issues in the code
- while safe_stop > 0 {
- safe_stop -= 1;
- let num_decoded = decoder
- .get(&mut buffer[total_decoded..total_decoded + 1])
- .expect("get() should be OK");
- if num_decoded == 0 {
- break;
- }
- total_decoded += num_decoded;
- }
- assert!(
- safe_stop > 0,
- "Failed to read values incrementally, reached safe stop"
- );
- assert_eq!(total_decoded, levels.len());
- assert_eq!(&buffer[0..levels.len()], levels);
- }
-
- // Tests encoding/decoding of values when output buffer is larger than number of
- // encoded values
- fn test_internal_roundtrip_underflow(
- enc: Encoding,
- levels: &[i16],
- max_level: i16,
- v2: bool,
- ) {
- let size = max_buffer_size(enc, max_level, levels.len());
- let mut encoder = if v2 {
- LevelEncoder::v2(max_level, vec![0; size])
- } else {
- LevelEncoder::v1(enc, max_level, vec![0; size])
- };
- // Encode only one value
- let num_encoded = encoder.put(&levels[0..1]).expect("put() should be OK");
- let encoded_levels = encoder.consume().expect("consume() should be OK");
- assert_eq!(num_encoded, 1);
-
- let byte_buf = ByteBufferPtr::new(encoded_levels);
- let mut decoder;
- // Set one encoded value as `num_buffered_values`
- if v2 {
- decoder = LevelDecoder::v2(max_level);
- decoder.set_data_range(1, &byte_buf, 0, byte_buf.len());
- } else {
- decoder = LevelDecoder::v1(enc, max_level);
- decoder.set_data(1, byte_buf);
- }
-
- let mut buffer = vec![0; levels.len()];
- let num_decoded = decoder.get(&mut buffer).expect("get() should be OK");
- assert_eq!(num_decoded, num_encoded);
- assert_eq!(buffer[0..num_decoded], levels[0..num_decoded]);
- }
-
- // Tests when encoded values are larger than encoder's buffer
- fn test_internal_roundtrip_overflow(
- enc: Encoding,
- levels: &[i16],
- max_level: i16,
- v2: bool,
- ) {
- let size = max_buffer_size(enc, max_level, levels.len());
- let mut encoder = if v2 {
- LevelEncoder::v2(max_level, vec![0; size])
- } else {
- LevelEncoder::v1(enc, max_level, vec![0; size])
- };
- let mut found_err = false;
- // Insert a large number of values, so we run out of space
- for _ in 0..100 {
- if let Err(err) = encoder.put(&levels) {
- assert!(format!("{}", err).contains("Not enough bytes left"));
- found_err = true;
- break;
- };
- }
- if !found_err {
- panic!("Failed test: no buffer overflow");
- }
- }
-
- #[test]
- fn test_roundtrip_one() {
- let levels = vec![0, 1, 1, 1, 1, 0, 0, 0, 0, 1];
- let max_level = 1;
- test_internal_roundtrip(Encoding::RLE, &levels, max_level, false);
- test_internal_roundtrip(Encoding::BIT_PACKED, &levels, max_level, false);
- test_internal_roundtrip(Encoding::RLE, &levels, max_level, true);
- }
-
- #[test]
- fn test_roundtrip() {
- let levels = vec![0, 1, 2, 3, 4, 5, 6, 7, 8, 9];
- let max_level = 10;
- test_internal_roundtrip(Encoding::RLE, &levels, max_level, false);
- test_internal_roundtrip(Encoding::BIT_PACKED, &levels, max_level, false);
- test_internal_roundtrip(Encoding::RLE, &levels, max_level, true);
- }
-
- #[test]
- fn test_roundtrip_incremental() {
- let levels = vec![0, 1, 2, 3, 4, 5, 6, 7, 8, 9];
- let max_level = 10;
- test_internal_roundtrip_incremental(Encoding::RLE, &levels, max_level, false);
- test_internal_roundtrip_incremental(
- Encoding::BIT_PACKED,
- &levels,
- max_level,
- false,
- );
- test_internal_roundtrip_incremental(Encoding::RLE, &levels, max_level, true);
- }
-
- #[test]
- fn test_roundtrip_all_zeros() {
- let levels = vec![0, 0, 0, 0, 0, 0, 0, 0, 0, 0];
- let max_level = 1;
- test_internal_roundtrip(Encoding::RLE, &levels, max_level, false);
- test_internal_roundtrip(Encoding::BIT_PACKED, &levels, max_level, false);
- test_internal_roundtrip(Encoding::RLE, &levels, max_level, true);
- }
-
- #[test]
- fn test_roundtrip_random() {
- // This test is mainly for bit packed level encoder/decoder
- let mut levels = Vec::new();
- let max_level = 5;
- random_numbers_range::<i16>(120, 0, max_level, &mut levels);
- test_internal_roundtrip(Encoding::RLE, &levels, max_level, false);
- test_internal_roundtrip(Encoding::BIT_PACKED, &levels, max_level, false);
- test_internal_roundtrip(Encoding::RLE, &levels, max_level, true);
- }
-
- #[test]
- fn test_roundtrip_underflow() {
- let levels = vec![1, 1, 2, 3, 2, 1, 1, 2, 3, 1];
- let max_level = 3;
- test_internal_roundtrip_underflow(Encoding::RLE, &levels, max_level, false);
- test_internal_roundtrip_underflow(
- Encoding::BIT_PACKED,
- &levels,
- max_level,
- false,
- );
- test_internal_roundtrip_underflow(Encoding::RLE, &levels, max_level, true);
- }
-
- #[test]
- fn test_roundtrip_overflow() {
- let levels = vec![1, 1, 2, 3, 2, 1, 1, 2, 3, 1];
- let max_level = 3;
- test_internal_roundtrip_overflow(Encoding::RLE, &levels, max_level, false);
- test_internal_roundtrip_overflow(Encoding::BIT_PACKED, &levels, max_level, false);
- test_internal_roundtrip_overflow(Encoding::RLE, &levels, max_level, true);
- }
-
- #[test]
- fn test_rle_decoder_set_data_range() {
- // Buffer containing both repetition and definition levels
- let buffer = ByteBufferPtr::new(vec![5, 198, 2, 5, 42, 168, 10, 0, 2, 3, 36, 73]);
-
- let max_rep_level = 1;
- let mut decoder = LevelDecoder::v2(max_rep_level);
- assert_eq!(decoder.set_data_range(10, &buffer, 0, 3), 3);
- let mut result = vec![0; 10];
- let num_decoded = decoder.get(&mut result).expect("get() should be OK");
- assert_eq!(num_decoded, 10);
- assert_eq!(result, vec![0, 1, 1, 0, 0, 0, 1, 1, 0, 1]);
-
- let max_def_level = 2;
- let mut decoder = LevelDecoder::v2(max_def_level);
- assert_eq!(decoder.set_data_range(10, &buffer, 3, 5), 5);
- let mut result = vec![0; 10];
- let num_decoded = decoder.get(&mut result).expect("get() should be OK");
- assert_eq!(num_decoded, 10);
- assert_eq!(result, vec![2, 2, 2, 0, 0, 2, 2, 2, 2, 2]);
- }
-
- #[test]
- #[should_panic(
- expected = "set_data_range() method is only supported by RLE v2 encoding type"
- )]
- fn test_bit_packed_decoder_set_data_range() {
- // Buffer containing both repetition and definition levels
- let buffer = ByteBufferPtr::new(vec![1, 2, 3, 4, 5]);
- let max_level = 1;
- let mut decoder = LevelDecoder::v1(Encoding::BIT_PACKED, max_level);
- decoder.set_data_range(10, &buffer, 0, 3);
- }
-
- #[test]
- fn test_bit_packed_decoder_set_data() {
- // Test the maximum size that is assigned based on number of values and buffer
- // length
- let buffer = ByteBufferPtr::new(vec![1, 2, 3, 4, 5]);
- let max_level = 1;
- let mut decoder = LevelDecoder::v1(Encoding::BIT_PACKED, max_level);
- // This should reset to entire buffer
- assert_eq!(decoder.set_data(1024, buffer.all()), buffer.len());
- // This should set smallest num bytes
- assert_eq!(decoder.set_data(3, buffer.all()), 1);
- }
-
- #[test]
- #[should_panic(expected = "No data set for decoding")]
- fn test_rle_level_decoder_get_no_set_data() {
- // `get()` normally panics because bit_reader is not set for RLE decoding
- // we have explicit check now in set_data
- let max_rep_level = 2;
- let mut decoder = LevelDecoder::v1(Encoding::RLE, max_rep_level);
- let mut buffer = vec![0; 16];
- decoder.get(&mut buffer).unwrap();
- }
-
- #[test]
- #[should_panic(expected = "No data set for decoding")]
- fn test_bit_packed_level_decoder_get_no_set_data() {
- let max_rep_level = 2;
- let mut decoder = LevelDecoder::v1(Encoding::BIT_PACKED, max_rep_level);
- let mut buffer = vec![0; 16];
- decoder.get(&mut buffer).unwrap();
- }
-}
diff --git a/parquet/src/encodings/mod.rs b/parquet/src/encodings/mod.rs
deleted file mode 100644
index 6046dda..0000000
--- a/parquet/src/encodings/mod.rs
+++ /dev/null
@@ -1,21 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-pub mod decoding;
-pub mod encoding;
-pub mod levels;
-pub(crate) mod rle;
diff --git a/parquet/src/encodings/rle.rs b/parquet/src/encodings/rle.rs
deleted file mode 100644
index b2a23da..0000000
--- a/parquet/src/encodings/rle.rs
+++ /dev/null
@@ -1,831 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-use std::{cmp, mem::size_of};
-
-use crate::errors::{ParquetError, Result};
-use crate::util::{
- bit_util::{self, from_ne_slice, BitReader, BitWriter, FromBytes},
- memory::ByteBufferPtr,
-};
-
-/// Rle/Bit-Packing Hybrid Encoding
-/// The grammar for this encoding looks like the following (copied verbatim
-/// from <https://github.com/Parquet/parquet-format/blob/master/Encodings.md>):
-///
-/// rle-bit-packed-hybrid: <length> <encoded-data>
-/// length := length of the <encoded-data> in bytes stored as 4 bytes little endian
-/// encoded-data := <run>*
-/// run := <bit-packed-run> | <rle-run>
-/// bit-packed-run := <bit-packed-header> <bit-packed-values>
-/// bit-packed-header := varint-encode(<bit-pack-count> << 1 | 1)
-/// we always bit-pack a multiple of 8 values at a time, so we only store the number of
-/// values / 8
-/// bit-pack-count := (number of values in this run) / 8
-/// bit-packed-values := *see 1 below*
-/// rle-run := <rle-header> <repeated-value>
-/// rle-header := varint-encode( (number of times repeated) << 1)
-/// repeated-value := value that is repeated, using a fixed-width of
-/// round-up-to-next-byte(bit-width)
-
-/// Maximum groups per bit-packed run. Current value is 64.
-const MAX_GROUPS_PER_BIT_PACKED_RUN: usize = 1 << 6;
-const MAX_VALUES_PER_BIT_PACKED_RUN: usize = MAX_GROUPS_PER_BIT_PACKED_RUN * 8;
-const MAX_WRITER_BUF_SIZE: usize = 1 << 10;
-
-/// A RLE/Bit-Packing hybrid encoder.
-// TODO: tracking memory usage
-pub struct RleEncoder {
- // Number of bits needed to encode the value. Must be in the range of [0, 64].
- bit_width: u8,
-
- // Underlying writer which holds an internal buffer.
- bit_writer: BitWriter,
-
- // The maximum byte size a single run can take.
- max_run_byte_size: usize,
-
- // Buffered values for bit-packed runs.
- buffered_values: [u64; 8],
-
- // Number of current buffered values. Must be less than 8.
- num_buffered_values: usize,
-
- // The current (also last) value that was written and the count of how many
- // times in a row that value has been seen.
- current_value: u64,
-
- // The number of repetitions for `current_value`. If this gets too high we'd
- // switch to use RLE encoding.
- repeat_count: usize,
-
- // Number of bit-packed values in the current run. This doesn't include values
- // in `buffered_values`.
- bit_packed_count: usize,
-
- // The position of the indicator byte in the `bit_writer`.
- indicator_byte_pos: i64,
-}
-
-impl RleEncoder {
- pub fn new(bit_width: u8, buffer_len: usize) -> Self {
- let buffer = vec![0; buffer_len];
- RleEncoder::new_from_buf(bit_width, buffer, 0)
- }
-
- /// Initialize the encoder from existing `buffer` and the starting offset `start`.
- pub fn new_from_buf(bit_width: u8, buffer: Vec<u8>, start: usize) -> Self {
- assert!(bit_width <= 64, "bit_width ({}) out of range.", bit_width);
- let max_run_byte_size = RleEncoder::min_buffer_size(bit_width);
- assert!(
- buffer.len() >= max_run_byte_size,
- "buffer length {} must be greater than {}",
- buffer.len(),
- max_run_byte_size
- );
- let bit_writer = BitWriter::new_from_buf(buffer, start);
- RleEncoder {
- bit_width,
- bit_writer,
- max_run_byte_size,
- buffered_values: [0; 8],
- num_buffered_values: 0,
- current_value: 0,
- repeat_count: 0,
- bit_packed_count: 0,
- indicator_byte_pos: -1,
- }
- }
-
- /// Returns the minimum buffer size needed to use the encoder for `bit_width`.
- /// This is the maximum length of a single run for `bit_width`.
- pub fn min_buffer_size(bit_width: u8) -> usize {
- let max_bit_packed_run_size = 1 + bit_util::ceil(
- (MAX_VALUES_PER_BIT_PACKED_RUN * bit_width as usize) as i64,
- 8,
- );
- let max_rle_run_size =
- bit_util::MAX_VLQ_BYTE_LEN + bit_util::ceil(bit_width as i64, 8) as usize;
- std::cmp::max(max_bit_packed_run_size as usize, max_rle_run_size)
- }
-
- /// Returns the maximum buffer size takes to encode `num_values` values with
- /// `bit_width`.
- pub fn max_buffer_size(bit_width: u8, num_values: usize) -> usize {
- // First the maximum size for bit-packed run
- let bytes_per_run = bit_width;
- let num_runs = bit_util::ceil(num_values as i64, 8) as usize;
- let bit_packed_max_size = num_runs + num_runs * bytes_per_run as usize;
-
- // Second the maximum size for RLE run
- let min_rle_run_size = 1 + bit_util::ceil(bit_width as i64, 8) as usize;
- let rle_max_size =
- bit_util::ceil(num_values as i64, 8) as usize * min_rle_run_size;
- std::cmp::max(bit_packed_max_size, rle_max_size) as usize
- }
-
- /// Encodes `value`, which must be representable with `bit_width` bits.
- /// Returns true if the value fits in buffer, false if it doesn't, or
- /// error if something is wrong.
- #[inline]
- pub fn put(&mut self, value: u64) -> Result<bool> {
- // This function buffers 8 values at a time. After seeing 8 values, it
- // decides whether the current run should be encoded in bit-packed or RLE.
- if self.current_value == value {
- self.repeat_count += 1;
- if self.repeat_count > 8 {
- // A continuation of last value. No need to buffer.
- return Ok(true);
- }
- } else {
- if self.repeat_count >= 8 {
- // The current RLE run has ended and we've gathered enough. Flush first.
- assert_eq!(self.bit_packed_count, 0);
- self.flush_rle_run()?;
- }
- self.repeat_count = 1;
- self.current_value = value;
- }
-
- self.buffered_values[self.num_buffered_values] = value;
- self.num_buffered_values += 1;
- if self.num_buffered_values == 8 {
- // Buffered values are full. Flush them.
- assert_eq!(self.bit_packed_count % 8, 0);
- self.flush_buffered_values()?;
- }
-
- Ok(true)
- }
-
- #[inline]
- pub fn buffer(&self) -> &[u8] {
- self.bit_writer.buffer()
- }
-
- #[inline]
- pub fn len(&self) -> usize {
- self.bit_writer.bytes_written()
- }
-
- #[inline]
- pub fn consume(mut self) -> Result<Vec<u8>> {
- self.flush()?;
- Ok(self.bit_writer.consume())
- }
-
- /// Borrow equivalent of the `consume` method.
- /// Call `clear()` after invoking this method.
- #[inline]
- pub fn flush_buffer(&mut self) -> Result<&[u8]> {
- self.flush()?;
- Ok(self.bit_writer.flush_buffer())
- }
-
- /// Clears the internal state so this encoder can be reused (e.g., after becoming
- /// full).
- #[inline]
- pub fn clear(&mut self) {
- self.bit_writer.clear();
- self.num_buffered_values = 0;
- self.current_value = 0;
- self.repeat_count = 0;
- self.bit_packed_count = 0;
- self.indicator_byte_pos = -1;
- }
-
- /// Flushes all remaining values and return the final byte buffer maintained by the
- /// internal writer.
- #[inline]
- pub fn flush(&mut self) -> Result<()> {
- if self.bit_packed_count > 0
- || self.repeat_count > 0
- || self.num_buffered_values > 0
- {
- let all_repeat = self.bit_packed_count == 0
- && (self.repeat_count == self.num_buffered_values
- || self.num_buffered_values == 0);
- if self.repeat_count > 0 && all_repeat {
- self.flush_rle_run()?;
- } else {
- // Buffer the last group of bit-packed values to 8 by padding with 0s.
- if self.num_buffered_values > 0 {
- while self.num_buffered_values < 8 {
- self.buffered_values[self.num_buffered_values] = 0;
- self.num_buffered_values += 1;
- }
- }
- self.bit_packed_count += self.num_buffered_values;
- self.flush_bit_packed_run(true)?;
- self.repeat_count = 0;
- }
- }
- Ok(())
- }
-
- fn flush_rle_run(&mut self) -> Result<()> {
- assert!(self.repeat_count > 0);
- let indicator_value = self.repeat_count << 1;
- let mut result = self.bit_writer.put_vlq_int(indicator_value as u64);
- result &= self.bit_writer.put_aligned(
- self.current_value,
- bit_util::ceil(self.bit_width as i64, 8) as usize,
- );
- if !result {
- return Err(general_err!("Failed to write RLE run"));
- }
- self.num_buffered_values = 0;
- self.repeat_count = 0;
- Ok(())
- }
-
- fn flush_bit_packed_run(&mut self, update_indicator_byte: bool) -> Result<()> {
- if self.indicator_byte_pos < 0 {
- self.indicator_byte_pos = self.bit_writer.skip(1)? as i64;
- }
-
- // Write all buffered values as bit-packed literals
- for i in 0..self.num_buffered_values {
- let _ = self
- .bit_writer
- .put_value(self.buffered_values[i], self.bit_width as usize);
- }
- self.num_buffered_values = 0;
- if update_indicator_byte {
- // Write the indicator byte to the reserved position in `bit_writer`
- let num_groups = self.bit_packed_count / 8;
- let indicator_byte = ((num_groups << 1) | 1) as u8;
- if !self.bit_writer.put_aligned_offset(
- indicator_byte,
- 1,
- self.indicator_byte_pos as usize,
- ) {
- return Err(general_err!("Not enough space to write indicator byte"));
- }
- self.indicator_byte_pos = -1;
- self.bit_packed_count = 0;
- }
- Ok(())
- }
-
- #[inline(never)]
- fn flush_buffered_values(&mut self) -> Result<()> {
- if self.repeat_count >= 8 {
- self.num_buffered_values = 0;
- if self.bit_packed_count > 0 {
- // In this case we choose RLE encoding. Flush the current buffered values
- // as bit-packed encoding.
- assert_eq!(self.bit_packed_count % 8, 0);
- self.flush_bit_packed_run(true)?
- }
- return Ok(());
- }
-
- self.bit_packed_count += self.num_buffered_values;
- let num_groups = self.bit_packed_count / 8;
- if num_groups + 1 >= MAX_GROUPS_PER_BIT_PACKED_RUN {
- // We've reached the maximum value that can be hold in a single bit-packed
- // run.
- assert!(self.indicator_byte_pos >= 0);
- self.flush_bit_packed_run(true)?;
- } else {
- self.flush_bit_packed_run(false)?;
- }
- self.repeat_count = 0;
- Ok(())
- }
-}
-
-/// A RLE/Bit-Packing hybrid decoder.
-pub struct RleDecoder {
- // Number of bits used to encode the value. Must be between [0, 64].
- bit_width: u8,
-
- // Bit reader loaded with input buffer.
- bit_reader: Option<BitReader>,
-
- // Buffer used when `bit_reader` is not `None`, for batch reading.
- index_buf: [i32; 1024],
-
- // The remaining number of values in RLE for this run
- rle_left: u32,
-
- // The remaining number of values in Bit-Packing for this run
- bit_packed_left: u32,
-
- // The current value for the case of RLE mode
- current_value: Option<u64>,
-}
-
-impl RleDecoder {
- pub fn new(bit_width: u8) -> Self {
- RleDecoder {
- bit_width,
- rle_left: 0,
- bit_packed_left: 0,
- bit_reader: None,
- index_buf: [0; 1024],
- current_value: None,
- }
- }
-
- #[inline]
- pub fn set_data(&mut self, data: ByteBufferPtr) {
- if let Some(ref mut bit_reader) = self.bit_reader {
- bit_reader.reset(data);
- } else {
- self.bit_reader = Some(BitReader::new(data));
- }
-
- let _ = self.reload();
- }
-
- // These functions inline badly, they tend to inline and then create very large loop unrolls
- // that damage L1d-cache occupancy. This results in a ~18% performance drop
- #[inline(never)]
- pub fn get<T: FromBytes>(&mut self) -> Result<Option<T>> {
- assert!(size_of::<T>() <= 8);
-
- while self.rle_left == 0 && self.bit_packed_left == 0 {
- if !self.reload() {
- return Ok(None);
- }
- }
-
- let value = if self.rle_left > 0 {
- let rle_value = from_ne_slice(
- &self
- .current_value
- .as_mut()
- .expect("current_value should be Some")
- .to_ne_bytes(),
- );
- self.rle_left -= 1;
- rle_value
- } else {
- // self.bit_packed_left > 0
- let bit_reader = self.bit_reader.as_mut().expect("bit_reader should be Some");
- let bit_packed_value = bit_reader
- .get_value(self.bit_width as usize)
- .ok_or_else(|| eof_err!("Not enough data for 'bit_packed_value'"))?;
- self.bit_packed_left -= 1;
- bit_packed_value
- };
-
- Ok(Some(value))
- }
-
- #[inline(never)]
- pub fn get_batch<T: FromBytes>(&mut self, buffer: &mut [T]) -> Result<usize> {
- assert!(size_of::<T>() <= 8);
-
- let mut values_read = 0;
- while values_read < buffer.len() {
- if self.rle_left > 0 {
- let num_values =
- cmp::min(buffer.len() - values_read, self.rle_left as usize);
- for i in 0..num_values {
- let repeated_value = from_ne_slice(
- &self.current_value.as_mut().unwrap().to_ne_bytes(),
- );
- buffer[values_read + i] = repeated_value;
- }
- self.rle_left -= num_values as u32;
- values_read += num_values;
- } else if self.bit_packed_left > 0 {
- let mut num_values =
- cmp::min(buffer.len() - values_read, self.bit_packed_left as usize);
- let bit_reader =
- self.bit_reader.as_mut().expect("bit_reader should be set");
-
- num_values = bit_reader.get_batch::<T>(
- &mut buffer[values_read..values_read + num_values],
- self.bit_width as usize,
- );
- self.bit_packed_left -= num_values as u32;
- values_read += num_values;
- } else if !self.reload() {
- break;
- }
- }
-
- Ok(values_read)
- }
-
- #[inline(never)]
- pub fn get_batch_with_dict<T>(
- &mut self,
- dict: &[T],
- buffer: &mut [T],
- max_values: usize,
- ) -> Result<usize>
- where
- T: Default + Clone,
- {
- assert!(buffer.len() >= max_values);
-
- let mut values_read = 0;
- while values_read < max_values {
- if self.rle_left > 0 {
- let num_values =
- cmp::min(max_values - values_read, self.rle_left as usize);
- let dict_idx = self.current_value.unwrap() as usize;
- for i in 0..num_values {
- buffer[values_read + i].clone_from(&dict[dict_idx]);
- }
- self.rle_left -= num_values as u32;
- values_read += num_values;
- } else if self.bit_packed_left > 0 {
- let bit_reader =
- self.bit_reader.as_mut().expect("bit_reader should be set");
-
- let mut num_values =
- cmp::min(max_values - values_read, self.bit_packed_left as usize);
-
- num_values = cmp::min(num_values, self.index_buf.len());
- loop {
- num_values = bit_reader.get_batch::<i32>(
- &mut self.index_buf[..num_values],
- self.bit_width as usize,
- );
- for i in 0..num_values {
- buffer[values_read + i]
- .clone_from(&dict[self.index_buf[i] as usize])
- }
- self.bit_packed_left -= num_values as u32;
- values_read += num_values;
- if num_values < self.index_buf.len() {
- break;
- }
- }
- } else if !self.reload() {
- break;
- }
- }
-
- Ok(values_read)
- }
-
- #[inline]
- fn reload(&mut self) -> bool {
- let bit_reader = self.bit_reader.as_mut().expect("bit_reader should be set");
-
- if let Some(indicator_value) = bit_reader.get_vlq_int() {
- if indicator_value & 1 == 1 {
- self.bit_packed_left = ((indicator_value >> 1) * 8) as u32;
- } else {
- self.rle_left = (indicator_value >> 1) as u32;
- let value_width = bit_util::ceil(self.bit_width as i64, 8);
- self.current_value = bit_reader.get_aligned::<u64>(value_width as usize);
- assert!(self.current_value.is_some());
- }
- true
- } else {
- false
- }
- }
-}
-
-#[cfg(test)]
-mod tests {
- use super::*;
-
- use rand::{self, distributions::Standard, thread_rng, Rng, SeedableRng};
-
- use crate::util::memory::ByteBufferPtr;
-
- const MAX_WIDTH: usize = 32;
-
- #[test]
- fn test_rle_decode_int32() {
- // Test data: 0-7 with bit width 3
- // 00000011 10001000 11000110 11111010
- let data = ByteBufferPtr::new(vec![0x03, 0x88, 0xC6, 0xFA]);
- let mut decoder: RleDecoder = RleDecoder::new(3);
- decoder.set_data(data);
- let mut buffer = vec![0; 8];
- let expected = vec![0, 1, 2, 3, 4, 5, 6, 7];
- let result = decoder.get_batch::<i32>(&mut buffer);
- assert!(result.is_ok());
- assert_eq!(buffer, expected);
- }
-
- #[test]
- fn test_rle_consume_flush_buffer() {
- let data = vec![1, 1, 1, 2, 2, 3, 3, 3];
- let mut encoder1 = RleEncoder::new(3, 256);
- let mut encoder2 = RleEncoder::new(3, 256);
- for value in data {
- encoder1.put(value as u64).unwrap();
- encoder2.put(value as u64).unwrap();
- }
- let res1 = encoder1.flush_buffer().unwrap();
- let res2 = encoder2.consume().unwrap();
- assert_eq!(res1, &res2[..]);
- }
-
- #[test]
- fn test_rle_decode_bool() {
- // RLE test data: 50 1s followed by 50 0s
- // 01100100 00000001 01100100 00000000
- let data1 = ByteBufferPtr::new(vec![0x64, 0x01, 0x64, 0x00]);
-
- // Bit-packing test data: alternating 1s and 0s, 100 total
- // 100 / 8 = 13 groups
- // 00011011 10101010 ... 00001010
- let data2 = ByteBufferPtr::new(vec![
- 0x1B, 0xAA, 0xAA, 0xAA, 0xAA, 0xAA, 0xAA, 0xAA, 0xAA, 0xAA, 0xAA, 0xAA, 0xAA,
- 0x0A,
- ]);
-
- let mut decoder: RleDecoder = RleDecoder::new(1);
- decoder.set_data(data1);
- let mut buffer = vec![false; 100];
- let mut expected = vec![];
- for i in 0..100 {
- if i < 50 {
- expected.push(true);
- } else {
- expected.push(false);
- }
- }
- let result = decoder.get_batch::<bool>(&mut buffer);
- assert!(result.is_ok());
- assert_eq!(buffer, expected);
-
- decoder.set_data(data2);
- let mut buffer = vec![false; 100];
- let mut expected = vec![];
- for i in 0..100 {
- if i % 2 == 0 {
- expected.push(false);
- } else {
- expected.push(true);
- }
- }
- let result = decoder.get_batch::<bool>(&mut buffer);
- assert!(result.is_ok());
- assert_eq!(buffer, expected);
- }
-
- #[test]
- fn test_rle_decode_with_dict_int32() {
- // Test RLE encoding: 3 0s followed by 4 1s followed by 5 2s
- // 00000110 00000000 00001000 00000001 00001010 00000010
- let dict = vec![10, 20, 30];
- let data = ByteBufferPtr::new(vec![0x06, 0x00, 0x08, 0x01, 0x0A, 0x02]);
- let mut decoder: RleDecoder = RleDecoder::new(3);
- decoder.set_data(data);
- let mut buffer = vec![0; 12];
- let expected = vec![10, 10, 10, 20, 20, 20, 20, 30, 30, 30, 30, 30];
- let result = decoder.get_batch_with_dict::<i32>(&dict, &mut buffer, 12);
- assert!(result.is_ok());
- assert_eq!(buffer, expected);
-
- // Test bit-pack encoding: 345345345455 (2 groups: 8 and 4)
- // 011 100 101 011 100 101 011 100 101 100 101 101
- // 00000011 01100011 11000111 10001110 00000011 01100101 00001011
- let dict = vec!["aaa", "bbb", "ccc", "ddd", "eee", "fff"];
- let data = ByteBufferPtr::new(vec![0x03, 0x63, 0xC7, 0x8E, 0x03, 0x65, 0x0B]);
- let mut decoder: RleDecoder = RleDecoder::new(3);
- decoder.set_data(data);
- let mut buffer = vec![""; 12];
- let expected = vec![
- "ddd", "eee", "fff", "ddd", "eee", "fff", "ddd", "eee", "fff", "eee", "fff",
- "fff",
- ];
- let result = decoder.get_batch_with_dict::<&str>(
- dict.as_slice(),
- buffer.as_mut_slice(),
- 12,
- );
- assert!(result.is_ok());
- assert_eq!(buffer, expected);
- }
-
- fn validate_rle(
- values: &[i64],
- bit_width: u8,
- expected_encoding: Option<&[u8]>,
- expected_len: i32,
- ) {
- let buffer_len = 64 * 1024;
- let mut encoder = RleEncoder::new(bit_width, buffer_len);
- for v in values {
- let result = encoder.put(*v as u64);
- assert!(result.is_ok());
- }
- let buffer = ByteBufferPtr::new(encoder.consume().expect("Expect consume() OK"));
- if expected_len != -1 {
- assert_eq!(buffer.len(), expected_len as usize);
- }
- if let Some(b) = expected_encoding {
- assert_eq!(buffer.as_ref(), b);
- }
-
- // Verify read
- let mut decoder = RleDecoder::new(bit_width);
- decoder.set_data(buffer.all());
- for v in values {
- let val: i64 = decoder
- .get()
- .expect("get() should be OK")
- .expect("get() should return more value");
- assert_eq!(val, *v);
- }
-
- // Verify batch read
- decoder.set_data(buffer);
- let mut values_read: Vec<i64> = vec![0; values.len()];
- decoder
- .get_batch(&mut values_read[..])
- .expect("get_batch() should be OK");
- assert_eq!(&values_read[..], values);
- }
-
- #[test]
- fn test_rle_specific_sequences() {
- let mut expected_buffer = Vec::new();
- let mut values = Vec::new();
- for _ in 0..50 {
- values.push(0);
- }
- for _ in 0..50 {
- values.push(1);
- }
- expected_buffer.push(50 << 1);
- expected_buffer.push(0);
- expected_buffer.push(50 << 1);
- expected_buffer.push(1);
-
- for width in 1..9 {
- validate_rle(&values[..], width, Some(&expected_buffer[..]), 4);
- }
- for width in 9..MAX_WIDTH + 1 {
- validate_rle(
- &values[..],
- width as u8,
- None,
- 2 * (1 + bit_util::ceil(width as i64, 8) as i32),
- );
- }
-
- // Test 100 0's and 1's alternating
- values.clear();
- expected_buffer.clear();
- for i in 0..101 {
- values.push(i % 2);
- }
- let num_groups = bit_util::ceil(100, 8) as u8;
- expected_buffer.push(((num_groups << 1) as u8) | 1);
- for _ in 1..(100 / 8) + 1 {
- expected_buffer.push(0b10101010);
- }
- // For the last 4 0 and 1's, padded with 0.
- expected_buffer.push(0b00001010);
- validate_rle(
- &values,
- 1,
- Some(&expected_buffer[..]),
- 1 + num_groups as i32,
- );
- for width in 2..MAX_WIDTH + 1 {
- let num_values = bit_util::ceil(100, 8) * 8;
- validate_rle(
- &values,
- width as u8,
- None,
- 1 + bit_util::ceil(width as i64 * num_values, 8) as i32,
- );
- }
- }
-
- // `validate_rle` on `num_vals` with width `bit_width`. If `value` is -1, that value
- // is used, otherwise alternating values are used.
- fn test_rle_values(bit_width: usize, num_vals: usize, value: i32) {
- let mod_val = if bit_width == 64 {
- 1
- } else {
- 1u64 << bit_width
- };
- let mut values: Vec<i64> = vec![];
- for v in 0..num_vals {
- let val = if value == -1 {
- v as i64 % mod_val as i64
- } else {
- value as i64
- };
- values.push(val);
- }
- validate_rle(&values, bit_width as u8, None, -1);
- }
-
- #[test]
- fn test_values() {
- for width in 1..MAX_WIDTH + 1 {
- test_rle_values(width, 1, -1);
- test_rle_values(width, 1024, -1);
- test_rle_values(width, 1024, 0);
- test_rle_values(width, 1024, 1);
- }
- }
-
- #[test]
- fn test_rle_specific_roundtrip() {
- let bit_width = 1;
- let buffer_len = RleEncoder::min_buffer_size(bit_width);
- let values: Vec<i16> = vec![0, 1, 1, 1, 1, 0, 0, 0, 0, 1];
- let mut encoder = RleEncoder::new(bit_width, buffer_len);
- for v in &values {
- assert!(encoder.put(*v as u64).expect("put() should be OK"));
- }
- let buffer = encoder.consume().expect("consume() should be OK");
- let mut decoder = RleDecoder::new(bit_width);
- decoder.set_data(ByteBufferPtr::new(buffer));
- let mut actual_values: Vec<i16> = vec![0; values.len()];
- decoder
- .get_batch(&mut actual_values)
- .expect("get_batch() should be OK");
- assert_eq!(actual_values, values);
- }
-
- fn test_round_trip(values: &[i32], bit_width: u8) {
- let buffer_len = 64 * 1024;
- let mut encoder = RleEncoder::new(bit_width, buffer_len);
- for v in values {
- let result = encoder.put(*v as u64).expect("put() should be OK");
- assert!(result, "put() should not return false");
- }
-
- let buffer =
- ByteBufferPtr::new(encoder.consume().expect("consume() should be OK"));
-
- // Verify read
- let mut decoder = RleDecoder::new(bit_width);
- decoder.set_data(buffer.all());
- for v in values {
- let val = decoder
- .get::<i32>()
- .expect("get() should be OK")
- .expect("get() should return value");
- assert_eq!(val, *v);
- }
-
- // Verify batch read
- let mut decoder = RleDecoder::new(bit_width);
- decoder.set_data(buffer);
- let mut values_read: Vec<i32> = vec![0; values.len()];
- decoder
- .get_batch(&mut values_read[..])
- .expect("get_batch() should be OK");
- assert_eq!(&values_read[..], values);
- }
-
- #[test]
- fn test_random() {
- let seed_len = 32;
- let niters = 50;
- let ngroups = 1000;
- let max_group_size = 15;
- let mut values = vec![];
-
- for _ in 0..niters {
- values.clear();
- let rng = thread_rng();
- let seed_vec: Vec<u8> =
- rng.sample_iter::<u8, _>(&Standard).take(seed_len).collect();
- let mut seed = [0u8; 32];
- seed.copy_from_slice(&seed_vec[0..seed_len]);
- let mut gen = rand::rngs::StdRng::from_seed(seed);
-
- let mut parity = false;
- for _ in 0..ngroups {
- let mut group_size = gen.gen_range(1..20);
- if group_size > max_group_size {
- group_size = 1;
- }
- for _ in 0..group_size {
- values.push(parity as i32);
- }
- parity = !parity;
- }
- let bit_width = bit_util::num_required_bits(values.len() as u64);
- assert!(bit_width < 64);
- test_round_trip(&values[..], bit_width as u8);
- }
- }
-}
diff --git a/parquet/src/errors.rs b/parquet/src/errors.rs
deleted file mode 100644
index be1a221..0000000
--- a/parquet/src/errors.rs
+++ /dev/null
@@ -1,146 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-//! Common Parquet errors and macros.
-
-use std::{cell, convert, io, result, str};
-
-#[cfg(any(feature = "arrow", test))]
-use arrow::error::ArrowError;
-
-#[derive(Debug, PartialEq, Clone)]
-pub enum ParquetError {
- /// General Parquet error.
- /// Returned when code violates normal workflow of working with Parquet files.
- General(String),
- /// "Not yet implemented" Parquet error.
- /// Returned when functionality is not yet available.
- NYI(String),
- /// "End of file" Parquet error.
- /// Returned when IO related failures occur, e.g. when there are not enough bytes to
- /// decode.
- EOF(String),
- #[cfg(any(feature = "arrow", test))]
- /// Arrow error.
- /// Returned when reading into arrow or writing from arrow.
- ArrowError(String),
- IndexOutOfBound(usize, usize),
-}
-
-impl std::fmt::Display for ParquetError {
- fn fmt(&self, fmt: &mut std::fmt::Formatter) -> std::fmt::Result {
- match *self {
- ParquetError::General(ref message) => {
- write!(fmt, "Parquet error: {}", message)
- }
- ParquetError::NYI(ref message) => write!(fmt, "NYI: {}", message),
- ParquetError::EOF(ref message) => write!(fmt, "EOF: {}", message),
- #[cfg(any(feature = "arrow", test))]
- ParquetError::ArrowError(ref message) => write!(fmt, "Arrow: {}", message),
- ParquetError::IndexOutOfBound(ref index, ref bound) => {
- write!(fmt, "Index {} out of bound: {}", index, bound)
- }
- }
- }
-}
-
-impl std::error::Error for ParquetError {
- fn cause(&self) -> Option<&dyn ::std::error::Error> {
- None
- }
-}
-
-impl From<io::Error> for ParquetError {
- fn from(e: io::Error) -> ParquetError {
- ParquetError::General(format!("underlying IO error: {}", e))
- }
-}
-
-#[cfg(any(feature = "snap", test))]
-impl From<snap::Error> for ParquetError {
- fn from(e: snap::Error) -> ParquetError {
- ParquetError::General(format!("underlying snap error: {}", e))
- }
-}
-
-impl From<thrift::Error> for ParquetError {
- fn from(e: thrift::Error) -> ParquetError {
- ParquetError::General(format!("underlying Thrift error: {}", e))
- }
-}
-
-impl From<cell::BorrowMutError> for ParquetError {
- fn from(e: cell::BorrowMutError) -> ParquetError {
- ParquetError::General(format!("underlying borrow error: {}", e))
- }
-}
-
-impl From<str::Utf8Error> for ParquetError {
- fn from(e: str::Utf8Error) -> ParquetError {
- ParquetError::General(format!("underlying utf8 error: {}", e))
- }
-}
-
-#[cfg(any(feature = "arrow", test))]
-impl From<ArrowError> for ParquetError {
- fn from(e: ArrowError) -> ParquetError {
- ParquetError::ArrowError(format!("underlying Arrow error: {}", e))
- }
-}
-
-/// A specialized `Result` for Parquet errors.
-pub type Result<T> = result::Result<T, ParquetError>;
-
-// ----------------------------------------------------------------------
-// Conversion from `ParquetError` to other types of `Error`s
-
-impl convert::From<ParquetError> for io::Error {
- fn from(e: ParquetError) -> Self {
- io::Error::new(io::ErrorKind::Other, e)
- }
-}
-
-// ----------------------------------------------------------------------
-// Convenient macros for different errors
-
-macro_rules! general_err {
- ($fmt:expr) => (ParquetError::General($fmt.to_owned()));
- ($fmt:expr, $($args:expr),*) => (ParquetError::General(format!($fmt, $($args),*)));
- ($e:expr, $fmt:expr) => (ParquetError::General($fmt.to_owned(), $e));
- ($e:ident, $fmt:expr, $($args:tt),*) => (
- ParquetError::General(&format!($fmt, $($args),*), $e));
-}
-
-macro_rules! nyi_err {
- ($fmt:expr) => (ParquetError::NYI($fmt.to_owned()));
- ($fmt:expr, $($args:expr),*) => (ParquetError::NYI(format!($fmt, $($args),*)));
-}
-
-macro_rules! eof_err {
- ($fmt:expr) => (ParquetError::EOF($fmt.to_owned()));
- ($fmt:expr, $($args:expr),*) => (ParquetError::EOF(format!($fmt, $($args),*)));
-}
-
-// ----------------------------------------------------------------------
-// Convert parquet error into other errors
-
-#[cfg(any(feature = "arrow", test))]
-impl Into<ArrowError> for ParquetError {
- fn into(self) -> ArrowError {
- ArrowError::ParquetError(format!("{}", self))
- }
-}
diff --git a/parquet/src/file/footer.rs b/parquet/src/file/footer.rs
deleted file mode 100644
index 2e57294..0000000
--- a/parquet/src/file/footer.rs
+++ /dev/null
@@ -1,263 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-use std::{
- cmp::min,
- io::{Cursor, Read, Seek, SeekFrom},
- sync::Arc,
-};
-
-use byteorder::{ByteOrder, LittleEndian};
-use parquet_format::{ColumnOrder as TColumnOrder, FileMetaData as TFileMetaData};
-use thrift::protocol::TCompactInputProtocol;
-
-use crate::basic::ColumnOrder;
-
-use crate::errors::{ParquetError, Result};
-use crate::file::{
- metadata::*, reader::ChunkReader, DEFAULT_FOOTER_READ_SIZE, FOOTER_SIZE,
- PARQUET_MAGIC,
-};
-
-use crate::schema::types::{self, SchemaDescriptor};
-
-/// Layout of Parquet file
-/// +---------------------------+-----+---+
-/// | Rest of file | B | A |
-/// +---------------------------+-----+---+
-/// where A: parquet footer, B: parquet metadata.
-///
-/// The reader first reads DEFAULT_FOOTER_SIZE bytes from the end of the file.
-/// If it is not enough according to the length indicated in the footer, it reads more bytes.
-pub fn parse_metadata<R: ChunkReader>(chunk_reader: &R) -> Result<ParquetMetaData> {
- // check file is large enough to hold footer
- let file_size = chunk_reader.len();
- if file_size < (FOOTER_SIZE as u64) {
- return Err(general_err!(
- "Invalid Parquet file. Size is smaller than footer"
- ));
- }
-
- // read and cache up to DEFAULT_FOOTER_READ_SIZE bytes from the end and process the footer
- let default_end_len = min(DEFAULT_FOOTER_READ_SIZE, chunk_reader.len() as usize);
- let mut default_end_reader = chunk_reader
- .get_read(chunk_reader.len() - default_end_len as u64, default_end_len)?;
- let mut default_len_end_buf = vec![0; default_end_len];
- default_end_reader.read_exact(&mut default_len_end_buf)?;
-
- // check this is indeed a parquet file
- if default_len_end_buf[default_end_len - 4..] != PARQUET_MAGIC {
- return Err(general_err!("Invalid Parquet file. Corrupt footer"));
- }
-
- // get the metadata length from the footer
- let metadata_len = LittleEndian::read_i32(
- &default_len_end_buf[default_end_len - 8..default_end_len - 4],
- ) as i64;
- if metadata_len < 0 {
- return Err(general_err!(
- "Invalid Parquet file. Metadata length is less than zero ({})",
- metadata_len
- ));
- }
- let footer_metadata_len = FOOTER_SIZE + metadata_len as usize;
-
- // build up the reader covering the entire metadata
- let mut default_end_cursor = Cursor::new(default_len_end_buf);
- let metadata_read: Box<dyn Read>;
- if footer_metadata_len > file_size as usize {
- return Err(general_err!(
- "Invalid Parquet file. Metadata start is less than zero ({})",
- file_size as i64 - footer_metadata_len as i64
- ));
- } else if footer_metadata_len < DEFAULT_FOOTER_READ_SIZE {
- // the whole metadata is in the bytes we already read
- default_end_cursor.seek(SeekFrom::End(-(footer_metadata_len as i64)))?;
- metadata_read = Box::new(default_end_cursor);
- } else {
- // the end of file read by default is not long enough, read missing bytes
- let complementary_end_read = chunk_reader.get_read(
- file_size - footer_metadata_len as u64,
- FOOTER_SIZE + metadata_len as usize - default_end_len,
- )?;
- metadata_read = Box::new(complementary_end_read.chain(default_end_cursor));
- }
-
- // TODO: row group filtering
- let mut prot = TCompactInputProtocol::new(metadata_read);
- let t_file_metadata: TFileMetaData = TFileMetaData::read_from_in_protocol(&mut prot)
- .map_err(|e| ParquetError::General(format!("Could not parse metadata: {}", e)))?;
- let schema = types::from_thrift(&t_file_metadata.schema)?;
- let schema_descr = Arc::new(SchemaDescriptor::new(schema));
- let mut row_groups = Vec::new();
- for rg in t_file_metadata.row_groups {
- row_groups.push(RowGroupMetaData::from_thrift(schema_descr.clone(), rg)?);
- }
- let column_orders = parse_column_orders(t_file_metadata.column_orders, &schema_descr);
-
- let file_metadata = FileMetaData::new(
- t_file_metadata.version,
- t_file_metadata.num_rows,
- t_file_metadata.created_by,
- t_file_metadata.key_value_metadata,
- schema_descr,
- column_orders,
- );
- Ok(ParquetMetaData::new(file_metadata, row_groups))
-}
-
-/// Parses column orders from Thrift definition.
-/// If no column orders are defined, returns `None`.
-fn parse_column_orders(
- t_column_orders: Option<Vec<TColumnOrder>>,
- schema_descr: &SchemaDescriptor,
-) -> Option<Vec<ColumnOrder>> {
- match t_column_orders {
- Some(orders) => {
- // Should always be the case
- assert_eq!(
- orders.len(),
- schema_descr.num_columns(),
- "Column order length mismatch"
- );
- let mut res = Vec::new();
- for (i, column) in schema_descr.columns().iter().enumerate() {
- match orders[i] {
- TColumnOrder::TYPEORDER(_) => {
- let sort_order = ColumnOrder::get_sort_order(
- column.logical_type(),
- column.converted_type(),
- column.physical_type(),
- );
- res.push(ColumnOrder::TYPE_DEFINED_ORDER(sort_order));
- }
- }
- }
- Some(res)
- }
- None => None,
- }
-}
-
-#[cfg(test)]
-mod tests {
- use super::*;
-
- use crate::basic::SortOrder;
- use crate::basic::Type;
- use crate::schema::types::Type as SchemaType;
- use crate::util::test_common::get_temp_file;
- use parquet_format::TypeDefinedOrder;
-
- #[test]
- fn test_parse_metadata_size_smaller_than_footer() {
- let test_file = get_temp_file("corrupt-1.parquet", &[]);
- let reader_result = parse_metadata(&test_file);
- assert!(reader_result.is_err());
- assert_eq!(
- reader_result.err().unwrap(),
- general_err!("Invalid Parquet file. Size is smaller than footer")
- );
- }
-
- #[test]
- fn test_parse_metadata_corrupt_footer() {
- let test_file = get_temp_file("corrupt-2.parquet", &[1, 2, 3, 4, 5, 6, 7, 8]);
- let reader_result = parse_metadata(&test_file);
- assert!(reader_result.is_err());
- assert_eq!(
- reader_result.err().unwrap(),
- general_err!("Invalid Parquet file. Corrupt footer")
- );
- }
-
- #[test]
- fn test_parse_metadata_invalid_length() {
- let test_file =
- get_temp_file("corrupt-3.parquet", &[0, 0, 0, 255, b'P', b'A', b'R', b'1']);
- let reader_result = parse_metadata(&test_file);
- assert!(reader_result.is_err());
- assert_eq!(
- reader_result.err().unwrap(),
- general_err!(
- "Invalid Parquet file. Metadata length is less than zero (-16777216)"
- )
- );
- }
-
- #[test]
- fn test_parse_metadata_invalid_start() {
- let test_file =
- get_temp_file("corrupt-4.parquet", &[255, 0, 0, 0, b'P', b'A', b'R', b'1']);
- let reader_result = parse_metadata(&test_file);
- assert!(reader_result.is_err());
- assert_eq!(
- reader_result.err().unwrap(),
- general_err!("Invalid Parquet file. Metadata start is less than zero (-255)")
- );
- }
-
- #[test]
- fn test_metadata_column_orders_parse() {
- // Define simple schema, we do not need to provide logical types.
- let mut fields = vec![
- Arc::new(
- SchemaType::primitive_type_builder("col1", Type::INT32)
- .build()
- .unwrap(),
- ),
- Arc::new(
- SchemaType::primitive_type_builder("col2", Type::FLOAT)
- .build()
- .unwrap(),
- ),
- ];
- let schema = SchemaType::group_type_builder("schema")
- .with_fields(&mut fields)
- .build()
- .unwrap();
- let schema_descr = SchemaDescriptor::new(Arc::new(schema));
-
- let t_column_orders = Some(vec![
- TColumnOrder::TYPEORDER(TypeDefinedOrder::new()),
- TColumnOrder::TYPEORDER(TypeDefinedOrder::new()),
- ]);
-
- assert_eq!(
- parse_column_orders(t_column_orders, &schema_descr),
- Some(vec![
- ColumnOrder::TYPE_DEFINED_ORDER(SortOrder::SIGNED),
- ColumnOrder::TYPE_DEFINED_ORDER(SortOrder::SIGNED)
- ])
- );
-
- // Test when no column orders are defined.
- assert_eq!(parse_column_orders(None, &schema_descr), None);
- }
-
- #[test]
- #[should_panic(expected = "Column order length mismatch")]
- fn test_metadata_column_orders_len_mismatch() {
- let schema = SchemaType::group_type_builder("schema").build().unwrap();
- let schema_descr = SchemaDescriptor::new(Arc::new(schema));
-
- let t_column_orders =
- Some(vec![TColumnOrder::TYPEORDER(TypeDefinedOrder::new())]);
-
- parse_column_orders(t_column_orders, &schema_descr);
- }
-}
diff --git a/parquet/src/file/metadata.rs b/parquet/src/file/metadata.rs
deleted file mode 100644
index 150c42c..0000000
--- a/parquet/src/file/metadata.rs
+++ /dev/null
@@ -1,789 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-//! Contains information about available Parquet metadata.
-//!
-//! The hierarchy of metadata is as follows:
-//!
-//! [`ParquetMetaData`](struct.ParquetMetaData.html) contains
-//! [`FileMetaData`](struct.FileMetaData.html) and zero or more
-//! [`RowGroupMetaData`](struct.RowGroupMetaData.html) for each row group.
-//!
-//! [`FileMetaData`](struct.FileMetaData.html) includes file version, application specific
-//! metadata.
-//!
-//! Each [`RowGroupMetaData`](struct.RowGroupMetaData.html) contains information about row
-//! group and one or more [`ColumnChunkMetaData`](struct.ColumnChunkMetaData.html) for
-//! each column chunk.
-//!
-//! [`ColumnChunkMetaData`](struct.ColumnChunkMetaData.html) has information about column
-//! chunk (primitive leaf column), including encoding/compression, number of values, etc.
-
-use std::sync::Arc;
-
-use parquet_format::{ColumnChunk, ColumnMetaData, RowGroup};
-
-use crate::basic::{ColumnOrder, Compression, Encoding, Type};
-use crate::errors::{ParquetError, Result};
-use crate::file::statistics::{self, Statistics};
-use crate::schema::types::{
- ColumnDescPtr, ColumnDescriptor, ColumnPath, SchemaDescPtr, SchemaDescriptor,
- Type as SchemaType,
-};
-
-/// Global Parquet metadata.
-#[derive(Debug, Clone)]
-pub struct ParquetMetaData {
- file_metadata: FileMetaData,
- row_groups: Vec<RowGroupMetaData>,
-}
-
-impl ParquetMetaData {
- /// Creates Parquet metadata from file metadata and a list of row group metadata `Arc`s
- /// for each available row group.
- pub fn new(file_metadata: FileMetaData, row_groups: Vec<RowGroupMetaData>) -> Self {
- ParquetMetaData {
- file_metadata,
- row_groups,
- }
- }
-
- /// Returns file metadata as reference.
- pub fn file_metadata(&self) -> &FileMetaData {
- &self.file_metadata
- }
-
- /// Returns number of row groups in this file.
- pub fn num_row_groups(&self) -> usize {
- self.row_groups.len()
- }
-
- /// Returns row group metadata for `i`th position.
- /// Position should be less than number of row groups `num_row_groups`.
- pub fn row_group(&self, i: usize) -> &RowGroupMetaData {
- &self.row_groups[i]
- }
-
- /// Returns slice of row groups in this file.
- pub fn row_groups(&self) -> &[RowGroupMetaData] {
- &self.row_groups
- }
-}
-
-pub type KeyValue = parquet_format::KeyValue;
-
-/// Reference counted pointer for [`FileMetaData`].
-pub type FileMetaDataPtr = Arc<FileMetaData>;
-
-/// Metadata for a Parquet file.
-#[derive(Debug, Clone)]
-pub struct FileMetaData {
- version: i32,
- num_rows: i64,
- created_by: Option<String>,
- key_value_metadata: Option<Vec<KeyValue>>,
- schema_descr: SchemaDescPtr,
- column_orders: Option<Vec<ColumnOrder>>,
-}
-
-impl FileMetaData {
- /// Creates new file metadata.
- pub fn new(
- version: i32,
- num_rows: i64,
- created_by: Option<String>,
- key_value_metadata: Option<Vec<KeyValue>>,
- schema_descr: SchemaDescPtr,
- column_orders: Option<Vec<ColumnOrder>>,
- ) -> Self {
- FileMetaData {
- version,
- num_rows,
- created_by,
- key_value_metadata,
- schema_descr,
- column_orders,
- }
- }
-
- /// Returns version of this file.
- pub fn version(&self) -> i32 {
- self.version
- }
-
- /// Returns number of rows in the file.
- pub fn num_rows(&self) -> i64 {
- self.num_rows
- }
-
- /// String message for application that wrote this file.
- ///
- /// This should have the following format:
- /// `<application> version <application version> (build <application build hash>)`.
- ///
- /// ```shell
- /// parquet-mr version 1.8.0 (build 0fda28af84b9746396014ad6a415b90592a98b3b)
- /// ```
- pub fn created_by(&self) -> &Option<String> {
- &self.created_by
- }
-
- /// Returns key_value_metadata of this file.
- pub fn key_value_metadata(&self) -> &Option<Vec<KeyValue>> {
- &self.key_value_metadata
- }
-
- /// Returns Parquet ['Type`] that describes schema in this file.
- pub fn schema(&self) -> &SchemaType {
- self.schema_descr.root_schema()
- }
-
- /// Returns a reference to schema descriptor.
- pub fn schema_descr(&self) -> &SchemaDescriptor {
- &self.schema_descr
- }
-
- /// Returns reference counted clone for schema descriptor.
- pub fn schema_descr_ptr(&self) -> SchemaDescPtr {
- self.schema_descr.clone()
- }
-
- /// Column (sort) order used for `min` and `max` values of each column in this file.
- ///
- /// Each column order corresponds to one column, determined by its position in the
- /// list, matching the position of the column in the schema.
- ///
- /// When `None` is returned, there are no column orders available, and each column
- /// should be assumed to have undefined (legacy) column order.
- pub fn column_orders(&self) -> Option<&Vec<ColumnOrder>> {
- self.column_orders.as_ref()
- }
-
- /// Returns column order for `i`th column in this file.
- /// If column orders are not available, returns undefined (legacy) column order.
- pub fn column_order(&self, i: usize) -> ColumnOrder {
- self.column_orders
- .as_ref()
- .map(|data| data[i])
- .unwrap_or(ColumnOrder::UNDEFINED)
- }
-}
-
-/// Reference counted pointer for [`RowGroupMetaData`].
-pub type RowGroupMetaDataPtr = Arc<RowGroupMetaData>;
-
-/// Metadata for a row group.
-#[derive(Debug, Clone)]
-pub struct RowGroupMetaData {
- columns: Vec<ColumnChunkMetaData>,
- num_rows: i64,
- total_byte_size: i64,
- schema_descr: SchemaDescPtr,
-}
-
-impl RowGroupMetaData {
- /// Returns builer for row group metadata.
- pub fn builder(schema_descr: SchemaDescPtr) -> RowGroupMetaDataBuilder {
- RowGroupMetaDataBuilder::new(schema_descr)
- }
-
- /// Number of columns in this row group.
- pub fn num_columns(&self) -> usize {
- self.columns.len()
- }
-
- /// Returns column chunk metadata for `i`th column.
- pub fn column(&self, i: usize) -> &ColumnChunkMetaData {
- &self.columns[i]
- }
-
- /// Returns slice of column chunk metadata.
- pub fn columns(&self) -> &[ColumnChunkMetaData] {
- &self.columns
- }
-
- /// Number of rows in this row group.
- pub fn num_rows(&self) -> i64 {
- self.num_rows
- }
-
- /// Total byte size of all uncompressed column data in this row group.
- pub fn total_byte_size(&self) -> i64 {
- self.total_byte_size
- }
-
- /// Total size of all compressed column data in this row group.
- pub fn compressed_size(&self) -> i64 {
- self.columns.iter().map(|c| c.total_compressed_size).sum()
- }
-
- /// Returns reference to a schema descriptor.
- pub fn schema_descr(&self) -> &SchemaDescriptor {
- self.schema_descr.as_ref()
- }
-
- /// Returns reference counted clone of schema descriptor.
- pub fn schema_descr_ptr(&self) -> SchemaDescPtr {
- self.schema_descr.clone()
- }
-
- /// Method to convert from Thrift.
- pub fn from_thrift(
- schema_descr: SchemaDescPtr,
- mut rg: RowGroup,
- ) -> Result<RowGroupMetaData> {
- assert_eq!(schema_descr.num_columns(), rg.columns.len());
- let total_byte_size = rg.total_byte_size;
- let num_rows = rg.num_rows;
- let mut columns = vec![];
- for (c, d) in rg.columns.drain(0..).zip(schema_descr.columns()) {
- let cc = ColumnChunkMetaData::from_thrift(d.clone(), c)?;
- columns.push(cc);
- }
- Ok(RowGroupMetaData {
- columns,
- num_rows,
- total_byte_size,
- schema_descr,
- })
- }
-
- /// Method to convert to Thrift.
- pub fn to_thrift(&self) -> RowGroup {
- RowGroup {
- columns: self.columns().iter().map(|v| v.to_thrift()).collect(),
- total_byte_size: self.total_byte_size,
- num_rows: self.num_rows,
- sorting_columns: None,
- }
- }
-}
-
-/// Builder for row group metadata.
-pub struct RowGroupMetaDataBuilder {
- columns: Vec<ColumnChunkMetaData>,
- schema_descr: SchemaDescPtr,
- num_rows: i64,
- total_byte_size: i64,
-}
-
-impl RowGroupMetaDataBuilder {
- /// Creates new builder from schema descriptor.
- fn new(schema_descr: SchemaDescPtr) -> Self {
- Self {
- columns: Vec::with_capacity(schema_descr.num_columns()),
- schema_descr,
- num_rows: 0,
- total_byte_size: 0,
- }
- }
-
- /// Sets number of rows in this row group.
- pub fn set_num_rows(mut self, value: i64) -> Self {
- self.num_rows = value;
- self
- }
-
- /// Sets total size in bytes for this row group.
- pub fn set_total_byte_size(mut self, value: i64) -> Self {
- self.total_byte_size = value;
- self
- }
-
- /// Sets column metadata for this row group.
- pub fn set_column_metadata(mut self, value: Vec<ColumnChunkMetaData>) -> Self {
- self.columns = value;
- self
- }
-
- /// Builds row group metadata.
- pub fn build(self) -> Result<RowGroupMetaData> {
- if self.schema_descr.num_columns() != self.columns.len() {
- return Err(general_err!(
- "Column length mismatch: {} != {}",
- self.schema_descr.num_columns(),
- self.columns.len()
- ));
- }
-
- Ok(RowGroupMetaData {
- columns: self.columns,
- num_rows: self.num_rows,
- total_byte_size: self.total_byte_size,
- schema_descr: self.schema_descr,
- })
- }
-}
-
-/// Metadata for a column chunk.
-#[derive(Debug, Clone)]
-pub struct ColumnChunkMetaData {
- column_type: Type,
- column_path: ColumnPath,
- column_descr: ColumnDescPtr,
- encodings: Vec<Encoding>,
- file_path: Option<String>,
- file_offset: i64,
- num_values: i64,
- compression: Compression,
- total_compressed_size: i64,
- total_uncompressed_size: i64,
- data_page_offset: i64,
- index_page_offset: Option<i64>,
- dictionary_page_offset: Option<i64>,
- statistics: Option<Statistics>,
-}
-
-/// Represents common operations for a column chunk.
-impl ColumnChunkMetaData {
- /// Returns builder for column chunk metadata.
- pub fn builder(column_descr: ColumnDescPtr) -> ColumnChunkMetaDataBuilder {
- ColumnChunkMetaDataBuilder::new(column_descr)
- }
-
- /// File where the column chunk is stored.
- ///
- /// If not set, assumed to belong to the same file as the metadata.
- /// This path is relative to the current file.
- pub fn file_path(&self) -> Option<&String> {
- self.file_path.as_ref()
- }
-
- /// Byte offset in `file_path()`.
- pub fn file_offset(&self) -> i64 {
- self.file_offset
- }
-
- /// Type of this column. Must be primitive.
- pub fn column_type(&self) -> Type {
- self.column_type
- }
-
- /// Path (or identifier) of this column.
- pub fn column_path(&self) -> &ColumnPath {
- &self.column_path
- }
-
- /// Descriptor for this column.
- pub fn column_descr(&self) -> &ColumnDescriptor {
- self.column_descr.as_ref()
- }
-
- /// Reference counted clone of descriptor for this column.
- pub fn column_descr_ptr(&self) -> ColumnDescPtr {
- self.column_descr.clone()
- }
-
- /// All encodings used for this column.
- pub fn encodings(&self) -> &Vec<Encoding> {
- &self.encodings
- }
-
- /// Total number of values in this column chunk.
- pub fn num_values(&self) -> i64 {
- self.num_values
- }
-
- /// Compression for this column.
- pub fn compression(&self) -> Compression {
- self.compression
- }
-
- /// Returns the total compressed data size of this column chunk.
- pub fn compressed_size(&self) -> i64 {
- self.total_compressed_size
- }
-
- /// Returns the total uncompressed data size of this column chunk.
- pub fn uncompressed_size(&self) -> i64 {
- self.total_uncompressed_size
- }
-
- /// Returns the offset for the column data.
- pub fn data_page_offset(&self) -> i64 {
- self.data_page_offset
- }
-
- /// Returns `true` if this column chunk contains a index page, `false` otherwise.
- pub fn has_index_page(&self) -> bool {
- self.index_page_offset.is_some()
- }
-
- /// Returns the offset for the index page.
- pub fn index_page_offset(&self) -> Option<i64> {
- self.index_page_offset
- }
-
- /// Returns `true` if this column chunk contains a dictionary page, `false` otherwise.
- pub fn has_dictionary_page(&self) -> bool {
- self.dictionary_page_offset.is_some()
- }
-
- /// Returns the offset for the dictionary page, if any.
- pub fn dictionary_page_offset(&self) -> Option<i64> {
- self.dictionary_page_offset
- }
-
- /// Returns the offset and length in bytes of the column chunk within the file
- pub fn byte_range(&self) -> (u64, u64) {
- let col_start = if self.has_dictionary_page() {
- self.dictionary_page_offset().unwrap()
- } else {
- self.data_page_offset()
- };
- let col_len = self.compressed_size();
- assert!(
- col_start >= 0 && col_len >= 0,
- "column start and length should not be negative"
- );
- (col_start as u64, col_len as u64)
- }
-
- /// Returns statistics that are set for this column chunk,
- /// or `None` if no statistics are available.
- pub fn statistics(&self) -> Option<&Statistics> {
- self.statistics.as_ref()
- }
-
- /// Method to convert from Thrift.
- pub fn from_thrift(column_descr: ColumnDescPtr, cc: ColumnChunk) -> Result<Self> {
- if cc.meta_data.is_none() {
- return Err(general_err!("Expected to have column metadata"));
- }
- let mut col_metadata: ColumnMetaData = cc.meta_data.unwrap();
- let column_type = Type::from(col_metadata.type_);
- let column_path = ColumnPath::new(col_metadata.path_in_schema);
- let encodings = col_metadata
- .encodings
- .drain(0..)
- .map(Encoding::from)
- .collect();
- let compression = Compression::from(col_metadata.codec);
- let file_path = cc.file_path;
- let file_offset = cc.file_offset;
- let num_values = col_metadata.num_values;
- let total_compressed_size = col_metadata.total_compressed_size;
- let total_uncompressed_size = col_metadata.total_uncompressed_size;
- let data_page_offset = col_metadata.data_page_offset;
- let index_page_offset = col_metadata.index_page_offset;
- let dictionary_page_offset = col_metadata.dictionary_page_offset;
- let statistics = statistics::from_thrift(column_type, col_metadata.statistics);
- let result = ColumnChunkMetaData {
- column_type,
- column_path,
- column_descr,
- encodings,
- file_path,
- file_offset,
- num_values,
- compression,
- total_compressed_size,
- total_uncompressed_size,
- data_page_offset,
- index_page_offset,
- dictionary_page_offset,
- statistics,
- };
- Ok(result)
- }
-
- /// Method to convert to Thrift.
- pub fn to_thrift(&self) -> ColumnChunk {
- let column_metadata = ColumnMetaData {
- type_: self.column_type.into(),
- encodings: self.encodings().iter().map(|&v| v.into()).collect(),
- path_in_schema: Vec::from(self.column_path.as_ref()),
- codec: self.compression.into(),
- num_values: self.num_values,
- total_uncompressed_size: self.total_uncompressed_size,
- total_compressed_size: self.total_compressed_size,
- key_value_metadata: None,
- data_page_offset: self.data_page_offset,
- index_page_offset: self.index_page_offset,
- dictionary_page_offset: self.dictionary_page_offset,
- statistics: statistics::to_thrift(self.statistics.as_ref()),
- encoding_stats: None,
- };
-
- ColumnChunk {
- file_path: self.file_path().cloned(),
- file_offset: self.file_offset,
- meta_data: Some(column_metadata),
- offset_index_offset: None,
- offset_index_length: None,
- column_index_offset: None,
- column_index_length: None,
- }
- }
-}
-
-/// Builder for column chunk metadata.
-pub struct ColumnChunkMetaDataBuilder {
- column_descr: ColumnDescPtr,
- encodings: Vec<Encoding>,
- file_path: Option<String>,
- file_offset: i64,
- num_values: i64,
- compression: Compression,
- total_compressed_size: i64,
- total_uncompressed_size: i64,
- data_page_offset: i64,
- index_page_offset: Option<i64>,
- dictionary_page_offset: Option<i64>,
- statistics: Option<Statistics>,
-}
-
-impl ColumnChunkMetaDataBuilder {
- /// Creates new column chunk metadata builder.
- fn new(column_descr: ColumnDescPtr) -> Self {
- Self {
- column_descr,
- encodings: Vec::new(),
- file_path: None,
- file_offset: 0,
- num_values: 0,
- compression: Compression::UNCOMPRESSED,
- total_compressed_size: 0,
- total_uncompressed_size: 0,
- data_page_offset: 0,
- index_page_offset: None,
- dictionary_page_offset: None,
- statistics: None,
- }
- }
-
- /// Sets list of encodings for this column chunk.
- pub fn set_encodings(mut self, encodings: Vec<Encoding>) -> Self {
- self.encodings = encodings;
- self
- }
-
- /// Sets optional file path for this column chunk.
- pub fn set_file_path(mut self, value: String) -> Self {
- self.file_path = Some(value);
- self
- }
-
- /// Sets file offset in bytes.
- pub fn set_file_offset(mut self, value: i64) -> Self {
- self.file_offset = value;
- self
- }
-
- /// Sets number of values.
- pub fn set_num_values(mut self, value: i64) -> Self {
- self.num_values = value;
- self
- }
-
- /// Sets compression.
- pub fn set_compression(mut self, value: Compression) -> Self {
- self.compression = value;
- self
- }
-
- /// Sets total compressed size in bytes.
- pub fn set_total_compressed_size(mut self, value: i64) -> Self {
- self.total_compressed_size = value;
- self
- }
-
- /// Sets total uncompressed size in bytes.
- pub fn set_total_uncompressed_size(mut self, value: i64) -> Self {
- self.total_uncompressed_size = value;
- self
- }
-
- /// Sets data page offset in bytes.
- pub fn set_data_page_offset(mut self, value: i64) -> Self {
- self.data_page_offset = value;
- self
- }
-
- /// Sets optional dictionary page ofset in bytes.
- pub fn set_dictionary_page_offset(mut self, value: Option<i64>) -> Self {
- self.dictionary_page_offset = value;
- self
- }
-
- /// Sets optional index page offset in bytes.
- pub fn set_index_page_offset(mut self, value: Option<i64>) -> Self {
- self.index_page_offset = value;
- self
- }
-
- /// Sets statistics for this column chunk.
- pub fn set_statistics(mut self, value: Statistics) -> Self {
- self.statistics = Some(value);
- self
- }
-
- /// Builds column chunk metadata.
- pub fn build(self) -> Result<ColumnChunkMetaData> {
- Ok(ColumnChunkMetaData {
- column_type: self.column_descr.physical_type(),
- column_path: self.column_descr.path().clone(),
- column_descr: self.column_descr,
- encodings: self.encodings,
- file_path: self.file_path,
- file_offset: self.file_offset,
- num_values: self.num_values,
- compression: self.compression,
- total_compressed_size: self.total_compressed_size,
- total_uncompressed_size: self.total_uncompressed_size,
- data_page_offset: self.data_page_offset,
- index_page_offset: self.index_page_offset,
- dictionary_page_offset: self.dictionary_page_offset,
- statistics: self.statistics,
- })
- }
-}
-
-#[cfg(test)]
-mod tests {
- use super::*;
-
- #[test]
- fn test_row_group_metadata_thrift_conversion() {
- let schema_descr = get_test_schema_descr();
-
- let mut columns = vec![];
- for ptr in schema_descr.columns() {
- let column = ColumnChunkMetaData::builder(ptr.clone()).build().unwrap();
- columns.push(column);
- }
- let row_group_meta = RowGroupMetaData::builder(schema_descr.clone())
- .set_num_rows(1000)
- .set_total_byte_size(2000)
- .set_column_metadata(columns)
- .build()
- .unwrap();
-
- let row_group_exp = row_group_meta.to_thrift();
- let row_group_res =
- RowGroupMetaData::from_thrift(schema_descr, row_group_exp.clone())
- .unwrap()
- .to_thrift();
-
- assert_eq!(row_group_res, row_group_exp);
- }
-
- #[test]
- fn test_row_group_metadata_thrift_conversion_empty() {
- let schema_descr = get_test_schema_descr();
-
- let row_group_meta = RowGroupMetaData::builder(schema_descr).build();
-
- assert!(row_group_meta.is_err());
- if let Err(e) = row_group_meta {
- assert_eq!(
- format!("{}", e),
- "Parquet error: Column length mismatch: 2 != 0"
- );
- }
- }
-
- #[test]
- fn test_column_chunk_metadata_thrift_conversion() {
- let column_descr = get_test_schema_descr().column(0);
-
- let col_metadata = ColumnChunkMetaData::builder(column_descr.clone())
- .set_encodings(vec![Encoding::PLAIN, Encoding::RLE])
- .set_file_path("file_path".to_owned())
- .set_file_offset(100)
- .set_num_values(1000)
- .set_compression(Compression::SNAPPY)
- .set_total_compressed_size(2000)
- .set_total_uncompressed_size(3000)
- .set_data_page_offset(4000)
- .set_dictionary_page_offset(Some(5000))
- .build()
- .unwrap();
-
- let col_chunk_exp = col_metadata.to_thrift();
-
- let col_chunk_res =
- ColumnChunkMetaData::from_thrift(column_descr, col_chunk_exp.clone())
- .unwrap()
- .to_thrift();
-
- assert_eq!(col_chunk_res, col_chunk_exp);
- }
-
- #[test]
- fn test_column_chunk_metadata_thrift_conversion_empty() {
- let column_descr = get_test_schema_descr().column(0);
-
- let col_metadata = ColumnChunkMetaData::builder(column_descr.clone())
- .build()
- .unwrap();
-
- let col_chunk_exp = col_metadata.to_thrift();
- let col_chunk_res =
- ColumnChunkMetaData::from_thrift(column_descr, col_chunk_exp.clone())
- .unwrap()
- .to_thrift();
-
- assert_eq!(col_chunk_res, col_chunk_exp);
- }
-
- #[test]
- fn test_compressed_size() {
- let schema_descr = get_test_schema_descr();
-
- let mut columns = vec![];
- for column_descr in schema_descr.columns() {
- let column = ColumnChunkMetaData::builder(column_descr.clone())
- .set_total_compressed_size(500)
- .set_total_uncompressed_size(700)
- .build()
- .unwrap();
- columns.push(column);
- }
- let row_group_meta = RowGroupMetaData::builder(schema_descr)
- .set_num_rows(1000)
- .set_column_metadata(columns)
- .build()
- .unwrap();
-
- let compressed_size_res: i64 = row_group_meta.compressed_size();
- let compressed_size_exp: i64 = 1000;
-
- assert_eq!(compressed_size_res, compressed_size_exp);
- }
-
- /// Returns sample schema descriptor so we can create column metadata.
- fn get_test_schema_descr() -> SchemaDescPtr {
- let schema = SchemaType::group_type_builder("schema")
- .with_fields(&mut vec![
- Arc::new(
- SchemaType::primitive_type_builder("a", Type::INT32)
- .build()
- .unwrap(),
- ),
- Arc::new(
- SchemaType::primitive_type_builder("b", Type::INT32)
- .build()
- .unwrap(),
- ),
- ])
- .build()
- .unwrap();
-
- Arc::new(SchemaDescriptor::new(Arc::new(schema)))
- }
-}
diff --git a/parquet/src/file/mod.rs b/parquet/src/file/mod.rs
deleted file mode 100644
index f85de98..0000000
--- a/parquet/src/file/mod.rs
+++ /dev/null
@@ -1,110 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-//! Main entrypoint for working with Parquet API.
-//!
-//! Provides access to file and row group readers and writers, record API, metadata, etc.
-//!
-//! See [`reader::SerializedFileReader`](reader/struct.SerializedFileReader.html) or
-//! [`writer::SerializedFileWriter`](writer/struct.SerializedFileWriter.html) for a
-//! starting reference, [`metadata::ParquetMetaData`](metadata/index.html) for file
-//! metadata, and [`statistics`](statistics/index.html) for working with statistics.
-//!
-//! # Example of writing a new file
-//!
-//! ```rust,no_run
-//! use std::{fs, path::Path, sync::Arc};
-//!
-//! use parquet::{
-//! file::{
-//! properties::WriterProperties,
-//! writer::{FileWriter, SerializedFileWriter},
-//! },
-//! schema::parser::parse_message_type,
-//! };
-//!
-//! let path = Path::new("/path/to/sample.parquet");
-//!
-//! let message_type = "
-//! message schema {
-//! REQUIRED INT32 b;
-//! }
-//! ";
-//! let schema = Arc::new(parse_message_type(message_type).unwrap());
-//! let props = Arc::new(WriterProperties::builder().build());
-//! let file = fs::File::create(&path).unwrap();
-//! let mut writer = SerializedFileWriter::new(file, schema, props).unwrap();
-//! let mut row_group_writer = writer.next_row_group().unwrap();
-//! while let Some(mut col_writer) = row_group_writer.next_column().unwrap() {
-//! // ... write values to a column writer
-//! row_group_writer.close_column(col_writer).unwrap();
-//! }
-//! writer.close_row_group(row_group_writer).unwrap();
-//! writer.close().unwrap();
-//!
-//! let bytes = fs::read(&path).unwrap();
-//! assert_eq!(&bytes[0..4], &[b'P', b'A', b'R', b'1']);
-//! ```
-//! # Example of reading an existing file
-//!
-//! ```rust,no_run
-//! use parquet::file::reader::{FileReader, SerializedFileReader};
-//! use std::{fs::File, path::Path};
-//!
-//! let path = Path::new("/path/to/sample.parquet");
-//! if let Ok(file) = File::open(&path) {
-//! let reader = SerializedFileReader::new(file).unwrap();
-//!
-//! let parquet_metadata = reader.metadata();
-//! assert_eq!(parquet_metadata.num_row_groups(), 1);
-//!
-//! let row_group_reader = reader.get_row_group(0).unwrap();
-//! assert_eq!(row_group_reader.num_columns(), 1);
-//! }
-//! ```
-//! # Example of reading multiple files
-//!
-//! ```rust,no_run
-//! use parquet::file::reader::SerializedFileReader;
-//! use std::convert::TryFrom;
-//!
-//! let paths = vec![
-//! "/path/to/sample.parquet/part-1.snappy.parquet",
-//! "/path/to/sample.parquet/part-2.snappy.parquet"
-//! ];
-//! // Create a reader for each file and flat map rows
-//! let rows = paths.iter()
-//! .map(|p| SerializedFileReader::try_from(*p).unwrap())
-//! .flat_map(|r| r.into_iter());
-//!
-//! for row in rows {
-//! println!("{}", row);
-//! }
-//! ```
-pub mod footer;
-pub mod metadata;
-pub mod properties;
-pub mod reader;
-pub mod serialized_reader;
-pub mod statistics;
-pub mod writer;
-
-const FOOTER_SIZE: usize = 8;
-const PARQUET_MAGIC: [u8; 4] = [b'P', b'A', b'R', b'1'];
-
-/// The number of bytes read at the end of the parquet file on first read
-const DEFAULT_FOOTER_READ_SIZE: usize = 64 * 1024;
diff --git a/parquet/src/file/properties.rs b/parquet/src/file/properties.rs
deleted file mode 100644
index 0d0cbef..0000000
--- a/parquet/src/file/properties.rs
+++ /dev/null
@@ -1,680 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-//! Writer properties.
-//!
-//! # Usage
-//!
-//! ```rust
-//! use parquet::{
-//! basic::{Compression, Encoding},
-//! file::properties::*,
-//! schema::types::ColumnPath,
-//! };
-//!
-//! // Create properties with default configuration.
-//! let props = WriterProperties::builder().build();
-//!
-//! // Use properties builder to set certain options and assemble the configuration.
-//! let props = WriterProperties::builder()
-//! .set_writer_version(WriterVersion::PARQUET_1_0)
-//! .set_encoding(Encoding::PLAIN)
-//! .set_column_encoding(ColumnPath::from("col1"), Encoding::DELTA_BINARY_PACKED)
-//! .set_compression(Compression::SNAPPY)
-//! .build();
-//!
-//! assert_eq!(props.writer_version(), WriterVersion::PARQUET_1_0);
-//! assert_eq!(
-//! props.encoding(&ColumnPath::from("col1")),
-//! Some(Encoding::DELTA_BINARY_PACKED)
-//! );
-//! assert_eq!(
-//! props.encoding(&ColumnPath::from("col2")),
-//! Some(Encoding::PLAIN)
-//! );
-//! ```
-
-use std::{collections::HashMap, sync::Arc};
-
-use crate::basic::{Compression, Encoding};
-use crate::file::metadata::KeyValue;
-use crate::schema::types::ColumnPath;
-
-const DEFAULT_PAGE_SIZE: usize = 1024 * 1024;
-const DEFAULT_WRITE_BATCH_SIZE: usize = 1024;
-const DEFAULT_WRITER_VERSION: WriterVersion = WriterVersion::PARQUET_1_0;
-const DEFAULT_COMPRESSION: Compression = Compression::UNCOMPRESSED;
-const DEFAULT_DICTIONARY_ENABLED: bool = true;
-const DEFAULT_DICTIONARY_PAGE_SIZE_LIMIT: usize = DEFAULT_PAGE_SIZE;
-const DEFAULT_STATISTICS_ENABLED: bool = true;
-const DEFAULT_MAX_STATISTICS_SIZE: usize = 4096;
-const DEFAULT_MAX_ROW_GROUP_SIZE: usize = 128 * 1024 * 1024;
-const DEFAULT_CREATED_BY: &str = env!("PARQUET_CREATED_BY");
-
-/// Parquet writer version.
-///
-/// Basic constant, which is not part of the Thrift definition.
-#[derive(Debug, Clone, Copy, PartialEq)]
-pub enum WriterVersion {
- PARQUET_1_0,
- PARQUET_2_0,
-}
-
-impl WriterVersion {
- /// Returns writer version as `i32`.
- pub fn as_num(&self) -> i32 {
- match self {
- WriterVersion::PARQUET_1_0 => 1,
- WriterVersion::PARQUET_2_0 => 2,
- }
- }
-}
-
-/// Reference counted writer properties.
-pub type WriterPropertiesPtr = Arc<WriterProperties>;
-
-/// Writer properties.
-///
-/// All properties except the key-value metadata are immutable,
-/// use [`WriterPropertiesBuilder`] to assemble these properties.
-#[derive(Debug, Clone)]
-pub struct WriterProperties {
- data_pagesize_limit: usize,
- dictionary_pagesize_limit: usize,
- write_batch_size: usize,
- max_row_group_size: usize,
- writer_version: WriterVersion,
- created_by: String,
- pub(crate) key_value_metadata: Option<Vec<KeyValue>>,
- default_column_properties: ColumnProperties,
- column_properties: HashMap<ColumnPath, ColumnProperties>,
-}
-
-impl WriterProperties {
- /// Returns builder for writer properties with default values.
- pub fn builder() -> WriterPropertiesBuilder {
- WriterPropertiesBuilder::with_defaults()
- }
-
- /// Returns data page size limit.
- pub fn data_pagesize_limit(&self) -> usize {
- self.data_pagesize_limit
- }
-
- /// Returns dictionary page size limit.
- pub fn dictionary_pagesize_limit(&self) -> usize {
- self.dictionary_pagesize_limit
- }
-
- /// Returns configured batch size for writes.
- ///
- /// When writing a batch of data, this setting allows to split it internally into
- /// smaller batches so we can better estimate the size of a page currently being
- /// written.
- pub fn write_batch_size(&self) -> usize {
- self.write_batch_size
- }
-
- /// Returns max size for a row group.
- pub fn max_row_group_size(&self) -> usize {
- self.max_row_group_size
- }
-
- /// Returns configured writer version.
- pub fn writer_version(&self) -> WriterVersion {
- self.writer_version
- }
-
- /// Returns `created_by` string.
- pub fn created_by(&self) -> &str {
- &self.created_by
- }
-
- /// Returns `key_value_metadata` KeyValue pairs.
- pub fn key_value_metadata(&self) -> &Option<Vec<KeyValue>> {
- &self.key_value_metadata
- }
-
- /// Returns encoding for a data page, when dictionary encoding is enabled.
- /// This is not configurable.
- #[inline]
- pub fn dictionary_data_page_encoding(&self) -> Encoding {
- // PLAIN_DICTIONARY encoding is deprecated in writer version 1.
- // Dictionary values are encoded using RLE_DICTIONARY encoding.
- Encoding::RLE_DICTIONARY
- }
-
- /// Returns encoding for dictionary page, when dictionary encoding is enabled.
- /// This is not configurable.
- #[inline]
- pub fn dictionary_page_encoding(&self) -> Encoding {
- // PLAIN_DICTIONARY is deprecated in writer version 1.
- // Dictionary is encoded using plain encoding.
- Encoding::PLAIN
- }
-
- /// Returns encoding for a column, if set.
- /// In case when dictionary is enabled, returns fallback encoding.
- ///
- /// If encoding is not set, then column writer will choose the best encoding
- /// based on the column type.
- pub fn encoding(&self, col: &ColumnPath) -> Option<Encoding> {
- self.column_properties
- .get(col)
- .and_then(|c| c.encoding())
- .or_else(|| self.default_column_properties.encoding())
- }
-
- /// Returns compression codec for a column.
- pub fn compression(&self, col: &ColumnPath) -> Compression {
- self.column_properties
- .get(col)
- .and_then(|c| c.compression())
- .or_else(|| self.default_column_properties.compression())
- .unwrap_or(DEFAULT_COMPRESSION)
- }
-
- /// Returns `true` if dictionary encoding is enabled for a column.
- pub fn dictionary_enabled(&self, col: &ColumnPath) -> bool {
- self.column_properties
- .get(col)
- .and_then(|c| c.dictionary_enabled())
- .or_else(|| self.default_column_properties.dictionary_enabled())
- .unwrap_or(DEFAULT_DICTIONARY_ENABLED)
- }
-
- /// Returns `true` if statistics are enabled for a column.
- pub fn statistics_enabled(&self, col: &ColumnPath) -> bool {
- self.column_properties
- .get(col)
- .and_then(|c| c.statistics_enabled())
- .or_else(|| self.default_column_properties.statistics_enabled())
- .unwrap_or(DEFAULT_STATISTICS_ENABLED)
- }
-
- /// Returns max size for statistics.
- /// Only applicable if statistics are enabled.
- pub fn max_statistics_size(&self, col: &ColumnPath) -> usize {
- self.column_properties
- .get(col)
- .and_then(|c| c.max_statistics_size())
- .or_else(|| self.default_column_properties.max_statistics_size())
- .unwrap_or(DEFAULT_MAX_STATISTICS_SIZE)
- }
-}
-
-/// Writer properties builder.
-pub struct WriterPropertiesBuilder {
- data_pagesize_limit: usize,
- dictionary_pagesize_limit: usize,
- write_batch_size: usize,
- max_row_group_size: usize,
- writer_version: WriterVersion,
- created_by: String,
- key_value_metadata: Option<Vec<KeyValue>>,
- default_column_properties: ColumnProperties,
- column_properties: HashMap<ColumnPath, ColumnProperties>,
-}
-
-impl WriterPropertiesBuilder {
- /// Returns default state of the builder.
- fn with_defaults() -> Self {
- Self {
- data_pagesize_limit: DEFAULT_PAGE_SIZE,
- dictionary_pagesize_limit: DEFAULT_DICTIONARY_PAGE_SIZE_LIMIT,
- write_batch_size: DEFAULT_WRITE_BATCH_SIZE,
- max_row_group_size: DEFAULT_MAX_ROW_GROUP_SIZE,
- writer_version: DEFAULT_WRITER_VERSION,
- created_by: DEFAULT_CREATED_BY.to_string(),
- key_value_metadata: None,
- default_column_properties: ColumnProperties::new(),
- column_properties: HashMap::new(),
- }
- }
-
- /// Finalizes the configuration and returns immutable writer properties struct.
- pub fn build(self) -> WriterProperties {
- WriterProperties {
- data_pagesize_limit: self.data_pagesize_limit,
- dictionary_pagesize_limit: self.dictionary_pagesize_limit,
- write_batch_size: self.write_batch_size,
- max_row_group_size: self.max_row_group_size,
- writer_version: self.writer_version,
- created_by: self.created_by,
- key_value_metadata: self.key_value_metadata,
- default_column_properties: self.default_column_properties,
- column_properties: self.column_properties,
- }
- }
-
- // ----------------------------------------------------------------------
- // Writer properties related to a file
-
- /// Sets writer version.
- pub fn set_writer_version(mut self, value: WriterVersion) -> Self {
- self.writer_version = value;
- self
- }
-
- /// Sets data page size limit.
- pub fn set_data_pagesize_limit(mut self, value: usize) -> Self {
- self.data_pagesize_limit = value;
- self
- }
-
- /// Sets dictionary page size limit.
- pub fn set_dictionary_pagesize_limit(mut self, value: usize) -> Self {
- self.dictionary_pagesize_limit = value;
- self
- }
-
- /// Sets write batch size.
- pub fn set_write_batch_size(mut self, value: usize) -> Self {
- self.write_batch_size = value;
- self
- }
-
- /// Sets max size for a row group.
- pub fn set_max_row_group_size(mut self, value: usize) -> Self {
- assert!(value > 0, "Cannot have a 0 max row group size");
- self.max_row_group_size = value;
- self
- }
-
- /// Sets "created by" property.
- pub fn set_created_by(mut self, value: String) -> Self {
- self.created_by = value;
- self
- }
-
- /// Sets "key_value_metadata" property.
- pub fn set_key_value_metadata(mut self, value: Option<Vec<KeyValue>>) -> Self {
- self.key_value_metadata = value;
- self
- }
-
- // ----------------------------------------------------------------------
- // Setters for any column (global)
-
- /// Sets encoding for any column.
- ///
- /// If dictionary is not enabled, this is treated as a primary encoding for all
- /// columns. In case when dictionary is enabled for any column, this value is
- /// considered to be a fallback encoding for that column.
- ///
- /// Panics if user tries to set dictionary encoding here, regardless of dictionary
- /// encoding flag being set.
- pub fn set_encoding(mut self, value: Encoding) -> Self {
- self.default_column_properties.set_encoding(value);
- self
- }
-
- /// Sets compression codec for any column.
- pub fn set_compression(mut self, value: Compression) -> Self {
- self.default_column_properties.set_compression(value);
- self
- }
-
- /// Sets flag to enable/disable dictionary encoding for any column.
- ///
- /// Use this method to set dictionary encoding, instead of explicitly specifying
- /// encoding in `set_encoding` method.
- pub fn set_dictionary_enabled(mut self, value: bool) -> Self {
- self.default_column_properties.set_dictionary_enabled(value);
- self
- }
-
- /// Sets flag to enable/disable statistics for any column.
- pub fn set_statistics_enabled(mut self, value: bool) -> Self {
- self.default_column_properties.set_statistics_enabled(value);
- self
- }
-
- /// Sets max statistics size for any column.
- /// Applicable only if statistics are enabled.
- pub fn set_max_statistics_size(mut self, value: usize) -> Self {
- self.default_column_properties
- .set_max_statistics_size(value);
- self
- }
-
- // ----------------------------------------------------------------------
- // Setters for a specific column
-
- /// Helper method to get existing or new mutable reference of column properties.
- #[inline]
- fn get_mut_props(&mut self, col: ColumnPath) -> &mut ColumnProperties {
- self.column_properties
- .entry(col)
- .or_insert(ColumnProperties::new())
- }
-
- /// Sets encoding for a column.
- /// Takes precedence over globally defined settings.
- ///
- /// If dictionary is not enabled, this is treated as a primary encoding for this
- /// column. In case when dictionary is enabled for this column, either through
- /// global defaults or explicitly, this value is considered to be a fallback
- /// encoding for this column.
- ///
- /// Panics if user tries to set dictionary encoding here, regardless of dictionary
- /// encoding flag being set.
- pub fn set_column_encoding(mut self, col: ColumnPath, value: Encoding) -> Self {
- self.get_mut_props(col).set_encoding(value);
- self
- }
-
- /// Sets compression codec for a column.
- /// Takes precedence over globally defined settings.
- pub fn set_column_compression(mut self, col: ColumnPath, value: Compression) -> Self {
- self.get_mut_props(col).set_compression(value);
- self
- }
-
- /// Sets flag to enable/disable dictionary encoding for a column.
- /// Takes precedence over globally defined settings.
- pub fn set_column_dictionary_enabled(mut self, col: ColumnPath, value: bool) -> Self {
- self.get_mut_props(col).set_dictionary_enabled(value);
- self
- }
-
- /// Sets flag to enable/disable statistics for a column.
- /// Takes precedence over globally defined settings.
- pub fn set_column_statistics_enabled(mut self, col: ColumnPath, value: bool) -> Self {
- self.get_mut_props(col).set_statistics_enabled(value);
- self
- }
-
- /// Sets max size for statistics for a column.
- /// Takes precedence over globally defined settings.
- pub fn set_column_max_statistics_size(
- mut self,
- col: ColumnPath,
- value: usize,
- ) -> Self {
- self.get_mut_props(col).set_max_statistics_size(value);
- self
- }
-}
-
-/// Container for column properties that can be changed as part of writer.
-///
-/// If a field is `None`, it means that no specific value has been set for this column,
-/// so some subsequent or default value must be used.
-#[derive(Debug, Clone, PartialEq)]
-struct ColumnProperties {
- encoding: Option<Encoding>,
- codec: Option<Compression>,
- dictionary_enabled: Option<bool>,
- statistics_enabled: Option<bool>,
- max_statistics_size: Option<usize>,
-}
-
-impl ColumnProperties {
- /// Initialise column properties with default values.
- fn new() -> Self {
- Self {
- encoding: None,
- codec: None,
- dictionary_enabled: None,
- statistics_enabled: None,
- max_statistics_size: None,
- }
- }
-
- /// Sets encoding for this column.
- ///
- /// If dictionary is not enabled, this is treated as a primary encoding for a column.
- /// In case when dictionary is enabled for a column, this value is considered to
- /// be a fallback encoding.
- ///
- /// Panics if user tries to set dictionary encoding here, regardless of dictionary
- /// encoding flag being set. Use `set_dictionary_enabled` method to enable dictionary
- /// for a column.
- fn set_encoding(&mut self, value: Encoding) {
- if value == Encoding::PLAIN_DICTIONARY || value == Encoding::RLE_DICTIONARY {
- panic!("Dictionary encoding can not be used as fallback encoding");
- }
- self.encoding = Some(value);
- }
-
- /// Sets compression codec for this column.
- fn set_compression(&mut self, value: Compression) {
- self.codec = Some(value);
- }
-
- /// Sets whether or not dictionary encoding is enabled for this column.
- fn set_dictionary_enabled(&mut self, enabled: bool) {
- self.dictionary_enabled = Some(enabled);
- }
-
- /// Sets whether or not statistics are enabled for this column.
- fn set_statistics_enabled(&mut self, enabled: bool) {
- self.statistics_enabled = Some(enabled);
- }
-
- /// Sets max size for statistics for this column.
- fn set_max_statistics_size(&mut self, value: usize) {
- self.max_statistics_size = Some(value);
- }
-
- /// Returns optional encoding for this column.
- fn encoding(&self) -> Option<Encoding> {
- self.encoding
- }
-
- /// Returns optional compression codec for this column.
- fn compression(&self) -> Option<Compression> {
- self.codec
- }
-
- /// Returns `Some(true)` if dictionary encoding is enabled for this column, if
- /// disabled then returns `Some(false)`. If result is `None`, then no setting has
- /// been provided.
- fn dictionary_enabled(&self) -> Option<bool> {
- self.dictionary_enabled
- }
-
- /// Returns `Some(true)` if statistics are enabled for this column, if disabled then
- /// returns `Some(false)`. If result is `None`, then no setting has been provided.
- fn statistics_enabled(&self) -> Option<bool> {
- self.statistics_enabled
- }
-
- /// Returns optional max size in bytes for statistics.
- fn max_statistics_size(&self) -> Option<usize> {
- self.max_statistics_size
- }
-}
-
-#[cfg(test)]
-mod tests {
- use super::*;
-
- #[test]
- fn test_writer_version() {
- assert_eq!(WriterVersion::PARQUET_1_0.as_num(), 1);
- assert_eq!(WriterVersion::PARQUET_2_0.as_num(), 2);
- }
-
- #[test]
- fn test_writer_properties_default_settings() {
- let props = WriterProperties::builder().build();
- assert_eq!(props.data_pagesize_limit(), DEFAULT_PAGE_SIZE);
- assert_eq!(
- props.dictionary_pagesize_limit(),
- DEFAULT_DICTIONARY_PAGE_SIZE_LIMIT
- );
- assert_eq!(props.write_batch_size(), DEFAULT_WRITE_BATCH_SIZE);
- assert_eq!(props.max_row_group_size(), DEFAULT_MAX_ROW_GROUP_SIZE);
- assert_eq!(props.writer_version(), DEFAULT_WRITER_VERSION);
- assert_eq!(props.created_by(), DEFAULT_CREATED_BY);
- assert_eq!(props.key_value_metadata(), &None);
- assert_eq!(props.encoding(&ColumnPath::from("col")), None);
- assert_eq!(
- props.compression(&ColumnPath::from("col")),
- DEFAULT_COMPRESSION
- );
- assert_eq!(
- props.dictionary_enabled(&ColumnPath::from("col")),
- DEFAULT_DICTIONARY_ENABLED
- );
- assert_eq!(
- props.statistics_enabled(&ColumnPath::from("col")),
- DEFAULT_STATISTICS_ENABLED
- );
- assert_eq!(
- props.max_statistics_size(&ColumnPath::from("col")),
- DEFAULT_MAX_STATISTICS_SIZE
- );
- }
-
- #[test]
- fn test_writer_properties_dictionary_encoding() {
- // dictionary encoding is not configurable, and it should be the same for both
- // writer version 1 and 2.
- for version in &[WriterVersion::PARQUET_1_0, WriterVersion::PARQUET_2_0] {
- let props = WriterProperties::builder()
- .set_writer_version(*version)
- .build();
- assert_eq!(props.dictionary_page_encoding(), Encoding::PLAIN);
- assert_eq!(
- props.dictionary_data_page_encoding(),
- Encoding::RLE_DICTIONARY
- );
- }
- }
-
- #[test]
- #[should_panic(expected = "Dictionary encoding can not be used as fallback encoding")]
- fn test_writer_properties_panic_when_plain_dictionary_is_fallback() {
- // Should panic when user specifies dictionary encoding as fallback encoding.
- WriterProperties::builder()
- .set_encoding(Encoding::PLAIN_DICTIONARY)
- .build();
- }
-
- #[test]
- #[should_panic(expected = "Dictionary encoding can not be used as fallback encoding")]
- fn test_writer_properties_panic_when_rle_dictionary_is_fallback() {
- // Should panic when user specifies dictionary encoding as fallback encoding.
- WriterProperties::builder()
- .set_encoding(Encoding::RLE_DICTIONARY)
- .build();
- }
-
- #[test]
- #[should_panic(expected = "Dictionary encoding can not be used as fallback encoding")]
- fn test_writer_properties_panic_when_dictionary_is_enabled() {
- WriterProperties::builder()
- .set_dictionary_enabled(true)
- .set_column_encoding(ColumnPath::from("col"), Encoding::RLE_DICTIONARY)
- .build();
- }
-
- #[test]
- #[should_panic(expected = "Dictionary encoding can not be used as fallback encoding")]
- fn test_writer_properties_panic_when_dictionary_is_disabled() {
- WriterProperties::builder()
- .set_dictionary_enabled(false)
- .set_column_encoding(ColumnPath::from("col"), Encoding::RLE_DICTIONARY)
- .build();
- }
-
- #[test]
- fn test_writer_properties_builder() {
- let props = WriterProperties::builder()
- // file settings
- .set_writer_version(WriterVersion::PARQUET_2_0)
- .set_data_pagesize_limit(10)
- .set_dictionary_pagesize_limit(20)
- .set_write_batch_size(30)
- .set_max_row_group_size(40)
- .set_created_by("default".to_owned())
- .set_key_value_metadata(Some(vec![KeyValue::new(
- "key".to_string(),
- "value".to_string(),
- )]))
- // global column settings
- .set_encoding(Encoding::DELTA_BINARY_PACKED)
- .set_compression(Compression::GZIP)
- .set_dictionary_enabled(false)
- .set_statistics_enabled(false)
- .set_max_statistics_size(50)
- // specific column settings
- .set_column_encoding(ColumnPath::from("col"), Encoding::RLE)
- .set_column_compression(ColumnPath::from("col"), Compression::SNAPPY)
- .set_column_dictionary_enabled(ColumnPath::from("col"), true)
- .set_column_statistics_enabled(ColumnPath::from("col"), true)
- .set_column_max_statistics_size(ColumnPath::from("col"), 123)
- .build();
-
- assert_eq!(props.writer_version(), WriterVersion::PARQUET_2_0);
- assert_eq!(props.data_pagesize_limit(), 10);
- assert_eq!(props.dictionary_pagesize_limit(), 20);
- assert_eq!(props.write_batch_size(), 30);
- assert_eq!(props.max_row_group_size(), 40);
- assert_eq!(props.created_by(), "default");
- assert_eq!(
- props.key_value_metadata(),
- &Some(vec![KeyValue::new("key".to_string(), "value".to_string(),)])
- );
-
- assert_eq!(
- props.encoding(&ColumnPath::from("a")),
- Some(Encoding::DELTA_BINARY_PACKED)
- );
- assert_eq!(props.compression(&ColumnPath::from("a")), Compression::GZIP);
- assert_eq!(props.dictionary_enabled(&ColumnPath::from("a")), false);
- assert_eq!(props.statistics_enabled(&ColumnPath::from("a")), false);
- assert_eq!(props.max_statistics_size(&ColumnPath::from("a")), 50);
-
- assert_eq!(
- props.encoding(&ColumnPath::from("col")),
- Some(Encoding::RLE)
- );
- assert_eq!(
- props.compression(&ColumnPath::from("col")),
- Compression::SNAPPY
- );
- assert_eq!(props.dictionary_enabled(&ColumnPath::from("col")), true);
- assert_eq!(props.statistics_enabled(&ColumnPath::from("col")), true);
- assert_eq!(props.max_statistics_size(&ColumnPath::from("col")), 123);
- }
-
- #[test]
- fn test_writer_properties_builder_partial_defaults() {
- let props = WriterProperties::builder()
- .set_encoding(Encoding::DELTA_BINARY_PACKED)
- .set_compression(Compression::GZIP)
- .set_column_encoding(ColumnPath::from("col"), Encoding::RLE)
- .build();
-
- assert_eq!(
- props.encoding(&ColumnPath::from("col")),
- Some(Encoding::RLE)
- );
- assert_eq!(
- props.compression(&ColumnPath::from("col")),
- Compression::GZIP
- );
- assert_eq!(
- props.dictionary_enabled(&ColumnPath::from("col")),
- DEFAULT_DICTIONARY_ENABLED
- );
- }
-}
diff --git a/parquet/src/file/reader.rs b/parquet/src/file/reader.rs
deleted file mode 100644
index aa8ba83..0000000
--- a/parquet/src/file/reader.rs
+++ /dev/null
@@ -1,206 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-//! Contains file reader API and provides methods to access file metadata, row group
-//! readers to read individual column chunks, or access record iterator.
-
-use std::{boxed::Box, io::Read, sync::Arc};
-
-use crate::column::page::PageIterator;
-use crate::column::{page::PageReader, reader::ColumnReader};
-use crate::errors::{ParquetError, Result};
-use crate::file::metadata::*;
-pub use crate::file::serialized_reader::{SerializedFileReader, SerializedPageReader};
-use crate::record::reader::RowIter;
-use crate::schema::types::{ColumnDescPtr, SchemaDescPtr, Type as SchemaType};
-
-use crate::basic::Type;
-
-use crate::column::reader::ColumnReaderImpl;
-
-/// Length should return the total number of bytes in the input source.
-/// It's mainly used to read the metadata, which is at the end of the source.
-#[allow(clippy::len_without_is_empty)]
-pub trait Length {
- /// Returns the amount of bytes of the inner source.
- fn len(&self) -> u64;
-}
-
-/// The ChunkReader trait generates readers of chunks of a source.
-/// For a file system reader, each chunk might contain a clone of File bounded on a given range.
-/// For an object store reader, each read can be mapped to a range request.
-pub trait ChunkReader: Length {
- type T: Read;
- /// get a serialy readeable slice of the current reader
- /// This should fail if the slice exceeds the current bounds
- fn get_read(&self, start: u64, length: usize) -> Result<Self::T>;
-}
-
-// ----------------------------------------------------------------------
-// APIs for file & row group readers
-
-/// Parquet file reader API. With this, user can get metadata information about the
-/// Parquet file, can get reader for each row group, and access record iterator.
-pub trait FileReader {
- /// Get metadata information about this file.
- fn metadata(&self) -> &ParquetMetaData;
-
- /// Get the total number of row groups for this file.
- fn num_row_groups(&self) -> usize;
-
- /// Get the `i`th row group reader. Note this doesn't do bound check.
- fn get_row_group(&self, i: usize) -> Result<Box<dyn RowGroupReader + '_>>;
-
- /// Get full iterator of `Row`s from a file (over all row groups).
- ///
- /// Iterator will automatically load the next row group to advance.
- ///
- /// Projected schema can be a subset of or equal to the file schema, when it is None,
- /// full file schema is assumed.
- fn get_row_iter(&self, projection: Option<SchemaType>) -> Result<RowIter>;
-}
-
-/// Parquet row group reader API. With this, user can get metadata information about the
-/// row group, as well as readers for each individual column chunk.
-pub trait RowGroupReader {
- /// Get metadata information about this row group.
- fn metadata(&self) -> &RowGroupMetaData;
-
- /// Get the total number of column chunks in this row group.
- fn num_columns(&self) -> usize;
-
- /// Get page reader for the `i`th column chunk.
- fn get_column_page_reader(&self, i: usize) -> Result<Box<dyn PageReader>>;
-
- /// Get value reader for the `i`th column chunk.
- fn get_column_reader(&self, i: usize) -> Result<ColumnReader> {
- let schema_descr = self.metadata().schema_descr();
- let col_descr = schema_descr.column(i);
- let col_page_reader = self.get_column_page_reader(i)?;
- let col_reader = match col_descr.physical_type() {
- Type::BOOLEAN => ColumnReader::BoolColumnReader(ColumnReaderImpl::new(
- col_descr,
- col_page_reader,
- )),
- Type::INT32 => ColumnReader::Int32ColumnReader(ColumnReaderImpl::new(
- col_descr,
- col_page_reader,
- )),
- Type::INT64 => ColumnReader::Int64ColumnReader(ColumnReaderImpl::new(
- col_descr,
- col_page_reader,
- )),
- Type::INT96 => ColumnReader::Int96ColumnReader(ColumnReaderImpl::new(
- col_descr,
- col_page_reader,
- )),
- Type::FLOAT => ColumnReader::FloatColumnReader(ColumnReaderImpl::new(
- col_descr,
- col_page_reader,
- )),
- Type::DOUBLE => ColumnReader::DoubleColumnReader(ColumnReaderImpl::new(
- col_descr,
- col_page_reader,
- )),
- Type::BYTE_ARRAY => ColumnReader::ByteArrayColumnReader(
- ColumnReaderImpl::new(col_descr, col_page_reader),
- ),
- Type::FIXED_LEN_BYTE_ARRAY => ColumnReader::FixedLenByteArrayColumnReader(
- ColumnReaderImpl::new(col_descr, col_page_reader),
- ),
- };
- Ok(col_reader)
- }
-
- /// Get iterator of `Row`s from this row group.
- ///
- /// Projected schema can be a subset of or equal to the file schema, when it is None,
- /// full file schema is assumed.
- fn get_row_iter(&self, projection: Option<SchemaType>) -> Result<RowIter>;
-}
-
-// ----------------------------------------------------------------------
-// Iterator
-
-/// Implementation of page iterator for parquet file.
-pub struct FilePageIterator {
- column_index: usize,
- row_group_indices: Box<dyn Iterator<Item = usize>>,
- file_reader: Arc<dyn FileReader>,
-}
-
-impl FilePageIterator {
- /// Creates a page iterator for all row groups in file.
- pub fn new(column_index: usize, file_reader: Arc<dyn FileReader>) -> Result<Self> {
- let num_row_groups = file_reader.metadata().num_row_groups();
-
- let row_group_indices = Box::new(0..num_row_groups);
-
- Self::with_row_groups(column_index, row_group_indices, file_reader)
- }
-
- /// Create page iterator from parquet file reader with only some row groups.
- pub fn with_row_groups(
- column_index: usize,
- row_group_indices: Box<dyn Iterator<Item = usize>>,
- file_reader: Arc<dyn FileReader>,
- ) -> Result<Self> {
- // Check that column_index is valid
- let num_columns = file_reader
- .metadata()
- .file_metadata()
- .schema_descr()
- .num_columns();
-
- if column_index >= num_columns {
- return Err(ParquetError::IndexOutOfBound(column_index, num_columns));
- }
-
- // We don't check iterators here because iterator may be infinite
- Ok(Self {
- column_index,
- row_group_indices,
- file_reader,
- })
- }
-}
-
-impl Iterator for FilePageIterator {
- type Item = Result<Box<dyn PageReader>>;
-
- fn next(&mut self) -> Option<Result<Box<dyn PageReader>>> {
- self.row_group_indices.next().map(|row_group_index| {
- self.file_reader
- .get_row_group(row_group_index)
- .and_then(|r| r.get_column_page_reader(self.column_index))
- })
- }
-}
-
-impl PageIterator for FilePageIterator {
- fn schema(&mut self) -> Result<SchemaDescPtr> {
- Ok(self
- .file_reader
- .metadata()
- .file_metadata()
- .schema_descr_ptr())
- }
-
- fn column_schema(&mut self) -> Result<ColumnDescPtr> {
- self.schema().map(|s| s.column(self.column_index))
- }
-}
diff --git a/parquet/src/file/serialized_reader.rs b/parquet/src/file/serialized_reader.rs
deleted file mode 100644
index a4d79a3..0000000
--- a/parquet/src/file/serialized_reader.rs
+++ /dev/null
@@ -1,779 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-//! Contains implementations of the reader traits FileReader, RowGroupReader and PageReader
-//! Also contains implementations of the ChunkReader for files (with buffering) and byte arrays (RAM)
-
-use std::{convert::TryFrom, fs::File, io::Read, path::Path, sync::Arc};
-
-use parquet_format::{PageHeader, PageType};
-use thrift::protocol::TCompactInputProtocol;
-
-use crate::basic::{Compression, Encoding, Type};
-use crate::column::page::{Page, PageReader};
-use crate::compression::{create_codec, Codec};
-use crate::errors::{ParquetError, Result};
-use crate::file::{footer, metadata::*, reader::*, statistics};
-use crate::record::reader::RowIter;
-use crate::record::Row;
-use crate::schema::types::Type as SchemaType;
-use crate::util::{io::TryClone, memory::ByteBufferPtr};
-
-// export `SliceableCursor` and `FileSource` publically so clients can
-// re-use the logic in their own ParquetFileWriter wrappers
-pub use crate::util::{cursor::SliceableCursor, io::FileSource};
-
-// ----------------------------------------------------------------------
-// Implementations of traits facilitating the creation of a new reader
-
-impl Length for File {
- fn len(&self) -> u64 {
- self.metadata().map(|m| m.len()).unwrap_or(0u64)
- }
-}
-
-impl TryClone for File {
- fn try_clone(&self) -> std::io::Result<Self> {
- self.try_clone()
- }
-}
-
-impl ChunkReader for File {
- type T = FileSource<File>;
-
- fn get_read(&self, start: u64, length: usize) -> Result<Self::T> {
- Ok(FileSource::new(self, start, length))
- }
-}
-
-impl Length for SliceableCursor {
- fn len(&self) -> u64 {
- SliceableCursor::len(self)
- }
-}
-
-impl ChunkReader for SliceableCursor {
- type T = SliceableCursor;
-
- fn get_read(&self, start: u64, length: usize) -> Result<Self::T> {
- self.slice(start, length).map_err(|e| e.into())
- }
-}
-
-impl TryFrom<File> for SerializedFileReader<File> {
- type Error = ParquetError;
-
- fn try_from(file: File) -> Result<Self> {
- Self::new(file)
- }
-}
-
-impl<'a> TryFrom<&'a Path> for SerializedFileReader<File> {
- type Error = ParquetError;
-
- fn try_from(path: &Path) -> Result<Self> {
- let file = File::open(path)?;
- Self::try_from(file)
- }
-}
-
-impl TryFrom<String> for SerializedFileReader<File> {
- type Error = ParquetError;
-
- fn try_from(path: String) -> Result<Self> {
- Self::try_from(Path::new(&path))
- }
-}
-
-impl<'a> TryFrom<&'a str> for SerializedFileReader<File> {
- type Error = ParquetError;
-
- fn try_from(path: &str) -> Result<Self> {
- Self::try_from(Path::new(&path))
- }
-}
-
-/// Conversion into a [`RowIter`](crate::record::reader::RowIter)
-/// using the full file schema over all row groups.
-impl IntoIterator for SerializedFileReader<File> {
- type Item = Row;
- type IntoIter = RowIter<'static>;
-
- fn into_iter(self) -> Self::IntoIter {
- RowIter::from_file_into(Box::new(self))
- }
-}
-
-// ----------------------------------------------------------------------
-// Implementations of file & row group readers
-
-/// A serialized implementation for Parquet [`FileReader`].
-pub struct SerializedFileReader<R: ChunkReader> {
- chunk_reader: Arc<R>,
- metadata: ParquetMetaData,
-}
-
-impl<R: 'static + ChunkReader> SerializedFileReader<R> {
- /// Creates file reader from a Parquet file.
- /// Returns error if Parquet file does not exist or is corrupt.
- pub fn new(chunk_reader: R) -> Result<Self> {
- let metadata = footer::parse_metadata(&chunk_reader)?;
- Ok(Self {
- chunk_reader: Arc::new(chunk_reader),
- metadata,
- })
- }
-
- /// Filters row group metadata to only those row groups,
- /// for which the predicate function returns true
- pub fn filter_row_groups(
- &mut self,
- predicate: &dyn Fn(&RowGroupMetaData, usize) -> bool,
- ) {
- let mut filtered_row_groups = Vec::<RowGroupMetaData>::new();
- for (i, row_group_metadata) in self.metadata.row_groups().iter().enumerate() {
- if predicate(row_group_metadata, i) {
- filtered_row_groups.push(row_group_metadata.clone());
- }
- }
- self.metadata = ParquetMetaData::new(
- self.metadata.file_metadata().clone(),
- filtered_row_groups,
- );
- }
-}
-
-impl<R: 'static + ChunkReader> FileReader for SerializedFileReader<R> {
- fn metadata(&self) -> &ParquetMetaData {
- &self.metadata
- }
-
- fn num_row_groups(&self) -> usize {
- self.metadata.num_row_groups()
- }
-
- fn get_row_group(&self, i: usize) -> Result<Box<dyn RowGroupReader + '_>> {
- let row_group_metadata = self.metadata.row_group(i);
- // Row groups should be processed sequentially.
- let f = Arc::clone(&self.chunk_reader);
- Ok(Box::new(SerializedRowGroupReader::new(
- f,
- row_group_metadata,
- )))
- }
-
- fn get_row_iter(&self, projection: Option<SchemaType>) -> Result<RowIter> {
- RowIter::from_file(projection, self)
- }
-}
-
-/// A serialized implementation for Parquet [`RowGroupReader`].
-pub struct SerializedRowGroupReader<'a, R: ChunkReader> {
- chunk_reader: Arc<R>,
- metadata: &'a RowGroupMetaData,
-}
-
-impl<'a, R: ChunkReader> SerializedRowGroupReader<'a, R> {
- /// Creates new row group reader from a file and row group metadata.
- fn new(chunk_reader: Arc<R>, metadata: &'a RowGroupMetaData) -> Self {
- Self {
- chunk_reader,
- metadata,
- }
- }
-}
-
-impl<'a, R: 'static + ChunkReader> RowGroupReader for SerializedRowGroupReader<'a, R> {
- fn metadata(&self) -> &RowGroupMetaData {
- &self.metadata
- }
-
- fn num_columns(&self) -> usize {
- self.metadata.num_columns()
- }
-
- // TODO: fix PARQUET-816
- fn get_column_page_reader(&self, i: usize) -> Result<Box<dyn PageReader>> {
- let col = self.metadata.column(i);
- let (col_start, col_length) = col.byte_range();
- let file_chunk = self.chunk_reader.get_read(col_start, col_length as usize)?;
- let page_reader = SerializedPageReader::new(
- file_chunk,
- col.num_values(),
- col.compression(),
- col.column_descr().physical_type(),
- )?;
- Ok(Box::new(page_reader))
- }
-
- fn get_row_iter(&self, projection: Option<SchemaType>) -> Result<RowIter> {
- RowIter::from_row_group(projection, self)
- }
-}
-
-/// A serialized implementation for Parquet [`PageReader`].
-pub struct SerializedPageReader<T: Read> {
- // The file source buffer which references exactly the bytes for the column trunk
- // to be read by this page reader.
- buf: T,
-
- // The compression codec for this column chunk. Only set for non-PLAIN codec.
- decompressor: Option<Box<dyn Codec>>,
-
- // The number of values we have seen so far.
- seen_num_values: i64,
-
- // The number of total values in this column chunk.
- total_num_values: i64,
-
- // Column chunk type.
- physical_type: Type,
-}
-
-impl<T: Read> SerializedPageReader<T> {
- /// Creates a new serialized page reader from file source.
- pub fn new(
- buf: T,
- total_num_values: i64,
- compression: Compression,
- physical_type: Type,
- ) -> Result<Self> {
- let decompressor = create_codec(compression)?;
- let result = Self {
- buf,
- total_num_values,
- seen_num_values: 0,
- decompressor,
- physical_type,
- };
- Ok(result)
- }
-
- /// Reads Page header from Thrift.
- fn read_page_header(&mut self) -> Result<PageHeader> {
- let mut prot = TCompactInputProtocol::new(&mut self.buf);
- let page_header = PageHeader::read_from_in_protocol(&mut prot)?;
- Ok(page_header)
- }
-}
-
-impl<T: Read> Iterator for SerializedPageReader<T> {
- type Item = Result<Page>;
-
- fn next(&mut self) -> Option<Self::Item> {
- self.get_next_page().transpose()
- }
-}
-
-impl<T: Read> PageReader for SerializedPageReader<T> {
- fn get_next_page(&mut self) -> Result<Option<Page>> {
- while self.seen_num_values < self.total_num_values {
- let page_header = self.read_page_header()?;
-
- // When processing data page v2, depending on enabled compression for the
- // page, we should account for uncompressed data ('offset') of
- // repetition and definition levels.
- //
- // We always use 0 offset for other pages other than v2, `true` flag means
- // that compression will be applied if decompressor is defined
- let mut offset: usize = 0;
- let mut can_decompress = true;
-
- if let Some(ref header_v2) = page_header.data_page_header_v2 {
- offset = (header_v2.definition_levels_byte_length
- + header_v2.repetition_levels_byte_length)
- as usize;
- // When is_compressed flag is missing the page is considered compressed
- can_decompress = header_v2.is_compressed.unwrap_or(true);
- }
-
- let compressed_len = page_header.compressed_page_size as usize - offset;
- let uncompressed_len = page_header.uncompressed_page_size as usize - offset;
- // We still need to read all bytes from buffered stream
- let mut buffer = vec![0; offset + compressed_len];
- self.buf.read_exact(&mut buffer)?;
-
- // TODO: page header could be huge because of statistics. We should set a
- // maximum page header size and abort if that is exceeded.
- if let Some(decompressor) = self.decompressor.as_mut() {
- if can_decompress {
- let mut decompressed_buffer = Vec::with_capacity(uncompressed_len);
- let decompressed_size = decompressor
- .decompress(&buffer[offset..], &mut decompressed_buffer)?;
- if decompressed_size != uncompressed_len {
- return Err(general_err!(
- "Actual decompressed size doesn't match the expected one ({} vs {})",
- decompressed_size,
- uncompressed_len
- ));
- }
- if offset == 0 {
- buffer = decompressed_buffer;
- } else {
- // Prepend saved offsets to the buffer
- buffer.truncate(offset);
- buffer.append(&mut decompressed_buffer);
- }
- }
- }
-
- let result = match page_header.type_ {
- PageType::DictionaryPage => {
- assert!(page_header.dictionary_page_header.is_some());
- let dict_header =
- page_header.dictionary_page_header.as_ref().unwrap();
- let is_sorted = dict_header.is_sorted.unwrap_or(false);
- Page::DictionaryPage {
- buf: ByteBufferPtr::new(buffer),
- num_values: dict_header.num_values as u32,
- encoding: Encoding::from(dict_header.encoding),
- is_sorted,
- }
- }
- PageType::DataPage => {
- assert!(page_header.data_page_header.is_some());
- let header = page_header.data_page_header.unwrap();
- self.seen_num_values += header.num_values as i64;
- Page::DataPage {
- buf: ByteBufferPtr::new(buffer),
- num_values: header.num_values as u32,
- encoding: Encoding::from(header.encoding),
- def_level_encoding: Encoding::from(
- header.definition_level_encoding,
- ),
- rep_level_encoding: Encoding::from(
- header.repetition_level_encoding,
- ),
- statistics: statistics::from_thrift(
- self.physical_type,
- header.statistics,
- ),
- }
- }
- PageType::DataPageV2 => {
- assert!(page_header.data_page_header_v2.is_some());
- let header = page_header.data_page_header_v2.unwrap();
- let is_compressed = header.is_compressed.unwrap_or(true);
- self.seen_num_values += header.num_values as i64;
- Page::DataPageV2 {
- buf: ByteBufferPtr::new(buffer),
- num_values: header.num_values as u32,
- encoding: Encoding::from(header.encoding),
- num_nulls: header.num_nulls as u32,
- num_rows: header.num_rows as u32,
- def_levels_byte_len: header.definition_levels_byte_length as u32,
- rep_levels_byte_len: header.repetition_levels_byte_length as u32,
- is_compressed,
- statistics: statistics::from_thrift(
- self.physical_type,
- header.statistics,
- ),
- }
- }
- _ => {
- // For unknown page type (e.g., INDEX_PAGE), skip and read next.
- continue;
- }
- };
- return Ok(Some(result));
- }
-
- // We are at the end of this column chunk and no more page left. Return None.
- Ok(None)
- }
-}
-
-#[cfg(test)]
-mod tests {
- use super::*;
- use crate::basic::ColumnOrder;
- use crate::record::RowAccessor;
- use crate::schema::parser::parse_message_type;
- use crate::util::test_common::{get_test_file, get_test_path};
- use std::sync::Arc;
-
- #[test]
- fn test_cursor_and_file_has_the_same_behaviour() {
- let mut buf: Vec<u8> = Vec::new();
- get_test_file("alltypes_plain.parquet")
- .read_to_end(&mut buf)
- .unwrap();
- let cursor = SliceableCursor::new(buf);
- let read_from_cursor = SerializedFileReader::new(cursor).unwrap();
-
- let test_file = get_test_file("alltypes_plain.parquet");
- let read_from_file = SerializedFileReader::new(test_file).unwrap();
-
- let file_iter = read_from_file.get_row_iter(None).unwrap();
- let cursor_iter = read_from_cursor.get_row_iter(None).unwrap();
-
- assert!(file_iter.eq(cursor_iter));
- }
-
- #[test]
- fn test_file_reader_try_from() {
- // Valid file path
- let test_file = get_test_file("alltypes_plain.parquet");
- let test_path_buf = get_test_path("alltypes_plain.parquet");
- let test_path = test_path_buf.as_path();
- let test_path_str = test_path.to_str().unwrap();
-
- let reader = SerializedFileReader::try_from(test_file);
- assert!(reader.is_ok());
-
- let reader = SerializedFileReader::try_from(test_path);
- assert!(reader.is_ok());
-
- let reader = SerializedFileReader::try_from(test_path_str);
- assert!(reader.is_ok());
-
- let reader = SerializedFileReader::try_from(test_path_str.to_string());
- assert!(reader.is_ok());
-
- // Invalid file path
- let test_path = Path::new("invalid.parquet");
- let test_path_str = test_path.to_str().unwrap();
-
- let reader = SerializedFileReader::try_from(test_path);
- assert!(reader.is_err());
-
- let reader = SerializedFileReader::try_from(test_path_str);
- assert!(reader.is_err());
-
- let reader = SerializedFileReader::try_from(test_path_str.to_string());
- assert!(reader.is_err());
- }
-
- #[test]
- fn test_file_reader_into_iter() {
- let path = get_test_path("alltypes_plain.parquet");
- let vec = vec![path.clone(), path]
- .iter()
- .map(|p| SerializedFileReader::try_from(p.as_path()).unwrap())
- .flat_map(|r| r.into_iter())
- .flat_map(|r| r.get_int(0))
- .collect::<Vec<_>>();
-
- // rows in the parquet file are not sorted by "id"
- // each file contains [id:4, id:5, id:6, id:7, id:2, id:3, id:0, id:1]
- assert_eq!(vec, vec![4, 5, 6, 7, 2, 3, 0, 1, 4, 5, 6, 7, 2, 3, 0, 1]);
- }
-
- #[test]
- fn test_file_reader_into_iter_project() {
- let path = get_test_path("alltypes_plain.parquet");
- let result = vec![path]
- .iter()
- .map(|p| SerializedFileReader::try_from(p.as_path()).unwrap())
- .flat_map(|r| {
- let schema = "message schema { OPTIONAL INT32 id; }";
- let proj = parse_message_type(&schema).ok();
-
- r.into_iter().project(proj).unwrap()
- })
- .map(|r| format!("{}", r))
- .collect::<Vec<_>>()
- .join(",");
-
- assert_eq!(
- result,
- "{id: 4},{id: 5},{id: 6},{id: 7},{id: 2},{id: 3},{id: 0},{id: 1}"
- );
- }
-
- #[test]
- fn test_reuse_file_chunk() {
- // This test covers the case of maintaining the correct start position in a file
- // stream for each column reader after initializing and moving to the next one
- // (without necessarily reading the entire column).
- let test_file = get_test_file("alltypes_plain.parquet");
- let reader = SerializedFileReader::new(test_file).unwrap();
- let row_group = reader.get_row_group(0).unwrap();
-
- let mut page_readers = Vec::new();
- for i in 0..row_group.num_columns() {
- page_readers.push(row_group.get_column_page_reader(i).unwrap());
- }
-
- // Now buffer each col reader, we do not expect any failures like:
- // General("underlying Thrift error: end of file")
- for mut page_reader in page_readers {
- assert!(page_reader.get_next_page().is_ok());
- }
- }
-
- #[test]
- fn test_file_reader() {
- let test_file = get_test_file("alltypes_plain.parquet");
- let reader_result = SerializedFileReader::new(test_file);
- assert!(reader_result.is_ok());
- let reader = reader_result.unwrap();
-
- // Test contents in Parquet metadata
- let metadata = reader.metadata();
- assert_eq!(metadata.num_row_groups(), 1);
-
- // Test contents in file metadata
- let file_metadata = metadata.file_metadata();
- assert!(file_metadata.created_by().is_some());
- assert_eq!(
- file_metadata.created_by().as_ref().unwrap(),
- "impala version 1.3.0-INTERNAL (build 8a48ddb1eff84592b3fc06bc6f51ec120e1fffc9)"
- );
- assert!(file_metadata.key_value_metadata().is_none());
- assert_eq!(file_metadata.num_rows(), 8);
- assert_eq!(file_metadata.version(), 1);
- assert_eq!(file_metadata.column_orders(), None);
-
- // Test contents in row group metadata
- let row_group_metadata = metadata.row_group(0);
- assert_eq!(row_group_metadata.num_columns(), 11);
- assert_eq!(row_group_metadata.num_rows(), 8);
- assert_eq!(row_group_metadata.total_byte_size(), 671);
- // Check each column order
- for i in 0..row_group_metadata.num_columns() {
- assert_eq!(file_metadata.column_order(i), ColumnOrder::UNDEFINED);
- }
-
- // Test row group reader
- let row_group_reader_result = reader.get_row_group(0);
- assert!(row_group_reader_result.is_ok());
- let row_group_reader: Box<dyn RowGroupReader> = row_group_reader_result.unwrap();
- assert_eq!(
- row_group_reader.num_columns(),
- row_group_metadata.num_columns()
- );
- assert_eq!(
- row_group_reader.metadata().total_byte_size(),
- row_group_metadata.total_byte_size()
- );
-
- // Test page readers
- // TODO: test for every column
- let page_reader_0_result = row_group_reader.get_column_page_reader(0);
- assert!(page_reader_0_result.is_ok());
- let mut page_reader_0: Box<dyn PageReader> = page_reader_0_result.unwrap();
- let mut page_count = 0;
- while let Ok(Some(page)) = page_reader_0.get_next_page() {
- let is_expected_page = match page {
- Page::DictionaryPage {
- buf,
- num_values,
- encoding,
- is_sorted,
- } => {
- assert_eq!(buf.len(), 32);
- assert_eq!(num_values, 8);
- assert_eq!(encoding, Encoding::PLAIN_DICTIONARY);
- assert_eq!(is_sorted, false);
- true
- }
- Page::DataPage {
- buf,
- num_values,
- encoding,
- def_level_encoding,
- rep_level_encoding,
- statistics,
- } => {
- assert_eq!(buf.len(), 11);
- assert_eq!(num_values, 8);
- assert_eq!(encoding, Encoding::PLAIN_DICTIONARY);
- assert_eq!(def_level_encoding, Encoding::RLE);
- assert_eq!(rep_level_encoding, Encoding::BIT_PACKED);
- assert!(statistics.is_none());
- true
- }
- _ => false,
- };
- assert!(is_expected_page);
- page_count += 1;
- }
- assert_eq!(page_count, 2);
- }
-
- #[test]
- fn test_file_reader_datapage_v2() {
- let test_file = get_test_file("datapage_v2.snappy.parquet");
- let reader_result = SerializedFileReader::new(test_file);
- assert!(reader_result.is_ok());
- let reader = reader_result.unwrap();
-
- // Test contents in Parquet metadata
- let metadata = reader.metadata();
- assert_eq!(metadata.num_row_groups(), 1);
-
- // Test contents in file metadata
- let file_metadata = metadata.file_metadata();
- assert!(file_metadata.created_by().is_some());
- assert_eq!(
- file_metadata.created_by().as_ref().unwrap(),
- "parquet-mr version 1.8.1 (build 4aba4dae7bb0d4edbcf7923ae1339f28fd3f7fcf)"
- );
- assert!(file_metadata.key_value_metadata().is_some());
- assert_eq!(
- file_metadata.key_value_metadata().to_owned().unwrap().len(),
- 1
- );
-
- assert_eq!(file_metadata.num_rows(), 5);
- assert_eq!(file_metadata.version(), 1);
- assert_eq!(file_metadata.column_orders(), None);
-
- let row_group_metadata = metadata.row_group(0);
-
- // Check each column order
- for i in 0..row_group_metadata.num_columns() {
- assert_eq!(file_metadata.column_order(i), ColumnOrder::UNDEFINED);
- }
-
- // Test row group reader
- let row_group_reader_result = reader.get_row_group(0);
- assert!(row_group_reader_result.is_ok());
- let row_group_reader: Box<dyn RowGroupReader> = row_group_reader_result.unwrap();
- assert_eq!(
- row_group_reader.num_columns(),
- row_group_metadata.num_columns()
- );
- assert_eq!(
- row_group_reader.metadata().total_byte_size(),
- row_group_metadata.total_byte_size()
- );
-
- // Test page readers
- // TODO: test for every column
- let page_reader_0_result = row_group_reader.get_column_page_reader(0);
- assert!(page_reader_0_result.is_ok());
- let mut page_reader_0: Box<dyn PageReader> = page_reader_0_result.unwrap();
- let mut page_count = 0;
- while let Ok(Some(page)) = page_reader_0.get_next_page() {
- let is_expected_page = match page {
- Page::DictionaryPage {
- buf,
- num_values,
- encoding,
- is_sorted,
- } => {
- assert_eq!(buf.len(), 7);
- assert_eq!(num_values, 1);
- assert_eq!(encoding, Encoding::PLAIN);
- assert_eq!(is_sorted, false);
- true
- }
- Page::DataPageV2 {
- buf,
- num_values,
- encoding,
- num_nulls,
- num_rows,
- def_levels_byte_len,
- rep_levels_byte_len,
- is_compressed,
- statistics,
- } => {
- assert_eq!(buf.len(), 4);
- assert_eq!(num_values, 5);
- assert_eq!(encoding, Encoding::RLE_DICTIONARY);
- assert_eq!(num_nulls, 1);
- assert_eq!(num_rows, 5);
- assert_eq!(def_levels_byte_len, 2);
- assert_eq!(rep_levels_byte_len, 0);
- assert_eq!(is_compressed, true);
- assert!(statistics.is_some());
- true
- }
- _ => false,
- };
- assert!(is_expected_page);
- page_count += 1;
- }
- assert_eq!(page_count, 2);
- }
-
- #[test]
- fn test_page_iterator() {
- let file = get_test_file("alltypes_plain.parquet");
- let file_reader = Arc::new(SerializedFileReader::new(file).unwrap());
-
- let mut page_iterator = FilePageIterator::new(0, file_reader.clone()).unwrap();
-
- // read first page
- let page = page_iterator.next();
- assert!(page.is_some());
- assert!(page.unwrap().is_ok());
-
- // reach end of file
- let page = page_iterator.next();
- assert!(page.is_none());
-
- let row_group_indices = Box::new(0..1);
- let mut page_iterator =
- FilePageIterator::with_row_groups(0, row_group_indices, file_reader).unwrap();
-
- // read first page
- let page = page_iterator.next();
- assert!(page.is_some());
- assert!(page.unwrap().is_ok());
-
- // reach end of file
- let page = page_iterator.next();
- assert!(page.is_none());
- }
-
- #[test]
- fn test_file_reader_key_value_metadata() {
- let file = get_test_file("binary.parquet");
- let file_reader = Arc::new(SerializedFileReader::new(file).unwrap());
-
- let metadata = file_reader
- .metadata
- .file_metadata()
- .key_value_metadata()
- .as_ref()
- .unwrap();
-
- assert_eq!(metadata.len(), 3);
-
- assert_eq!(metadata.get(0).unwrap().key, "parquet.proto.descriptor");
-
- assert_eq!(metadata.get(1).unwrap().key, "writer.model.name");
- assert_eq!(metadata.get(1).unwrap().value, Some("protobuf".to_owned()));
-
- assert_eq!(metadata.get(2).unwrap().key, "parquet.proto.class");
- assert_eq!(
- metadata.get(2).unwrap().value,
- Some("foo.baz.Foobaz$Event".to_owned())
- );
- }
-
- #[test]
- fn test_file_reader_filter_row_groups() -> Result<()> {
- let test_file = get_test_file("alltypes_plain.parquet");
- let mut reader = SerializedFileReader::new(test_file)?;
-
- // test initial number of row groups
- let metadata = reader.metadata();
- assert_eq!(metadata.num_row_groups(), 1);
-
- // test filtering out all row groups
- reader.filter_row_groups(&|_, _| false);
- let metadata = reader.metadata();
- assert_eq!(metadata.num_row_groups(), 0);
-
- Ok(())
- }
-}
diff --git a/parquet/src/file/statistics.rs b/parquet/src/file/statistics.rs
deleted file mode 100644
index 4f5d0e9..0000000
--- a/parquet/src/file/statistics.rs
+++ /dev/null
@@ -1,664 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-//! Contains definitions for working with Parquet statistics.
-//!
-//! Though some common methods are available on enum, use pattern match to extract
-//! actual min and max values from statistics, see below:
-//!
-//! ```rust
-//! use parquet::file::statistics::Statistics;
-//!
-//! let stats = Statistics::int32(Some(1), Some(10), None, 3, true);
-//! assert_eq!(stats.null_count(), 3);
-//! assert!(stats.has_min_max_set());
-//! assert!(stats.is_min_max_deprecated());
-//!
-//! match stats {
-//! Statistics::Int32(ref typed) => {
-//! assert_eq!(*typed.min(), 1);
-//! assert_eq!(*typed.max(), 10);
-//! }
-//! _ => {}
-//! }
-//! ```
-
-use std::{cmp, fmt};
-
-use byteorder::{ByteOrder, LittleEndian};
-use parquet_format::Statistics as TStatistics;
-
-use crate::basic::Type;
-use crate::data_type::*;
-use crate::util::bit_util::from_ne_slice;
-
-// Macro to generate methods create Statistics.
-macro_rules! statistics_new_func {
- ($func:ident, $vtype:ty, $stat:ident) => {
- pub fn $func(
- min: $vtype,
- max: $vtype,
- distinct: Option<u64>,
- nulls: u64,
- is_deprecated: bool,
- ) -> Self {
- Statistics::$stat(TypedStatistics::new(
- min,
- max,
- distinct,
- nulls,
- is_deprecated,
- ))
- }
- };
-}
-
-// Macro to generate getter functions for Statistics.
-macro_rules! statistics_enum_func {
- ($self:ident, $func:ident) => {{
- match *$self {
- Statistics::Boolean(ref typed) => typed.$func(),
- Statistics::Int32(ref typed) => typed.$func(),
- Statistics::Int64(ref typed) => typed.$func(),
- Statistics::Int96(ref typed) => typed.$func(),
- Statistics::Float(ref typed) => typed.$func(),
- Statistics::Double(ref typed) => typed.$func(),
- Statistics::ByteArray(ref typed) => typed.$func(),
- Statistics::FixedLenByteArray(ref typed) => typed.$func(),
- }
- }};
-}
-
-/// Converts Thrift definition into `Statistics`.
-pub fn from_thrift(
- physical_type: Type,
- thrift_stats: Option<TStatistics>,
-) -> Option<Statistics> {
- match thrift_stats {
- Some(stats) => {
- // Number of nulls recorded, when it is not available, we just mark it as 0.
- let null_count = stats.null_count.unwrap_or(0);
- assert!(
- null_count >= 0,
- "Statistics null count is negative ({})",
- null_count
- );
-
- // Generic null count.
- let null_count = null_count as u64;
- // Generic distinct count (count of distinct values occurring)
- let distinct_count = stats.distinct_count.map(|value| value as u64);
- // Whether or not statistics use deprecated min/max fields.
- let old_format = stats.min_value.is_none() && stats.max_value.is_none();
- // Generic min value as bytes.
- let min = if old_format {
- stats.min
- } else {
- stats.min_value
- };
- // Generic max value as bytes.
- let max = if old_format {
- stats.max
- } else {
- stats.max_value
- };
-
- // Values are encoded using PLAIN encoding definition, except that
- // variable-length byte arrays do not include a length prefix.
- //
- // Instead of using actual decoder, we manually convert values.
- let res = match physical_type {
- Type::BOOLEAN => Statistics::boolean(
- min.map(|data| data[0] != 0),
- max.map(|data| data[0] != 0),
- distinct_count,
- null_count,
- old_format,
- ),
- Type::INT32 => Statistics::int32(
- min.map(|data| LittleEndian::read_i32(&data)),
- max.map(|data| LittleEndian::read_i32(&data)),
- distinct_count,
- null_count,
- old_format,
- ),
- Type::INT64 => Statistics::int64(
- min.map(|data| LittleEndian::read_i64(&data)),
- max.map(|data| LittleEndian::read_i64(&data)),
- distinct_count,
- null_count,
- old_format,
- ),
- Type::INT96 => {
- // INT96 statistics may not be correct, because comparison is signed
- // byte-wise, not actual timestamps. It is recommended to ignore
- // min/max statistics for INT96 columns.
- let min = min.map(|data| {
- assert_eq!(data.len(), 12);
- from_ne_slice::<Int96>(&data)
- });
- let max = max.map(|data| {
- assert_eq!(data.len(), 12);
- from_ne_slice::<Int96>(&data)
- });
- Statistics::int96(min, max, distinct_count, null_count, old_format)
- }
- Type::FLOAT => Statistics::float(
- min.map(|data| LittleEndian::read_f32(&data)),
- max.map(|data| LittleEndian::read_f32(&data)),
- distinct_count,
- null_count,
- old_format,
- ),
- Type::DOUBLE => Statistics::double(
- min.map(|data| LittleEndian::read_f64(&data)),
- max.map(|data| LittleEndian::read_f64(&data)),
- distinct_count,
- null_count,
- old_format,
- ),
- Type::BYTE_ARRAY => Statistics::byte_array(
- min.map(ByteArray::from),
- max.map(ByteArray::from),
- distinct_count,
- null_count,
- old_format,
- ),
- Type::FIXED_LEN_BYTE_ARRAY => Statistics::fixed_len_byte_array(
- min.map(ByteArray::from).map(FixedLenByteArray::from),
- max.map(ByteArray::from).map(FixedLenByteArray::from),
- distinct_count,
- null_count,
- old_format,
- ),
- };
-
- Some(res)
- }
- None => None,
- }
-}
-
-// Convert Statistics into Thrift definition.
-pub fn to_thrift(stats: Option<&Statistics>) -> Option<TStatistics> {
- let stats = stats?;
-
- let mut thrift_stats = TStatistics {
- max: None,
- min: None,
- null_count: if stats.has_nulls() {
- Some(stats.null_count() as i64)
- } else {
- None
- },
- distinct_count: stats.distinct_count().map(|value| value as i64),
- max_value: None,
- min_value: None,
- };
-
- // Get min/max if set.
- let (min, max) = if stats.has_min_max_set() {
- (
- Some(stats.min_bytes().to_vec()),
- Some(stats.max_bytes().to_vec()),
- )
- } else {
- (None, None)
- };
-
- if stats.is_min_max_deprecated() {
- thrift_stats.min = min;
- thrift_stats.max = max;
- } else {
- thrift_stats.min_value = min;
- thrift_stats.max_value = max;
- }
-
- Some(thrift_stats)
-}
-
-/// Statistics for a column chunk and data page.
-#[derive(Debug, Clone, PartialEq)]
-pub enum Statistics {
- Boolean(TypedStatistics<BoolType>),
- Int32(TypedStatistics<Int32Type>),
- Int64(TypedStatistics<Int64Type>),
- Int96(TypedStatistics<Int96Type>),
- Float(TypedStatistics<FloatType>),
- Double(TypedStatistics<DoubleType>),
- ByteArray(TypedStatistics<ByteArrayType>),
- FixedLenByteArray(TypedStatistics<FixedLenByteArrayType>),
-}
-
-impl Statistics {
- statistics_new_func![boolean, Option<bool>, Boolean];
-
- statistics_new_func![int32, Option<i32>, Int32];
-
- statistics_new_func![int64, Option<i64>, Int64];
-
- statistics_new_func![int96, Option<Int96>, Int96];
-
- statistics_new_func![float, Option<f32>, Float];
-
- statistics_new_func![double, Option<f64>, Double];
-
- statistics_new_func![byte_array, Option<ByteArray>, ByteArray];
-
- statistics_new_func![
- fixed_len_byte_array,
- Option<FixedLenByteArray>,
- FixedLenByteArray
- ];
-
- /// Returns `true` if statistics have old `min` and `max` fields set.
- /// This means that the column order is likely to be undefined, which, for old files
- /// could mean a signed sort order of values.
- ///
- /// Refer to [`ColumnOrder`](crate::basic::ColumnOrder) and
- /// [`SortOrder`](crate::basic::SortOrder) for more information.
- pub fn is_min_max_deprecated(&self) -> bool {
- statistics_enum_func![self, is_min_max_deprecated]
- }
-
- /// Returns optional value of number of distinct values occurring.
- /// When it is `None`, the value should be ignored.
- pub fn distinct_count(&self) -> Option<u64> {
- statistics_enum_func![self, distinct_count]
- }
-
- /// Returns number of null values for the column.
- /// Note that this includes all nulls when column is part of the complex type.
- pub fn null_count(&self) -> u64 {
- statistics_enum_func![self, null_count]
- }
-
- /// Returns `true` if statistics collected any null values, `false` otherwise.
- pub fn has_nulls(&self) -> bool {
- self.null_count() > 0
- }
-
- /// Returns `true` if min value and max value are set.
- /// Normally both min/max values will be set to `Some(value)` or `None`.
- pub fn has_min_max_set(&self) -> bool {
- statistics_enum_func![self, has_min_max_set]
- }
-
- /// Returns slice of bytes that represent min value.
- /// Panics if min value is not set.
- pub fn min_bytes(&self) -> &[u8] {
- statistics_enum_func![self, min_bytes]
- }
-
- /// Returns slice of bytes that represent max value.
- /// Panics if max value is not set.
- pub fn max_bytes(&self) -> &[u8] {
- statistics_enum_func![self, max_bytes]
- }
-
- /// Returns physical type associated with statistics.
- pub fn physical_type(&self) -> Type {
- match self {
- Statistics::Boolean(_) => Type::BOOLEAN,
- Statistics::Int32(_) => Type::INT32,
- Statistics::Int64(_) => Type::INT64,
- Statistics::Int96(_) => Type::INT96,
- Statistics::Float(_) => Type::FLOAT,
- Statistics::Double(_) => Type::DOUBLE,
- Statistics::ByteArray(_) => Type::BYTE_ARRAY,
- Statistics::FixedLenByteArray(_) => Type::FIXED_LEN_BYTE_ARRAY,
- }
- }
-}
-
-impl fmt::Display for Statistics {
- fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
- match self {
- Statistics::Boolean(typed) => write!(f, "{}", typed),
- Statistics::Int32(typed) => write!(f, "{}", typed),
- Statistics::Int64(typed) => write!(f, "{}", typed),
- Statistics::Int96(typed) => write!(f, "{}", typed),
- Statistics::Float(typed) => write!(f, "{}", typed),
- Statistics::Double(typed) => write!(f, "{}", typed),
- Statistics::ByteArray(typed) => write!(f, "{}", typed),
- Statistics::FixedLenByteArray(typed) => write!(f, "{}", typed),
- }
- }
-}
-
-/// Typed implementation for [`Statistics`].
-#[derive(Clone)]
-pub struct TypedStatistics<T: DataType> {
- min: Option<T::T>,
- max: Option<T::T>,
- // Distinct count could be omitted in some cases
- distinct_count: Option<u64>,
- null_count: u64,
- is_min_max_deprecated: bool,
-}
-
-impl<T: DataType> TypedStatistics<T> {
- /// Creates new typed statistics.
- pub fn new(
- min: Option<T::T>,
- max: Option<T::T>,
- distinct_count: Option<u64>,
- null_count: u64,
- is_min_max_deprecated: bool,
- ) -> Self {
- Self {
- min,
- max,
- distinct_count,
- null_count,
- is_min_max_deprecated,
- }
- }
-
- /// Returns min value of the statistics.
- ///
- /// Panics if min value is not set, e.g. all values are `null`.
- /// Use `has_min_max_set` method to check that.
- pub fn min(&self) -> &T::T {
- self.min.as_ref().unwrap()
- }
-
- /// Returns max value of the statistics.
- ///
- /// Panics if max value is not set, e.g. all values are `null`.
- /// Use `has_min_max_set` method to check that.
- pub fn max(&self) -> &T::T {
- self.max.as_ref().unwrap()
- }
-
- /// Returns min value as bytes of the statistics.
- ///
- /// Panics if min value is not set, use `has_min_max_set` method to check
- /// if values are set.
- pub fn min_bytes(&self) -> &[u8] {
- self.min().as_bytes()
- }
-
- /// Returns max value as bytes of the statistics.
- ///
- /// Panics if max value is not set, use `has_min_max_set` method to check
- /// if values are set.
- pub fn max_bytes(&self) -> &[u8] {
- self.max().as_bytes()
- }
-
- /// Whether or not min and max values are set.
- /// Normally both min/max values will be set to `Some(value)` or `None`.
- fn has_min_max_set(&self) -> bool {
- self.min.is_some() && self.max.is_some()
- }
-
- /// Returns optional value of number of distinct values occurring.
- fn distinct_count(&self) -> Option<u64> {
- self.distinct_count
- }
-
- /// Returns null count.
- fn null_count(&self) -> u64 {
- self.null_count
- }
-
- /// Returns `true` if statistics were created using old min/max fields.
- fn is_min_max_deprecated(&self) -> bool {
- self.is_min_max_deprecated
- }
-}
-
-impl<T: DataType> fmt::Display for TypedStatistics<T> {
- fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
- write!(f, "{{")?;
- write!(f, "min: ")?;
- match self.min {
- Some(ref value) => write!(f, "{}", value)?,
- None => write!(f, "N/A")?,
- }
- write!(f, ", max: ")?;
- match self.max {
- Some(ref value) => write!(f, "{}", value)?,
- None => write!(f, "N/A")?,
- }
- write!(f, ", distinct_count: ")?;
- match self.distinct_count {
- Some(value) => write!(f, "{}", value)?,
- None => write!(f, "N/A")?,
- }
- write!(f, ", null_count: {}", self.null_count)?;
- write!(f, ", min_max_deprecated: {}", self.is_min_max_deprecated)?;
- write!(f, "}}")
- }
-}
-
-impl<T: DataType> fmt::Debug for TypedStatistics<T> {
- fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
- write!(
- f,
- "{{min: {:?}, max: {:?}, distinct_count: {:?}, null_count: {}, \
- min_max_deprecated: {}}}",
- self.min,
- self.max,
- self.distinct_count,
- self.null_count,
- self.is_min_max_deprecated
- )
- }
-}
-
-impl<T: DataType> cmp::PartialEq for TypedStatistics<T> {
- fn eq(&self, other: &TypedStatistics<T>) -> bool {
- self.min == other.min
- && self.max == other.max
- && self.distinct_count == other.distinct_count
- && self.null_count == other.null_count
- && self.is_min_max_deprecated == other.is_min_max_deprecated
- }
-}
-
-#[cfg(test)]
-mod tests {
- use super::*;
-
- #[test]
- fn test_statistics_min_max_bytes() {
- let stats = Statistics::int32(Some(-123), Some(234), None, 1, false);
- assert!(stats.has_min_max_set());
- assert_eq!(stats.min_bytes(), (-123).as_bytes());
- assert_eq!(stats.max_bytes(), 234.as_bytes());
-
- let stats = Statistics::byte_array(
- Some(ByteArray::from(vec![1, 2, 3])),
- Some(ByteArray::from(vec![3, 4, 5])),
- None,
- 1,
- true,
- );
- assert!(stats.has_min_max_set());
- assert_eq!(stats.min_bytes(), &[1, 2, 3]);
- assert_eq!(stats.max_bytes(), &[3, 4, 5]);
- }
-
- #[test]
- #[should_panic(expected = "Statistics null count is negative (-10)")]
- fn test_statistics_negative_null_count() {
- let thrift_stats = TStatistics {
- max: None,
- min: None,
- null_count: Some(-10),
- distinct_count: None,
- max_value: None,
- min_value: None,
- };
-
- from_thrift(Type::INT32, Some(thrift_stats));
- }
-
- #[test]
- fn test_statistics_thrift_none() {
- assert_eq!(from_thrift(Type::INT32, None), None);
- assert_eq!(from_thrift(Type::BYTE_ARRAY, None), None);
- }
-
- #[test]
- fn test_statistics_debug() {
- let stats = Statistics::int32(Some(1), Some(12), None, 12, true);
- assert_eq!(
- format!("{:?}", stats),
- "Int32({min: Some(1), max: Some(12), distinct_count: None, null_count: 12, \
- min_max_deprecated: true})"
- );
-
- let stats = Statistics::int32(None, None, None, 7, false);
- assert_eq!(
- format!("{:?}", stats),
- "Int32({min: None, max: None, distinct_count: None, null_count: 7, \
- min_max_deprecated: false})"
- )
- }
-
- #[test]
- fn test_statistics_display() {
- let stats = Statistics::int32(Some(1), Some(12), None, 12, true);
- assert_eq!(
- format!("{}", stats),
- "{min: 1, max: 12, distinct_count: N/A, null_count: 12, min_max_deprecated: true}"
- );
-
- let stats = Statistics::int64(None, None, None, 7, false);
- assert_eq!(
- format!("{}", stats),
- "{min: N/A, max: N/A, distinct_count: N/A, null_count: 7, min_max_deprecated: \
- false}"
- );
-
- let stats = Statistics::int96(
- Some(Int96::from(vec![1, 0, 0])),
- Some(Int96::from(vec![2, 3, 4])),
- None,
- 3,
- true,
- );
- assert_eq!(
- format!("{}", stats),
- "{min: [1, 0, 0], max: [2, 3, 4], distinct_count: N/A, null_count: 3, \
- min_max_deprecated: true}"
- );
-
- let stats = Statistics::byte_array(
- Some(ByteArray::from(vec![1u8])),
- Some(ByteArray::from(vec![2u8])),
- Some(5),
- 7,
- false,
- );
- assert_eq!(
- format!("{}", stats),
- "{min: [1], max: [2], distinct_count: 5, null_count: 7, min_max_deprecated: false}"
- );
- }
-
- #[test]
- fn test_statistics_partial_eq() {
- let expected = Statistics::int32(Some(12), Some(45), None, 11, true);
-
- assert!(Statistics::int32(Some(12), Some(45), None, 11, true) == expected);
- assert!(Statistics::int32(Some(11), Some(45), None, 11, true) != expected);
- assert!(Statistics::int32(Some(12), Some(44), None, 11, true) != expected);
- assert!(Statistics::int32(Some(12), Some(45), None, 23, true) != expected);
- assert!(Statistics::int32(Some(12), Some(45), None, 11, false) != expected);
-
- assert!(
- Statistics::int32(Some(12), Some(45), None, 11, false)
- != Statistics::int64(Some(12), Some(45), None, 11, false)
- );
-
- assert!(
- Statistics::boolean(Some(false), Some(true), None, 0, true)
- != Statistics::double(Some(1.2), Some(4.5), None, 0, true)
- );
-
- assert!(
- Statistics::byte_array(
- Some(ByteArray::from(vec![1, 2, 3])),
- Some(ByteArray::from(vec![1, 2, 3])),
- None,
- 0,
- true
- ) != Statistics::fixed_len_byte_array(
- Some(ByteArray::from(vec![1, 2, 3]).into()),
- Some(ByteArray::from(vec![1, 2, 3]).into()),
- None,
- 0,
- true
- )
- );
- }
-
- #[test]
- fn test_statistics_from_thrift() {
- // Helper method to check statistics conversion.
- fn check_stats(stats: Statistics) {
- let tpe = stats.physical_type();
- let thrift_stats = to_thrift(Some(&stats));
- assert_eq!(from_thrift(tpe, thrift_stats), Some(stats));
- }
-
- check_stats(Statistics::boolean(Some(false), Some(true), None, 7, true));
- check_stats(Statistics::boolean(Some(false), Some(true), None, 7, true));
- check_stats(Statistics::boolean(Some(false), Some(true), None, 0, false));
- check_stats(Statistics::boolean(Some(true), Some(true), None, 7, true));
- check_stats(Statistics::boolean(Some(false), Some(false), None, 7, true));
- check_stats(Statistics::boolean(None, None, None, 7, true));
-
- check_stats(Statistics::int32(Some(-100), Some(500), None, 7, true));
- check_stats(Statistics::int32(Some(-100), Some(500), None, 0, false));
- check_stats(Statistics::int32(None, None, None, 7, true));
-
- check_stats(Statistics::int64(Some(-100), Some(200), None, 7, true));
- check_stats(Statistics::int64(Some(-100), Some(200), None, 0, false));
- check_stats(Statistics::int64(None, None, None, 7, true));
-
- check_stats(Statistics::float(Some(1.2), Some(3.4), None, 7, true));
- check_stats(Statistics::float(Some(1.2), Some(3.4), None, 0, false));
- check_stats(Statistics::float(None, None, None, 7, true));
-
- check_stats(Statistics::double(Some(1.2), Some(3.4), None, 7, true));
- check_stats(Statistics::double(Some(1.2), Some(3.4), None, 0, false));
- check_stats(Statistics::double(None, None, None, 7, true));
-
- check_stats(Statistics::byte_array(
- Some(ByteArray::from(vec![1, 2, 3])),
- Some(ByteArray::from(vec![3, 4, 5])),
- None,
- 7,
- true,
- ));
- check_stats(Statistics::byte_array(None, None, None, 7, true));
-
- check_stats(Statistics::fixed_len_byte_array(
- Some(ByteArray::from(vec![1, 2, 3]).into()),
- Some(ByteArray::from(vec![3, 4, 5]).into()),
- None,
- 7,
- true,
- ));
- check_stats(Statistics::fixed_len_byte_array(None, None, None, 7, true));
- }
-}
diff --git a/parquet/src/file/writer.rs b/parquet/src/file/writer.rs
deleted file mode 100644
index e1c2dc6..0000000
--- a/parquet/src/file/writer.rs
+++ /dev/null
@@ -1,1164 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-//! Contains file writer API, and provides methods to write row groups and columns by
-//! using row group writers and column writers respectively.
-
-use std::{
- io::{Seek, SeekFrom, Write},
- sync::Arc,
-};
-
-use byteorder::{ByteOrder, LittleEndian};
-use parquet_format as parquet;
-use thrift::protocol::{TCompactOutputProtocol, TOutputProtocol};
-
-use crate::basic::PageType;
-use crate::column::{
- page::{CompressedPage, Page, PageWriteSpec, PageWriter},
- writer::{get_column_writer, ColumnWriter},
-};
-use crate::errors::{ParquetError, Result};
-use crate::file::{
- metadata::*, properties::WriterPropertiesPtr,
- statistics::to_thrift as statistics_to_thrift, FOOTER_SIZE, PARQUET_MAGIC,
-};
-use crate::schema::types::{self, SchemaDescPtr, SchemaDescriptor, TypePtr};
-use crate::util::io::{FileSink, Position};
-
-// Exposed publically so client code can implement [`ParquetWriter`]
-pub use crate::util::io::TryClone;
-
-// Exposed publically for convenience of writing Parquet to a buffer of bytes
-pub use crate::util::cursor::InMemoryWriteableCursor;
-
-// ----------------------------------------------------------------------
-// APIs for file & row group writers
-
-/// Parquet file writer API.
-/// Provides methods to write row groups sequentially.
-///
-/// The main workflow should be as following:
-/// - Create file writer, this will open a new file and potentially write some metadata.
-/// - Request a new row group writer by calling `next_row_group`.
-/// - Once finished writing row group, close row group writer by passing it into
-/// `close_row_group` method - this will finalise row group metadata and update metrics.
-/// - Write subsequent row groups, if necessary.
-/// - After all row groups have been written, close the file writer using `close` method.
-pub trait FileWriter {
- /// Creates new row group from this file writer.
- /// In case of IO error or Thrift error, returns `Err`.
- ///
- /// There is no limit on a number of row groups in a file; however, row groups have
- /// to be written sequentially. Every time the next row group is requested, the
- /// previous row group must be finalised and closed using `close_row_group` method.
- fn next_row_group(&mut self) -> Result<Box<dyn RowGroupWriter>>;
-
- /// Finalises and closes row group that was created using `next_row_group` method.
- /// After calling this method, the next row group is available for writes.
- fn close_row_group(
- &mut self,
- row_group_writer: Box<dyn RowGroupWriter>,
- ) -> Result<()>;
-
- /// Closes and finalises file writer, returning the file metadata.
- ///
- /// All row groups must be appended before this method is called.
- /// No writes are allowed after this point.
- ///
- /// Can be called multiple times. It is up to implementation to either result in
- /// no-op, or return an `Err` for subsequent calls.
- fn close(&mut self) -> Result<parquet::FileMetaData>;
-}
-
-/// Parquet row group writer API.
-/// Provides methods to access column writers in an iterator-like fashion, order is
-/// guaranteed to match the order of schema leaves (column descriptors).
-///
-/// All columns should be written sequentially; the main workflow is:
-/// - Request the next column using `next_column` method - this will return `None` if no
-/// more columns are available to write.
-/// - Once done writing a column, close column writer with `close_column` method - this
-/// will finalise column chunk metadata and update row group metrics.
-/// - Once all columns have been written, close row group writer with `close` method -
-/// it will return row group metadata and is no-op on already closed row group.
-pub trait RowGroupWriter {
- /// Returns the next column writer, if available; otherwise returns `None`.
- /// In case of any IO error or Thrift error, or if row group writer has already been
- /// closed returns `Err`.
- ///
- /// To request the next column writer, the previous one must be finalised and closed
- /// using `close_column`.
- fn next_column(&mut self) -> Result<Option<ColumnWriter>>;
-
- /// Closes column writer that was created using `next_column` method.
- /// This should be called before requesting the next column writer.
- fn close_column(&mut self, column_writer: ColumnWriter) -> Result<()>;
-
- /// Closes this row group writer and returns row group metadata.
- /// After calling this method row group writer must not be used.
- ///
- /// It is recommended to call this method before requesting another row group, but it
- /// will be closed automatically before returning a new row group.
- ///
- /// Can be called multiple times. In subsequent calls will result in no-op and return
- /// already created row group metadata.
- fn close(&mut self) -> Result<RowGroupMetaDataPtr>;
-}
-
-// ----------------------------------------------------------------------
-// Serialized impl for file & row group writers
-
-pub trait ParquetWriter: Write + Seek + TryClone {}
-impl<T: Write + Seek + TryClone> ParquetWriter for T {}
-
-/// A serialized implementation for Parquet [`FileWriter`].
-/// See documentation on file writer for more information.
-pub struct SerializedFileWriter<W: ParquetWriter> {
- buf: W,
- schema: TypePtr,
- descr: SchemaDescPtr,
- props: WriterPropertiesPtr,
- total_num_rows: i64,
- row_groups: Vec<RowGroupMetaDataPtr>,
- previous_writer_closed: bool,
- is_closed: bool,
-}
-
-impl<W: ParquetWriter> SerializedFileWriter<W> {
- /// Creates new file writer.
- pub fn new(
- mut buf: W,
- schema: TypePtr,
- properties: WriterPropertiesPtr,
- ) -> Result<Self> {
- Self::start_file(&mut buf)?;
- Ok(Self {
- buf,
- schema: schema.clone(),
- descr: Arc::new(SchemaDescriptor::new(schema)),
- props: properties,
- total_num_rows: 0,
- row_groups: Vec::new(),
- previous_writer_closed: true,
- is_closed: false,
- })
- }
-
- /// Writes magic bytes at the beginning of the file.
- fn start_file(buf: &mut W) -> Result<()> {
- buf.write_all(&PARQUET_MAGIC)?;
- Ok(())
- }
-
- /// Finalises active row group writer, otherwise no-op.
- fn finalise_row_group_writer(
- &mut self,
- mut row_group_writer: Box<dyn RowGroupWriter>,
- ) -> Result<()> {
- let row_group_metadata = row_group_writer.close()?;
- self.total_num_rows += row_group_metadata.num_rows();
- self.row_groups.push(row_group_metadata);
- Ok(())
- }
-
- /// Assembles and writes metadata at the end of the file.
- fn write_metadata(&mut self) -> Result<parquet::FileMetaData> {
- let file_metadata = parquet::FileMetaData {
- version: self.props.writer_version().as_num(),
- schema: types::to_thrift(self.schema.as_ref())?,
- num_rows: self.total_num_rows as i64,
- row_groups: self
- .row_groups
- .as_slice()
- .iter()
- .map(|v| v.to_thrift())
- .collect(),
- key_value_metadata: self.props.key_value_metadata().to_owned(),
- created_by: Some(self.props.created_by().to_owned()),
- column_orders: None,
- };
-
- // Write file metadata
- let start_pos = self.buf.seek(SeekFrom::Current(0))?;
- {
- let mut protocol = TCompactOutputProtocol::new(&mut self.buf);
- file_metadata.write_to_out_protocol(&mut protocol)?;
- protocol.flush()?;
- }
- let end_pos = self.buf.seek(SeekFrom::Current(0))?;
-
- // Write footer
- let mut footer_buffer: [u8; FOOTER_SIZE] = [0; FOOTER_SIZE];
- let metadata_len = (end_pos - start_pos) as i32;
- LittleEndian::write_i32(&mut footer_buffer, metadata_len);
- (&mut footer_buffer[4..]).write_all(&PARQUET_MAGIC)?;
- self.buf.write_all(&footer_buffer)?;
- Ok(file_metadata)
- }
-
- #[inline]
- fn assert_closed(&self) -> Result<()> {
- if self.is_closed {
- Err(general_err!("File writer is closed"))
- } else {
- Ok(())
- }
- }
-
- #[inline]
- fn assert_previous_writer_closed(&self) -> Result<()> {
- if !self.previous_writer_closed {
- Err(general_err!("Previous row group writer was not closed"))
- } else {
- Ok(())
- }
- }
-}
-
-impl<W: 'static + ParquetWriter> FileWriter for SerializedFileWriter<W> {
- #[inline]
- fn next_row_group(&mut self) -> Result<Box<dyn RowGroupWriter>> {
- self.assert_closed()?;
- self.assert_previous_writer_closed()?;
- let row_group_writer = SerializedRowGroupWriter::new(
- self.descr.clone(),
- self.props.clone(),
- &self.buf,
- );
- self.previous_writer_closed = false;
- Ok(Box::new(row_group_writer))
- }
-
- #[inline]
- fn close_row_group(
- &mut self,
- row_group_writer: Box<dyn RowGroupWriter>,
- ) -> Result<()> {
- self.assert_closed()?;
- let res = self.finalise_row_group_writer(row_group_writer);
- self.previous_writer_closed = res.is_ok();
- res
- }
-
- #[inline]
- fn close(&mut self) -> Result<parquet::FileMetaData> {
- self.assert_closed()?;
- self.assert_previous_writer_closed()?;
- let metadata = self.write_metadata()?;
- self.is_closed = true;
- Ok(metadata)
- }
-}
-
-/// A serialized implementation for Parquet [`RowGroupWriter`].
-/// Coordinates writing of a row group with column writers.
-/// See documentation on row group writer for more information.
-pub struct SerializedRowGroupWriter<W: ParquetWriter> {
- descr: SchemaDescPtr,
- props: WriterPropertiesPtr,
- buf: W,
- total_rows_written: Option<u64>,
- total_bytes_written: u64,
- column_index: usize,
- previous_writer_closed: bool,
- row_group_metadata: Option<RowGroupMetaDataPtr>,
- column_chunks: Vec<ColumnChunkMetaData>,
-}
-
-impl<W: 'static + ParquetWriter> SerializedRowGroupWriter<W> {
- pub fn new(
- schema_descr: SchemaDescPtr,
- properties: WriterPropertiesPtr,
- buf: &W,
- ) -> Self {
- let num_columns = schema_descr.num_columns();
- Self {
- descr: schema_descr,
- props: properties,
- buf: buf.try_clone().unwrap(),
- total_rows_written: None,
- total_bytes_written: 0,
- column_index: 0,
- previous_writer_closed: true,
- row_group_metadata: None,
- column_chunks: Vec::with_capacity(num_columns),
- }
- }
-
- /// Checks and finalises current column writer.
- fn finalise_column_writer(&mut self, writer: ColumnWriter) -> Result<()> {
- let (bytes_written, rows_written, metadata) = match writer {
- ColumnWriter::BoolColumnWriter(typed) => typed.close()?,
- ColumnWriter::Int32ColumnWriter(typed) => typed.close()?,
- ColumnWriter::Int64ColumnWriter(typed) => typed.close()?,
- ColumnWriter::Int96ColumnWriter(typed) => typed.close()?,
- ColumnWriter::FloatColumnWriter(typed) => typed.close()?,
- ColumnWriter::DoubleColumnWriter(typed) => typed.close()?,
- ColumnWriter::ByteArrayColumnWriter(typed) => typed.close()?,
- ColumnWriter::FixedLenByteArrayColumnWriter(typed) => typed.close()?,
- };
-
- // Update row group writer metrics
- self.total_bytes_written += bytes_written;
- self.column_chunks.push(metadata);
- if let Some(rows) = self.total_rows_written {
- if rows != rows_written {
- return Err(general_err!(
- "Incorrect number of rows, expected {} != {} rows",
- rows,
- rows_written
- ));
- }
- } else {
- self.total_rows_written = Some(rows_written);
- }
-
- Ok(())
- }
-
- #[inline]
- fn assert_closed(&self) -> Result<()> {
- if self.row_group_metadata.is_some() {
- Err(general_err!("Row group writer is closed"))
- } else {
- Ok(())
- }
- }
-
- #[inline]
- fn assert_previous_writer_closed(&self) -> Result<()> {
- if !self.previous_writer_closed {
- Err(general_err!("Previous column writer was not closed"))
- } else {
- Ok(())
- }
- }
-}
-
-impl<W: 'static + ParquetWriter> RowGroupWriter for SerializedRowGroupWriter<W> {
- #[inline]
- fn next_column(&mut self) -> Result<Option<ColumnWriter>> {
- self.assert_closed()?;
- self.assert_previous_writer_closed()?;
-
- if self.column_index >= self.descr.num_columns() {
- return Ok(None);
- }
- let sink = FileSink::new(&self.buf);
- let page_writer = Box::new(SerializedPageWriter::new(sink));
- let column_writer = get_column_writer(
- self.descr.column(self.column_index),
- self.props.clone(),
- page_writer,
- );
- self.column_index += 1;
- self.previous_writer_closed = false;
-
- Ok(Some(column_writer))
- }
-
- #[inline]
- fn close_column(&mut self, column_writer: ColumnWriter) -> Result<()> {
- let res = self.finalise_column_writer(column_writer);
- self.previous_writer_closed = res.is_ok();
- res
- }
-
- #[inline]
- fn close(&mut self) -> Result<RowGroupMetaDataPtr> {
- if self.row_group_metadata.is_none() {
- self.assert_previous_writer_closed()?;
-
- let column_chunks = std::mem::take(&mut self.column_chunks);
- let row_group_metadata = RowGroupMetaData::builder(self.descr.clone())
- .set_column_metadata(column_chunks)
- .set_total_byte_size(self.total_bytes_written as i64)
- .set_num_rows(self.total_rows_written.unwrap_or(0) as i64)
- .build()?;
-
- self.row_group_metadata = Some(Arc::new(row_group_metadata));
- }
-
- let metadata = self.row_group_metadata.as_ref().unwrap().clone();
- Ok(metadata)
- }
-}
-
-/// A serialized implementation for Parquet [`PageWriter`].
-/// Writes and serializes pages and metadata into output stream.
-///
-/// `SerializedPageWriter` should not be used after calling `close()`.
-pub struct SerializedPageWriter<T: Write + Position> {
- sink: T,
-}
-
-impl<T: Write + Position> SerializedPageWriter<T> {
- /// Creates new page writer.
- pub fn new(sink: T) -> Self {
- Self { sink }
- }
-
- /// Serializes page header into Thrift.
- /// Returns number of bytes that have been written into the sink.
- #[inline]
- fn serialize_page_header(&mut self, header: parquet::PageHeader) -> Result<usize> {
- let start_pos = self.sink.pos();
- {
- let mut protocol = TCompactOutputProtocol::new(&mut self.sink);
- header.write_to_out_protocol(&mut protocol)?;
- protocol.flush()?;
- }
- Ok((self.sink.pos() - start_pos) as usize)
- }
-
- /// Serializes column chunk into Thrift.
- /// Returns Ok() if there are not errors serializing and writing data into the sink.
- #[inline]
- fn serialize_column_chunk(&mut self, chunk: parquet::ColumnChunk) -> Result<()> {
- let mut protocol = TCompactOutputProtocol::new(&mut self.sink);
- chunk.write_to_out_protocol(&mut protocol)?;
- protocol.flush()?;
- Ok(())
- }
-}
-
-impl<T: Write + Position> PageWriter for SerializedPageWriter<T> {
- fn write_page(&mut self, page: CompressedPage) -> Result<PageWriteSpec> {
- let uncompressed_size = page.uncompressed_size();
- let compressed_size = page.compressed_size();
- let num_values = page.num_values();
- let encoding = page.encoding();
- let page_type = page.page_type();
-
- let mut page_header = parquet::PageHeader {
- type_: page_type.into(),
- uncompressed_page_size: uncompressed_size as i32,
- compressed_page_size: compressed_size as i32,
- // TODO: Add support for crc checksum
- crc: None,
- data_page_header: None,
- index_page_header: None,
- dictionary_page_header: None,
- data_page_header_v2: None,
- };
-
- match *page.compressed_page() {
- Page::DataPage {
- def_level_encoding,
- rep_level_encoding,
- ref statistics,
- ..
- } => {
- let data_page_header = parquet::DataPageHeader {
- num_values: num_values as i32,
- encoding: encoding.into(),
- definition_level_encoding: def_level_encoding.into(),
- repetition_level_encoding: rep_level_encoding.into(),
- statistics: statistics_to_thrift(statistics.as_ref()),
- };
- page_header.data_page_header = Some(data_page_header);
- }
- Page::DataPageV2 {
- num_nulls,
- num_rows,
- def_levels_byte_len,
- rep_levels_byte_len,
- is_compressed,
- ref statistics,
- ..
- } => {
- let data_page_header_v2 = parquet::DataPageHeaderV2 {
- num_values: num_values as i32,
- num_nulls: num_nulls as i32,
- num_rows: num_rows as i32,
- encoding: encoding.into(),
- definition_levels_byte_length: def_levels_byte_len as i32,
- repetition_levels_byte_length: rep_levels_byte_len as i32,
- is_compressed: Some(is_compressed),
- statistics: statistics_to_thrift(statistics.as_ref()),
- };
- page_header.data_page_header_v2 = Some(data_page_header_v2);
- }
- Page::DictionaryPage { is_sorted, .. } => {
- let dictionary_page_header = parquet::DictionaryPageHeader {
- num_values: num_values as i32,
- encoding: encoding.into(),
- is_sorted: Some(is_sorted),
- };
- page_header.dictionary_page_header = Some(dictionary_page_header);
- }
- }
-
- let start_pos = self.sink.pos();
-
- let header_size = self.serialize_page_header(page_header)?;
- self.sink.write_all(page.data())?;
-
- let mut spec = PageWriteSpec::new();
- spec.page_type = page_type;
- spec.uncompressed_size = uncompressed_size + header_size;
- spec.compressed_size = compressed_size + header_size;
- spec.offset = start_pos;
- spec.bytes_written = self.sink.pos() - start_pos;
- // Number of values is incremented for data pages only
- if page_type == PageType::DATA_PAGE || page_type == PageType::DATA_PAGE_V2 {
- spec.num_values = num_values;
- }
-
- Ok(spec)
- }
-
- fn write_metadata(&mut self, metadata: &ColumnChunkMetaData) -> Result<()> {
- self.serialize_column_chunk(metadata.to_thrift())
- }
-
- fn close(&mut self) -> Result<()> {
- self.sink.flush()?;
- Ok(())
- }
-}
-
-#[cfg(test)]
-mod tests {
- use super::*;
-
- use std::{fs::File, io::Cursor};
-
- use crate::basic::{Compression, Encoding, IntType, LogicalType, Repetition, Type};
- use crate::column::page::PageReader;
- use crate::compression::{create_codec, Codec};
- use crate::file::{
- properties::{WriterProperties, WriterVersion},
- reader::{FileReader, SerializedFileReader, SerializedPageReader},
- statistics::{from_thrift, to_thrift, Statistics},
- };
- use crate::record::RowAccessor;
- use crate::util::{memory::ByteBufferPtr, test_common::get_temp_file};
-
- #[test]
- fn test_file_writer_error_after_close() {
- let file = get_temp_file("test_file_writer_error_after_close", &[]);
- let schema = Arc::new(types::Type::group_type_builder("schema").build().unwrap());
- let props = Arc::new(WriterProperties::builder().build());
- let mut writer = SerializedFileWriter::new(file, schema, props).unwrap();
- writer.close().unwrap();
- {
- let res = writer.next_row_group();
- assert!(res.is_err());
- if let Err(err) = res {
- assert_eq!(format!("{}", err), "Parquet error: File writer is closed");
- }
- }
- {
- let res = writer.close();
- assert!(res.is_err());
- if let Err(err) = res {
- assert_eq!(format!("{}", err), "Parquet error: File writer is closed");
- }
- }
- }
-
- #[test]
- fn test_row_group_writer_error_after_close() {
- let file = get_temp_file("test_file_writer_row_group_error_after_close", &[]);
- let schema = Arc::new(types::Type::group_type_builder("schema").build().unwrap());
- let props = Arc::new(WriterProperties::builder().build());
- let mut writer = SerializedFileWriter::new(file, schema, props).unwrap();
- let mut row_group_writer = writer.next_row_group().unwrap();
- row_group_writer.close().unwrap();
-
- let res = row_group_writer.next_column();
- assert!(res.is_err());
- if let Err(err) = res {
- assert_eq!(
- format!("{}", err),
- "Parquet error: Row group writer is closed"
- );
- }
- }
-
- #[test]
- fn test_row_group_writer_error_not_all_columns_written() {
- let file =
- get_temp_file("test_row_group_writer_error_not_all_columns_written", &[]);
- let schema = Arc::new(
- types::Type::group_type_builder("schema")
- .with_fields(&mut vec![Arc::new(
- types::Type::primitive_type_builder("col1", Type::INT32)
- .build()
- .unwrap(),
- )])
- .build()
- .unwrap(),
- );
- let props = Arc::new(WriterProperties::builder().build());
- let mut writer = SerializedFileWriter::new(file, schema, props).unwrap();
- let mut row_group_writer = writer.next_row_group().unwrap();
- let res = row_group_writer.close();
- assert!(res.is_err());
- if let Err(err) = res {
- assert_eq!(
- format!("{}", err),
- "Parquet error: Column length mismatch: 1 != 0"
- );
- }
- }
-
- #[test]
- fn test_row_group_writer_num_records_mismatch() {
- let file = get_temp_file("test_row_group_writer_num_records_mismatch", &[]);
- let schema = Arc::new(
- types::Type::group_type_builder("schema")
- .with_fields(&mut vec![
- Arc::new(
- types::Type::primitive_type_builder("col1", Type::INT32)
- .with_repetition(Repetition::REQUIRED)
- .build()
- .unwrap(),
- ),
- Arc::new(
- types::Type::primitive_type_builder("col2", Type::INT32)
- .with_repetition(Repetition::REQUIRED)
- .build()
- .unwrap(),
- ),
- ])
- .build()
- .unwrap(),
- );
- let props = Arc::new(WriterProperties::builder().build());
- let mut writer = SerializedFileWriter::new(file, schema, props).unwrap();
- let mut row_group_writer = writer.next_row_group().unwrap();
-
- let mut col_writer = row_group_writer.next_column().unwrap().unwrap();
- if let ColumnWriter::Int32ColumnWriter(ref mut typed) = col_writer {
- typed.write_batch(&[1, 2, 3], None, None).unwrap();
- }
- row_group_writer.close_column(col_writer).unwrap();
-
- let mut col_writer = row_group_writer.next_column().unwrap().unwrap();
- if let ColumnWriter::Int32ColumnWriter(ref mut typed) = col_writer {
- typed.write_batch(&[1, 2], None, None).unwrap();
- }
-
- let res = row_group_writer.close_column(col_writer);
- assert!(res.is_err());
- if let Err(err) = res {
- assert_eq!(
- format!("{}", err),
- "Parquet error: Incorrect number of rows, expected 3 != 2 rows"
- );
- }
- }
-
- #[test]
- fn test_file_writer_empty_file() {
- let file = get_temp_file("test_file_writer_write_empty_file", &[]);
-
- let schema = Arc::new(
- types::Type::group_type_builder("schema")
- .with_fields(&mut vec![Arc::new(
- types::Type::primitive_type_builder("col1", Type::INT32)
- .build()
- .unwrap(),
- )])
- .build()
- .unwrap(),
- );
- let props = Arc::new(WriterProperties::builder().build());
- let mut writer =
- SerializedFileWriter::new(file.try_clone().unwrap(), schema, props).unwrap();
- writer.close().unwrap();
-
- let reader = SerializedFileReader::new(file).unwrap();
- assert_eq!(reader.get_row_iter(None).unwrap().count(), 0);
- }
-
- #[test]
- fn test_file_writer_with_metadata() {
- let file = get_temp_file("test_file_writer_write_with_metadata", &[]);
-
- let schema = Arc::new(
- types::Type::group_type_builder("schema")
- .with_fields(&mut vec![Arc::new(
- types::Type::primitive_type_builder("col1", Type::INT32)
- .build()
- .unwrap(),
- )])
- .build()
- .unwrap(),
- );
- let props = Arc::new(
- WriterProperties::builder()
- .set_key_value_metadata(Some(vec![KeyValue::new(
- "key".to_string(),
- "value".to_string(),
- )]))
- .build(),
- );
- let mut writer =
- SerializedFileWriter::new(file.try_clone().unwrap(), schema, props).unwrap();
- writer.close().unwrap();
-
- let reader = SerializedFileReader::new(file).unwrap();
- assert_eq!(
- reader
- .metadata()
- .file_metadata()
- .key_value_metadata()
- .to_owned()
- .unwrap()
- .len(),
- 1
- );
- }
-
- #[test]
- fn test_file_writer_v2_with_metadata() {
- let file = get_temp_file("test_file_writer_v2_write_with_metadata", &[]);
- let field_logical_type = Some(LogicalType::INTEGER(IntType {
- bit_width: 8,
- is_signed: false,
- }));
- let field = Arc::new(
- types::Type::primitive_type_builder("col1", Type::INT32)
- .with_logical_type(field_logical_type.clone())
- .with_converted_type(field_logical_type.into())
- .build()
- .unwrap(),
- );
- let schema = Arc::new(
- types::Type::group_type_builder("schema")
- .with_fields(&mut vec![field.clone()])
- .build()
- .unwrap(),
- );
- let props = Arc::new(
- WriterProperties::builder()
- .set_key_value_metadata(Some(vec![KeyValue::new(
- "key".to_string(),
- "value".to_string(),
- )]))
- .set_writer_version(WriterVersion::PARQUET_2_0)
- .build(),
- );
- let mut writer =
- SerializedFileWriter::new(file.try_clone().unwrap(), schema, props).unwrap();
- writer.close().unwrap();
-
- let reader = SerializedFileReader::new(file).unwrap();
-
- assert_eq!(
- reader
- .metadata()
- .file_metadata()
- .key_value_metadata()
- .to_owned()
- .unwrap()
- .len(),
- 1
- );
-
- // ARROW-11803: Test that the converted and logical types have been populated
- let fields = reader.metadata().file_metadata().schema().get_fields();
- assert_eq!(fields.len(), 1);
- let read_field = fields.get(0).unwrap();
- assert_eq!(read_field, &field);
- }
-
- #[test]
- fn test_file_writer_empty_row_groups() {
- let file = get_temp_file("test_file_writer_write_empty_row_groups", &[]);
- test_file_roundtrip(file, vec![]);
- }
-
- #[test]
- fn test_file_writer_single_row_group() {
- let file = get_temp_file("test_file_writer_write_single_row_group", &[]);
- test_file_roundtrip(file, vec![vec![1, 2, 3, 4, 5]]);
- }
-
- #[test]
- fn test_file_writer_multiple_row_groups() {
- let file = get_temp_file("test_file_writer_write_multiple_row_groups", &[]);
- test_file_roundtrip(
- file,
- vec![
- vec![1, 2, 3, 4, 5],
- vec![1, 2, 3],
- vec![1],
- vec![1, 2, 3, 4, 5, 6],
- ],
- );
- }
-
- #[test]
- fn test_file_writer_multiple_large_row_groups() {
- let file = get_temp_file("test_file_writer_multiple_large_row_groups", &[]);
- test_file_roundtrip(
- file,
- vec![vec![123; 1024], vec![124; 1000], vec![125; 15], vec![]],
- );
- }
-
- #[test]
- fn test_page_writer_data_pages() {
- let pages = vec![
- Page::DataPage {
- buf: ByteBufferPtr::new(vec![1, 2, 3, 4, 5, 6, 7, 8]),
- num_values: 10,
- encoding: Encoding::DELTA_BINARY_PACKED,
- def_level_encoding: Encoding::RLE,
- rep_level_encoding: Encoding::RLE,
- statistics: Some(Statistics::int32(Some(1), Some(3), None, 7, true)),
- },
- Page::DataPageV2 {
- buf: ByteBufferPtr::new(vec![4; 128]),
- num_values: 10,
- encoding: Encoding::DELTA_BINARY_PACKED,
- num_nulls: 2,
- num_rows: 12,
- def_levels_byte_len: 24,
- rep_levels_byte_len: 32,
- is_compressed: false,
- statistics: Some(Statistics::int32(Some(1), Some(3), None, 7, true)),
- },
- ];
-
- test_page_roundtrip(&pages[..], Compression::SNAPPY, Type::INT32);
- test_page_roundtrip(&pages[..], Compression::UNCOMPRESSED, Type::INT32);
- }
-
- #[test]
- fn test_page_writer_dict_pages() {
- let pages = vec![
- Page::DictionaryPage {
- buf: ByteBufferPtr::new(vec![1, 2, 3, 4, 5]),
- num_values: 5,
- encoding: Encoding::RLE_DICTIONARY,
- is_sorted: false,
- },
- Page::DataPage {
- buf: ByteBufferPtr::new(vec![1, 2, 3, 4, 5, 6, 7, 8]),
- num_values: 10,
- encoding: Encoding::DELTA_BINARY_PACKED,
- def_level_encoding: Encoding::RLE,
- rep_level_encoding: Encoding::RLE,
- statistics: Some(Statistics::int32(Some(1), Some(3), None, 7, true)),
- },
- Page::DataPageV2 {
- buf: ByteBufferPtr::new(vec![4; 128]),
- num_values: 10,
- encoding: Encoding::DELTA_BINARY_PACKED,
- num_nulls: 2,
- num_rows: 12,
- def_levels_byte_len: 24,
- rep_levels_byte_len: 32,
- is_compressed: false,
- statistics: None,
- },
- ];
-
- test_page_roundtrip(&pages[..], Compression::SNAPPY, Type::INT32);
- test_page_roundtrip(&pages[..], Compression::UNCOMPRESSED, Type::INT32);
- }
-
- /// Tests writing and reading pages.
- /// Physical type is for statistics only, should match any defined statistics type in
- /// pages.
- fn test_page_roundtrip(pages: &[Page], codec: Compression, physical_type: Type) {
- let mut compressed_pages = vec![];
- let mut total_num_values = 0i64;
- let mut compressor = create_codec(codec).unwrap();
-
- for page in pages {
- let uncompressed_len = page.buffer().len();
-
- let compressed_page = match *page {
- Page::DataPage {
- ref buf,
- num_values,
- encoding,
- def_level_encoding,
- rep_level_encoding,
- ref statistics,
- } => {
- total_num_values += num_values as i64;
- let output_buf = compress_helper(compressor.as_mut(), buf.data());
-
- Page::DataPage {
- buf: ByteBufferPtr::new(output_buf),
- num_values,
- encoding,
- def_level_encoding,
- rep_level_encoding,
- statistics: from_thrift(
- physical_type,
- to_thrift(statistics.as_ref()),
- ),
- }
- }
- Page::DataPageV2 {
- ref buf,
- num_values,
- encoding,
- num_nulls,
- num_rows,
- def_levels_byte_len,
- rep_levels_byte_len,
- ref statistics,
- ..
- } => {
- total_num_values += num_values as i64;
- let offset = (def_levels_byte_len + rep_levels_byte_len) as usize;
- let cmp_buf =
- compress_helper(compressor.as_mut(), &buf.data()[offset..]);
- let mut output_buf = Vec::from(&buf.data()[..offset]);
- output_buf.extend_from_slice(&cmp_buf[..]);
-
- Page::DataPageV2 {
- buf: ByteBufferPtr::new(output_buf),
- num_values,
- encoding,
- num_nulls,
- num_rows,
- def_levels_byte_len,
- rep_levels_byte_len,
- is_compressed: compressor.is_some(),
- statistics: from_thrift(
- physical_type,
- to_thrift(statistics.as_ref()),
- ),
- }
- }
- Page::DictionaryPage {
- ref buf,
- num_values,
- encoding,
- is_sorted,
- } => {
- let output_buf = compress_helper(compressor.as_mut(), buf.data());
-
- Page::DictionaryPage {
- buf: ByteBufferPtr::new(output_buf),
- num_values,
- encoding,
- is_sorted,
- }
- }
- };
-
- let compressed_page = CompressedPage::new(compressed_page, uncompressed_len);
- compressed_pages.push(compressed_page);
- }
-
- let mut buffer: Vec<u8> = vec![];
- let mut result_pages: Vec<Page> = vec![];
- {
- let cursor = Cursor::new(&mut buffer);
- let mut page_writer = SerializedPageWriter::new(cursor);
-
- for page in compressed_pages {
- page_writer.write_page(page).unwrap();
- }
- page_writer.close().unwrap();
- }
- {
- let mut page_reader = SerializedPageReader::new(
- Cursor::new(&buffer),
- total_num_values,
- codec,
- physical_type,
- )
- .unwrap();
-
- while let Some(page) = page_reader.get_next_page().unwrap() {
- result_pages.push(page);
- }
- }
-
- assert_eq!(result_pages.len(), pages.len());
- for i in 0..result_pages.len() {
- assert_page(&result_pages[i], &pages[i]);
- }
- }
-
- /// Helper function to compress a slice
- fn compress_helper(compressor: Option<&mut Box<dyn Codec>>, data: &[u8]) -> Vec<u8> {
- let mut output_buf = vec![];
- if let Some(cmpr) = compressor {
- cmpr.compress(data, &mut output_buf).unwrap();
- } else {
- output_buf.extend_from_slice(data);
- }
- output_buf
- }
-
- /// Check if pages match.
- fn assert_page(left: &Page, right: &Page) {
- assert_eq!(left.page_type(), right.page_type());
- assert_eq!(left.buffer().data(), right.buffer().data());
- assert_eq!(left.num_values(), right.num_values());
- assert_eq!(left.encoding(), right.encoding());
- assert_eq!(to_thrift(left.statistics()), to_thrift(right.statistics()));
- }
-
- /// File write-read roundtrip.
- /// `data` consists of arrays of values for each row group.
- fn test_file_roundtrip(file: File, data: Vec<Vec<i32>>) {
- let schema = Arc::new(
- types::Type::group_type_builder("schema")
- .with_fields(&mut vec![Arc::new(
- types::Type::primitive_type_builder("col1", Type::INT32)
- .with_repetition(Repetition::REQUIRED)
- .build()
- .unwrap(),
- )])
- .build()
- .unwrap(),
- );
- let props = Arc::new(WriterProperties::builder().build());
- let mut file_writer = assert_send(
- SerializedFileWriter::new(file.try_clone().unwrap(), schema, props).unwrap(),
- );
- let mut rows: i64 = 0;
-
- for subset in &data {
- let mut row_group_writer = file_writer.next_row_group().unwrap();
- let col_writer = row_group_writer.next_column().unwrap();
- if let Some(mut writer) = col_writer {
- match writer {
- ColumnWriter::Int32ColumnWriter(ref mut typed) => {
- rows +=
- typed.write_batch(&subset[..], None, None).unwrap() as i64;
- }
- _ => {
- unimplemented!();
- }
- }
- row_group_writer.close_column(writer).unwrap();
- }
- file_writer.close_row_group(row_group_writer).unwrap();
- }
-
- file_writer.close().unwrap();
-
- let reader = assert_send(SerializedFileReader::new(file).unwrap());
- assert_eq!(reader.num_row_groups(), data.len());
- assert_eq!(
- reader.metadata().file_metadata().num_rows(),
- rows,
- "row count in metadata not equal to number of rows written"
- );
- for i in 0..reader.num_row_groups() {
- let row_group_reader = reader.get_row_group(i).unwrap();
- let iter = row_group_reader.get_row_iter(None).unwrap();
- let res = iter
- .map(|elem| elem.get_int(0).unwrap())
- .collect::<Vec<i32>>();
- assert_eq!(res, data[i]);
- }
- }
-
- fn assert_send<T: Send>(t: T) -> T {
- t
- }
-
- #[test]
- fn test_bytes_writer_empty_row_groups() {
- test_bytes_roundtrip(vec![]);
- }
-
- #[test]
- fn test_bytes_writer_single_row_group() {
- test_bytes_roundtrip(vec![vec![1, 2, 3, 4, 5]]);
- }
-
- #[test]
- fn test_bytes_writer_multiple_row_groups() {
- test_bytes_roundtrip(vec![
- vec![1, 2, 3, 4, 5],
- vec![1, 2, 3],
- vec![1],
- vec![1, 2, 3, 4, 5, 6],
- ]);
- }
-
- fn test_bytes_roundtrip(data: Vec<Vec<i32>>) {
- let cursor = InMemoryWriteableCursor::default();
-
- let schema = Arc::new(
- types::Type::group_type_builder("schema")
- .with_fields(&mut vec![Arc::new(
- types::Type::primitive_type_builder("col1", Type::INT32)
- .with_repetition(Repetition::REQUIRED)
- .build()
- .unwrap(),
- )])
- .build()
- .unwrap(),
- );
-
- let mut rows: i64 = 0;
- {
- let props = Arc::new(WriterProperties::builder().build());
- let mut writer =
- SerializedFileWriter::new(cursor.clone(), schema, props).unwrap();
-
- for subset in &data {
- let mut row_group_writer = writer.next_row_group().unwrap();
- let col_writer = row_group_writer.next_column().unwrap();
- if let Some(mut writer) = col_writer {
- match writer {
- ColumnWriter::Int32ColumnWriter(ref mut typed) => {
- rows += typed.write_batch(&subset[..], None, None).unwrap()
- as i64;
- }
- _ => {
- unimplemented!();
- }
- }
- row_group_writer.close_column(writer).unwrap();
- }
- writer.close_row_group(row_group_writer).unwrap();
- }
-
- writer.close().unwrap();
- }
-
- let buffer = cursor.into_inner().unwrap();
-
- let reading_cursor = crate::file::serialized_reader::SliceableCursor::new(buffer);
- let reader = SerializedFileReader::new(reading_cursor).unwrap();
-
- assert_eq!(reader.num_row_groups(), data.len());
- assert_eq!(
- reader.metadata().file_metadata().num_rows(),
- rows,
- "row count in metadata not equal to number of rows written"
- );
- for i in 0..reader.num_row_groups() {
- let row_group_reader = reader.get_row_group(i).unwrap();
- let iter = row_group_reader.get_row_iter(None).unwrap();
- let res = iter
- .map(|elem| elem.get_int(0).unwrap())
- .collect::<Vec<i32>>();
- assert_eq!(res, data[i]);
- }
- }
-}
diff --git a/parquet/src/lib.rs b/parquet/src/lib.rs
deleted file mode 100644
index 900e2b5..0000000
--- a/parquet/src/lib.rs
+++ /dev/null
@@ -1,57 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-#![allow(incomplete_features)]
-#![allow(dead_code)]
-#![allow(non_camel_case_types)]
-#![allow(
- clippy::approx_constant,
- clippy::cast_ptr_alignment,
- clippy::float_cmp,
- clippy::float_equality_without_abs,
- clippy::from_over_into,
- clippy::many_single_char_names,
- clippy::needless_range_loop,
- clippy::new_without_default,
- clippy::or_fun_call,
- clippy::same_item_push,
- clippy::too_many_arguments,
- clippy::transmute_ptr_to_ptr,
- clippy::upper_case_acronyms,
- clippy::vec_init_then_push
-)]
-
-#[macro_use]
-pub mod errors;
-pub mod basic;
-#[macro_use]
-pub mod data_type;
-
-// Exported for external use, such as benchmarks
-pub use self::encodings::{decoding, encoding};
-pub use self::util::memory;
-
-#[macro_use]
-pub mod util;
-#[cfg(any(feature = "arrow", test))]
-pub mod arrow;
-pub mod column;
-pub mod compression;
-mod encodings;
-pub mod file;
-pub mod record;
-pub mod schema;
diff --git a/parquet/src/record/api.rs b/parquet/src/record/api.rs
deleted file mode 100644
index 411016e..0000000
--- a/parquet/src/record/api.rs
+++ /dev/null
@@ -1,1846 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-//! Contains Row enum that is used to represent record in Rust.
-
-use std::fmt;
-
-use chrono::{TimeZone, Utc};
-use num_bigint::{BigInt, Sign};
-
-use crate::basic::{ConvertedType, Type as PhysicalType};
-use crate::data_type::{ByteArray, Decimal, Int96};
-use crate::errors::{ParquetError, Result};
-use crate::schema::types::ColumnDescPtr;
-
-#[cfg(feature = "cli")]
-use serde_json::Value;
-
-/// Macro as a shortcut to generate 'not yet implemented' panic error.
-macro_rules! nyi {
- ($column_descr:ident, $value:ident) => {{
- unimplemented!(
- "Conversion for physical type {}, converted type {}, value {:?}",
- $column_descr.physical_type(),
- $column_descr.converted_type(),
- $value
- );
- }};
-}
-
-/// `Row` represents a nested Parquet record.
-#[derive(Clone, Debug, PartialEq)]
-pub struct Row {
- fields: Vec<(String, Field)>,
-}
-
-#[allow(clippy::len_without_is_empty)]
-impl Row {
- /// Get the number of fields in this row.
- pub fn len(&self) -> usize {
- self.fields.len()
- }
-
- /// Get an iterator to go through all columns in the row.
- ///
- /// # Example
- ///
- /// ```no_run
- /// use std::fs::File;
- /// use parquet::record::Row;
- /// use parquet::file::reader::{FileReader, SerializedFileReader};
- ///
- /// let file = File::open("/path/to/file").unwrap();
- /// let reader = SerializedFileReader::new(file).unwrap();
- /// let row: Row = reader.get_row_iter(None).unwrap().next().unwrap();
- /// for (idx, (name, field)) in row.get_column_iter().enumerate() {
- /// println!("column index: {}, column name: {}, column value: {}", idx, name, field);
- /// }
- /// ```
- pub fn get_column_iter(&self) -> RowColumnIter {
- RowColumnIter {
- fields: &self.fields,
- curr: 0,
- count: self.fields.len(),
- }
- }
-
- #[cfg(feature = "cli")]
- pub fn to_json_value(&self) -> Value {
- Value::Object(
- self.fields
- .iter()
- .map(|(key, field)| (key.to_owned(), field.to_json_value()))
- .collect(),
- )
- }
-}
-
-pub struct RowColumnIter<'a> {
- fields: &'a Vec<(String, Field)>,
- curr: usize,
- count: usize,
-}
-
-impl<'a> Iterator for RowColumnIter<'a> {
- type Item = (&'a String, &'a Field);
-
- fn next(&mut self) -> Option<Self::Item> {
- let idx = self.curr;
- if idx >= self.count {
- return None;
- }
- self.curr += 1;
- Some((&self.fields[idx].0, &self.fields[idx].1))
- }
-}
-
-/// Trait for type-safe convenient access to fields within a Row.
-pub trait RowAccessor {
- fn get_bool(&self, i: usize) -> Result<bool>;
- fn get_byte(&self, i: usize) -> Result<i8>;
- fn get_short(&self, i: usize) -> Result<i16>;
- fn get_int(&self, i: usize) -> Result<i32>;
- fn get_long(&self, i: usize) -> Result<i64>;
- fn get_ubyte(&self, i: usize) -> Result<u8>;
- fn get_ushort(&self, i: usize) -> Result<u16>;
- fn get_uint(&self, i: usize) -> Result<u32>;
- fn get_ulong(&self, i: usize) -> Result<u64>;
- fn get_float(&self, i: usize) -> Result<f32>;
- fn get_double(&self, i: usize) -> Result<f64>;
- fn get_timestamp_millis(&self, i: usize) -> Result<u64>;
- fn get_timestamp_micros(&self, i: usize) -> Result<u64>;
- fn get_decimal(&self, i: usize) -> Result<&Decimal>;
- fn get_string(&self, i: usize) -> Result<&String>;
- fn get_bytes(&self, i: usize) -> Result<&ByteArray>;
- fn get_group(&self, i: usize) -> Result<&Row>;
- fn get_list(&self, i: usize) -> Result<&List>;
- fn get_map(&self, i: usize) -> Result<&Map>;
-}
-
-/// Trait for formating fields within a Row.
-pub trait RowFormatter {
- fn fmt(&self, i: usize) -> &dyn fmt::Display;
-}
-
-/// Macro to generate type-safe get_xxx methods for primitive types,
-/// e.g. `get_bool`, `get_short`.
-macro_rules! row_primitive_accessor {
- ($METHOD:ident, $VARIANT:ident, $TY:ty) => {
- fn $METHOD(&self, i: usize) -> Result<$TY> {
- match self.fields[i].1 {
- Field::$VARIANT(v) => Ok(v),
- _ => Err(general_err!(
- "Cannot access {} as {}",
- self.fields[i].1.get_type_name(),
- stringify!($VARIANT)
- )),
- }
- }
- };
-}
-
-/// Macro to generate type-safe get_xxx methods for reference types,
-/// e.g. `get_list`, `get_map`.
-macro_rules! row_complex_accessor {
- ($METHOD:ident, $VARIANT:ident, $TY:ty) => {
- fn $METHOD(&self, i: usize) -> Result<&$TY> {
- match self.fields[i].1 {
- Field::$VARIANT(ref v) => Ok(v),
- _ => Err(general_err!(
- "Cannot access {} as {}",
- self.fields[i].1.get_type_name(),
- stringify!($VARIANT)
- )),
- }
- }
- };
-}
-
-impl RowFormatter for Row {
- /// Get Display reference for a given field.
- fn fmt(&self, i: usize) -> &dyn fmt::Display {
- &self.fields[i].1
- }
-}
-
-impl RowAccessor for Row {
- row_primitive_accessor!(get_bool, Bool, bool);
-
- row_primitive_accessor!(get_byte, Byte, i8);
-
- row_primitive_accessor!(get_short, Short, i16);
-
- row_primitive_accessor!(get_int, Int, i32);
-
- row_primitive_accessor!(get_long, Long, i64);
-
- row_primitive_accessor!(get_ubyte, UByte, u8);
-
- row_primitive_accessor!(get_ushort, UShort, u16);
-
- row_primitive_accessor!(get_uint, UInt, u32);
-
- row_primitive_accessor!(get_ulong, ULong, u64);
-
- row_primitive_accessor!(get_float, Float, f32);
-
- row_primitive_accessor!(get_double, Double, f64);
-
- row_primitive_accessor!(get_timestamp_millis, TimestampMillis, u64);
-
- row_primitive_accessor!(get_timestamp_micros, TimestampMicros, u64);
-
- row_complex_accessor!(get_decimal, Decimal, Decimal);
-
- row_complex_accessor!(get_string, Str, String);
-
- row_complex_accessor!(get_bytes, Bytes, ByteArray);
-
- row_complex_accessor!(get_group, Group, Row);
-
- row_complex_accessor!(get_list, ListInternal, List);
-
- row_complex_accessor!(get_map, MapInternal, Map);
-}
-
-/// Constructs a `Row` from the list of `fields` and returns it.
-#[inline]
-pub fn make_row(fields: Vec<(String, Field)>) -> Row {
- Row { fields }
-}
-
-impl fmt::Display for Row {
- fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
- write!(f, "{{")?;
- for (i, &(ref key, ref value)) in self.fields.iter().enumerate() {
- key.fmt(f)?;
- write!(f, ": ")?;
- value.fmt(f)?;
- if i < self.fields.len() - 1 {
- write!(f, ", ")?;
- }
- }
- write!(f, "}}")
- }
-}
-
-/// `List` represents a list which contains an array of elements.
-#[derive(Clone, Debug, PartialEq)]
-pub struct List {
- elements: Vec<Field>,
-}
-
-#[allow(clippy::len_without_is_empty)]
-impl List {
- /// Get the number of fields in this row
- pub fn len(&self) -> usize {
- self.elements.len()
- }
-
- pub fn elements(&self) -> &[Field] {
- self.elements.as_slice()
- }
-}
-
-/// Constructs a `List` from the list of `fields` and returns it.
-#[inline]
-pub fn make_list(elements: Vec<Field>) -> List {
- List { elements }
-}
-
-/// Trait for type-safe access of an index for a `List`.
-/// Note that the get_XXX methods do not do bound checking.
-pub trait ListAccessor {
- fn get_bool(&self, i: usize) -> Result<bool>;
- fn get_byte(&self, i: usize) -> Result<i8>;
- fn get_short(&self, i: usize) -> Result<i16>;
- fn get_int(&self, i: usize) -> Result<i32>;
- fn get_long(&self, i: usize) -> Result<i64>;
- fn get_ubyte(&self, i: usize) -> Result<u8>;
- fn get_ushort(&self, i: usize) -> Result<u16>;
- fn get_uint(&self, i: usize) -> Result<u32>;
- fn get_ulong(&self, i: usize) -> Result<u64>;
- fn get_float(&self, i: usize) -> Result<f32>;
- fn get_double(&self, i: usize) -> Result<f64>;
- fn get_timestamp_millis(&self, i: usize) -> Result<u64>;
- fn get_timestamp_micros(&self, i: usize) -> Result<u64>;
- fn get_decimal(&self, i: usize) -> Result<&Decimal>;
- fn get_string(&self, i: usize) -> Result<&String>;
- fn get_bytes(&self, i: usize) -> Result<&ByteArray>;
- fn get_group(&self, i: usize) -> Result<&Row>;
- fn get_list(&self, i: usize) -> Result<&List>;
- fn get_map(&self, i: usize) -> Result<&Map>;
-}
-
-/// Macro to generate type-safe get_xxx methods for primitive types,
-/// e.g. get_bool, get_short
-macro_rules! list_primitive_accessor {
- ($METHOD:ident, $VARIANT:ident, $TY:ty) => {
- fn $METHOD(&self, i: usize) -> Result<$TY> {
- match self.elements[i] {
- Field::$VARIANT(v) => Ok(v),
- _ => Err(general_err!(
- "Cannot access {} as {}",
- self.elements[i].get_type_name(),
- stringify!($VARIANT)
- )),
- }
- }
- };
-}
-
-/// Macro to generate type-safe get_xxx methods for reference types
-/// e.g. get_list, get_map
-macro_rules! list_complex_accessor {
- ($METHOD:ident, $VARIANT:ident, $TY:ty) => {
- fn $METHOD(&self, i: usize) -> Result<&$TY> {
- match self.elements[i] {
- Field::$VARIANT(ref v) => Ok(v),
- _ => Err(general_err!(
- "Cannot access {} as {}",
- self.elements[i].get_type_name(),
- stringify!($VARIANT)
- )),
- }
- }
- };
-}
-
-impl ListAccessor for List {
- list_primitive_accessor!(get_bool, Bool, bool);
-
- list_primitive_accessor!(get_byte, Byte, i8);
-
- list_primitive_accessor!(get_short, Short, i16);
-
- list_primitive_accessor!(get_int, Int, i32);
-
- list_primitive_accessor!(get_long, Long, i64);
-
- list_primitive_accessor!(get_ubyte, UByte, u8);
-
- list_primitive_accessor!(get_ushort, UShort, u16);
-
- list_primitive_accessor!(get_uint, UInt, u32);
-
- list_primitive_accessor!(get_ulong, ULong, u64);
-
- list_primitive_accessor!(get_float, Float, f32);
-
- list_primitive_accessor!(get_double, Double, f64);
-
- list_primitive_accessor!(get_timestamp_millis, TimestampMillis, u64);
-
- list_primitive_accessor!(get_timestamp_micros, TimestampMicros, u64);
-
- list_complex_accessor!(get_decimal, Decimal, Decimal);
-
- list_complex_accessor!(get_string, Str, String);
-
- list_complex_accessor!(get_bytes, Bytes, ByteArray);
-
- list_complex_accessor!(get_group, Group, Row);
-
- list_complex_accessor!(get_list, ListInternal, List);
-
- list_complex_accessor!(get_map, MapInternal, Map);
-}
-
-/// `Map` represents a map which contains a list of key->value pairs.
-#[derive(Clone, Debug, PartialEq)]
-pub struct Map {
- entries: Vec<(Field, Field)>,
-}
-
-#[allow(clippy::len_without_is_empty)]
-impl Map {
- /// Get the number of fields in this row
- pub fn len(&self) -> usize {
- self.entries.len()
- }
-
- pub fn entries(&self) -> &[(Field, Field)] {
- self.entries.as_slice()
- }
-}
-
-/// Constructs a `Map` from the list of `entries` and returns it.
-#[inline]
-pub fn make_map(entries: Vec<(Field, Field)>) -> Map {
- Map { entries }
-}
-
-/// Trait for type-safe access of an index for a `Map`
-pub trait MapAccessor {
- fn get_keys<'a>(&'a self) -> Box<dyn ListAccessor + 'a>;
- fn get_values<'a>(&'a self) -> Box<dyn ListAccessor + 'a>;
-}
-
-struct MapList<'a> {
- elements: Vec<&'a Field>,
-}
-
-/// Macro to generate type-safe get_xxx methods for primitive types,
-/// e.g. get_bool, get_short
-macro_rules! map_list_primitive_accessor {
- ($METHOD:ident, $VARIANT:ident, $TY:ty) => {
- fn $METHOD(&self, i: usize) -> Result<$TY> {
- match self.elements[i] {
- Field::$VARIANT(v) => Ok(*v),
- _ => Err(general_err!(
- "Cannot access {} as {}",
- self.elements[i].get_type_name(),
- stringify!($VARIANT)
- )),
- }
- }
- };
-}
-
-impl<'a> ListAccessor for MapList<'a> {
- map_list_primitive_accessor!(get_bool, Bool, bool);
-
- map_list_primitive_accessor!(get_byte, Byte, i8);
-
- map_list_primitive_accessor!(get_short, Short, i16);
-
- map_list_primitive_accessor!(get_int, Int, i32);
-
- map_list_primitive_accessor!(get_long, Long, i64);
-
- map_list_primitive_accessor!(get_ubyte, UByte, u8);
-
- map_list_primitive_accessor!(get_ushort, UShort, u16);
-
- map_list_primitive_accessor!(get_uint, UInt, u32);
-
- map_list_primitive_accessor!(get_ulong, ULong, u64);
-
- map_list_primitive_accessor!(get_float, Float, f32);
-
- map_list_primitive_accessor!(get_double, Double, f64);
-
- map_list_primitive_accessor!(get_timestamp_millis, TimestampMillis, u64);
-
- map_list_primitive_accessor!(get_timestamp_micros, TimestampMicros, u64);
-
- list_complex_accessor!(get_decimal, Decimal, Decimal);
-
- list_complex_accessor!(get_string, Str, String);
-
- list_complex_accessor!(get_bytes, Bytes, ByteArray);
-
- list_complex_accessor!(get_group, Group, Row);
-
- list_complex_accessor!(get_list, ListInternal, List);
-
- list_complex_accessor!(get_map, MapInternal, Map);
-}
-
-impl MapAccessor for Map {
- fn get_keys<'a>(&'a self) -> Box<dyn ListAccessor + 'a> {
- let map_list = MapList {
- elements: self.entries.iter().map(|v| &v.0).collect(),
- };
- Box::new(map_list)
- }
-
- fn get_values<'a>(&'a self) -> Box<dyn ListAccessor + 'a> {
- let map_list = MapList {
- elements: self.entries.iter().map(|v| &v.1).collect(),
- };
- Box::new(map_list)
- }
-}
-
-/// API to represent a single field in a `Row`.
-#[derive(Clone, Debug, PartialEq)]
-pub enum Field {
- // Primitive types
- /// Null value.
- Null,
- /// Boolean value (`true`, `false`).
- Bool(bool),
- /// Signed integer INT_8.
- Byte(i8),
- /// Signed integer INT_16.
- Short(i16),
- /// Signed integer INT_32.
- Int(i32),
- /// Signed integer INT_64.
- Long(i64),
- // Unsigned integer UINT_8.
- UByte(u8),
- // Unsigned integer UINT_16.
- UShort(u16),
- // Unsigned integer UINT_32.
- UInt(u32),
- // Unsigned integer UINT_64.
- ULong(u64),
- /// IEEE 32-bit floating point value.
- Float(f32),
- /// IEEE 64-bit floating point value.
- Double(f64),
- /// Decimal value.
- Decimal(Decimal),
- /// UTF-8 encoded character string.
- Str(String),
- /// General binary value.
- Bytes(ByteArray),
- /// Date without a time of day, stores the number of days from the
- /// Unix epoch, 1 January 1970.
- Date(u32),
- /// Milliseconds from the Unix epoch, 1 January 1970.
- TimestampMillis(u64),
- /// Microseconds from the Unix epoch, 1 Janiary 1970.
- TimestampMicros(u64),
-
- // ----------------------------------------------------------------------
- // Complex types
- /// Struct, child elements are tuples of field-value pairs.
- Group(Row),
- /// List of elements.
- ListInternal(List),
- /// List of key-value pairs.
- MapInternal(Map),
-}
-
-impl Field {
- /// Get the type name.
- fn get_type_name(&self) -> &'static str {
- match *self {
- Field::Null => "Null",
- Field::Bool(_) => "Bool",
- Field::Byte(_) => "Byte",
- Field::Short(_) => "Short",
- Field::Int(_) => "Int",
- Field::Long(_) => "Long",
- Field::UByte(_) => "UByte",
- Field::UShort(_) => "UShort",
- Field::UInt(_) => "UInt",
- Field::ULong(_) => "ULong",
- Field::Float(_) => "Float",
- Field::Double(_) => "Double",
- Field::Decimal(_) => "Decimal",
- Field::Date(_) => "Date",
- Field::Str(_) => "Str",
- Field::Bytes(_) => "Bytes",
- Field::TimestampMillis(_) => "TimestampMillis",
- Field::TimestampMicros(_) => "TimestampMicros",
- Field::Group(_) => "Group",
- Field::ListInternal(_) => "ListInternal",
- Field::MapInternal(_) => "MapInternal",
- }
- }
-
- /// Determines if this Row represents a primitive value.
- pub fn is_primitive(&self) -> bool {
- !matches!(
- *self,
- Field::Group(_) | Field::ListInternal(_) | Field::MapInternal(_)
- )
- }
-
- /// Converts Parquet BOOLEAN type with logical type into `bool` value.
- #[inline]
- pub fn convert_bool(_descr: &ColumnDescPtr, value: bool) -> Self {
- Field::Bool(value)
- }
-
- /// Converts Parquet INT32 type with converted type into `i32` value.
- #[inline]
- pub fn convert_int32(descr: &ColumnDescPtr, value: i32) -> Self {
- match descr.converted_type() {
- ConvertedType::INT_8 => Field::Byte(value as i8),
- ConvertedType::INT_16 => Field::Short(value as i16),
- ConvertedType::INT_32 | ConvertedType::NONE => Field::Int(value),
- ConvertedType::UINT_8 => Field::UByte(value as u8),
- ConvertedType::UINT_16 => Field::UShort(value as u16),
- ConvertedType::UINT_32 => Field::UInt(value as u32),
- ConvertedType::DATE => Field::Date(value as u32),
- ConvertedType::DECIMAL => Field::Decimal(Decimal::from_i32(
- value,
- descr.type_precision(),
- descr.type_scale(),
- )),
- _ => nyi!(descr, value),
- }
- }
-
- /// Converts Parquet INT64 type with converted type into `i64` value.
- #[inline]
- pub fn convert_int64(descr: &ColumnDescPtr, value: i64) -> Self {
- match descr.converted_type() {
- ConvertedType::INT_64 | ConvertedType::NONE => Field::Long(value),
- ConvertedType::UINT_64 => Field::ULong(value as u64),
- ConvertedType::TIMESTAMP_MILLIS => Field::TimestampMillis(value as u64),
- ConvertedType::TIMESTAMP_MICROS => Field::TimestampMicros(value as u64),
- ConvertedType::DECIMAL => Field::Decimal(Decimal::from_i64(
- value,
- descr.type_precision(),
- descr.type_scale(),
- )),
- _ => nyi!(descr, value),
- }
- }
-
- /// Converts Parquet INT96 (nanosecond timestamps) type and logical type into
- /// `Timestamp` value.
- #[inline]
- pub fn convert_int96(_descr: &ColumnDescPtr, value: Int96) -> Self {
- Field::TimestampMillis(value.to_i64() as u64)
- }
-
- /// Converts Parquet FLOAT type with logical type into `f32` value.
- #[inline]
- pub fn convert_float(_descr: &ColumnDescPtr, value: f32) -> Self {
- Field::Float(value)
- }
-
- /// Converts Parquet DOUBLE type with converted type into `f64` value.
- #[inline]
- pub fn convert_double(_descr: &ColumnDescPtr, value: f64) -> Self {
- Field::Double(value)
- }
-
- /// Converts Parquet BYTE_ARRAY type with converted type into either UTF8 string or
- /// array of bytes.
- #[inline]
- pub fn convert_byte_array(descr: &ColumnDescPtr, value: ByteArray) -> Self {
- match descr.physical_type() {
- PhysicalType::BYTE_ARRAY => match descr.converted_type() {
- ConvertedType::UTF8 | ConvertedType::ENUM | ConvertedType::JSON => {
- let value = String::from_utf8(value.data().to_vec()).unwrap();
- Field::Str(value)
- }
- ConvertedType::BSON | ConvertedType::NONE => Field::Bytes(value),
- ConvertedType::DECIMAL => Field::Decimal(Decimal::from_bytes(
- value,
- descr.type_precision(),
- descr.type_scale(),
- )),
- _ => nyi!(descr, value),
- },
- PhysicalType::FIXED_LEN_BYTE_ARRAY => match descr.converted_type() {
- ConvertedType::DECIMAL => Field::Decimal(Decimal::from_bytes(
- value,
- descr.type_precision(),
- descr.type_scale(),
- )),
- ConvertedType::NONE => Field::Bytes(value),
- _ => nyi!(descr, value),
- },
- _ => nyi!(descr, value),
- }
- }
-
- #[cfg(feature = "cli")]
- pub fn to_json_value(&self) -> Value {
- match &self {
- Field::Null => Value::Null,
- Field::Bool(b) => Value::Bool(*b),
- Field::Byte(n) => Value::Number(serde_json::Number::from(*n)),
- Field::Short(n) => Value::Number(serde_json::Number::from(*n)),
- Field::Int(n) => Value::Number(serde_json::Number::from(*n)),
- Field::Long(n) => Value::Number(serde_json::Number::from(*n)),
- Field::UByte(n) => Value::Number(serde_json::Number::from(*n)),
- Field::UShort(n) => Value::Number(serde_json::Number::from(*n)),
- Field::UInt(n) => Value::Number(serde_json::Number::from(*n)),
- Field::ULong(n) => Value::Number(serde_json::Number::from(*n)),
- Field::Float(n) => serde_json::Number::from_f64(f64::from(*n))
- .map(Value::Number)
- .unwrap_or(Value::Null),
- Field::Double(n) => serde_json::Number::from_f64(*n)
- .map(Value::Number)
- .unwrap_or(Value::Null),
- Field::Decimal(n) => Value::String(convert_decimal_to_string(&n)),
- Field::Str(s) => Value::String(s.to_owned()),
- Field::Bytes(b) => Value::String(base64::encode(b.data())),
- Field::Date(d) => Value::String(convert_date_to_string(*d)),
- Field::TimestampMillis(ts) => {
- Value::String(convert_timestamp_millis_to_string(*ts))
- }
- Field::TimestampMicros(ts) => {
- Value::String(convert_timestamp_micros_to_string(*ts))
- }
- Field::Group(row) => row.to_json_value(),
- Field::ListInternal(fields) => {
- Value::Array(fields.elements.iter().map(|f| f.to_json_value()).collect())
- }
- Field::MapInternal(map) => Value::Object(
- map.entries
- .iter()
- .map(|(key_field, value_field)| {
- let key_val = key_field.to_json_value();
- let key_str = key_val
- .as_str()
- .map(|s| s.to_owned())
- .unwrap_or_else(|| key_val.to_string());
- (key_str, value_field.to_json_value())
- })
- .collect(),
- ),
- }
- }
-}
-
-impl fmt::Display for Field {
- fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
- match *self {
- Field::Null => write!(f, "null"),
- Field::Bool(value) => write!(f, "{}", value),
- Field::Byte(value) => write!(f, "{}", value),
- Field::Short(value) => write!(f, "{}", value),
- Field::Int(value) => write!(f, "{}", value),
- Field::Long(value) => write!(f, "{}", value),
- Field::UByte(value) => write!(f, "{}", value),
- Field::UShort(value) => write!(f, "{}", value),
- Field::UInt(value) => write!(f, "{}", value),
- Field::ULong(value) => write!(f, "{}", value),
- Field::Float(value) => {
- if !(1e-15..=1e19).contains(&value) {
- write!(f, "{:E}", value)
- } else {
- write!(f, "{:?}", value)
- }
- }
- Field::Double(value) => {
- if !(1e-15..=1e19).contains(&value) {
- write!(f, "{:E}", value)
- } else {
- write!(f, "{:?}", value)
- }
- }
- Field::Decimal(ref value) => {
- write!(f, "{}", convert_decimal_to_string(value))
- }
- Field::Str(ref value) => write!(f, "\"{}\"", value),
- Field::Bytes(ref value) => write!(f, "{:?}", value.data()),
- Field::Date(value) => write!(f, "{}", convert_date_to_string(value)),
- Field::TimestampMillis(value) => {
- write!(f, "{}", convert_timestamp_millis_to_string(value))
- }
- Field::TimestampMicros(value) => {
- write!(f, "{}", convert_timestamp_micros_to_string(value))
- }
- Field::Group(ref fields) => write!(f, "{}", fields),
- Field::ListInternal(ref list) => {
- let elems = &list.elements;
- write!(f, "[")?;
- for (i, field) in elems.iter().enumerate() {
- field.fmt(f)?;
- if i < elems.len() - 1 {
- write!(f, ", ")?;
- }
- }
- write!(f, "]")
- }
- Field::MapInternal(ref map) => {
- let entries = &map.entries;
- write!(f, "{{")?;
- for (i, &(ref key, ref value)) in entries.iter().enumerate() {
- key.fmt(f)?;
- write!(f, " -> ")?;
- value.fmt(f)?;
- if i < entries.len() - 1 {
- write!(f, ", ")?;
- }
- }
- write!(f, "}}")
- }
- }
- }
-}
-
-/// Helper method to convert Parquet date into a string.
-/// Input `value` is a number of days since the epoch in UTC.
-/// Date is displayed in local timezone.
-#[inline]
-fn convert_date_to_string(value: u32) -> String {
- static NUM_SECONDS_IN_DAY: i64 = 60 * 60 * 24;
- let dt = Utc.timestamp(value as i64 * NUM_SECONDS_IN_DAY, 0).date();
- format!("{}", dt.format("%Y-%m-%d %:z"))
-}
-
-/// Helper method to convert Parquet timestamp into a string.
-/// Input `value` is a number of milliseconds since the epoch in UTC.
-/// Datetime is displayed in local timezone.
-#[inline]
-fn convert_timestamp_millis_to_string(value: u64) -> String {
- let dt = Utc.timestamp((value / 1000) as i64, 0);
- format!("{}", dt.format("%Y-%m-%d %H:%M:%S %:z"))
-}
-
-/// Helper method to convert Parquet timestamp into a string.
-/// Input `value` is a number of microseconds since the epoch in UTC.
-/// Datetime is displayed in local timezone.
-#[inline]
-fn convert_timestamp_micros_to_string(value: u64) -> String {
- convert_timestamp_millis_to_string(value / 1000)
-}
-
-/// Helper method to convert Parquet decimal into a string.
-/// We assert that `scale >= 0` and `precision > scale`, but this will be enforced
-/// when constructing Parquet schema.
-#[inline]
-fn convert_decimal_to_string(decimal: &Decimal) -> String {
- assert!(decimal.scale() >= 0 && decimal.precision() > decimal.scale());
-
- // Specify as signed bytes to resolve sign as part of conversion.
- let num = BigInt::from_signed_bytes_be(decimal.data());
-
- // Offset of the first digit in a string.
- let negative = if num.sign() == Sign::Minus { 1 } else { 0 };
- let mut num_str = num.to_string();
- let mut point = num_str.len() as i32 - decimal.scale() - negative;
-
- // Convert to string form without scientific notation.
- if point <= 0 {
- // Zeros need to be prepended to the unscaled value.
- while point < 0 {
- num_str.insert(negative as usize, '0');
- point += 1;
- }
- num_str.insert_str(negative as usize, "0.");
- } else {
- // No zeroes need to be prepended to the unscaled value, simply insert decimal
- // point.
- num_str.insert((point + negative) as usize, '.');
- }
-
- num_str
-}
-
-#[cfg(test)]
-#[allow(clippy::approx_constant, clippy::many_single_char_names)]
-mod tests {
- use super::*;
-
- use std::sync::Arc;
-
- use crate::schema::types::{ColumnDescriptor, ColumnPath, PrimitiveTypeBuilder};
-
- /// Creates test column descriptor based on provided type parameters.
- macro_rules! make_column_descr {
- ($physical_type:expr, $logical_type:expr) => {{
- let tpe = PrimitiveTypeBuilder::new("col", $physical_type)
- .with_converted_type($logical_type)
- .build()
- .unwrap();
- Arc::new(ColumnDescriptor::new(
- Arc::new(tpe),
- 0,
- 0,
- ColumnPath::from("col"),
- ))
- }};
- ($physical_type:expr, $logical_type:expr, $len:expr, $prec:expr, $scale:expr) => {{
- let tpe = PrimitiveTypeBuilder::new("col", $physical_type)
- .with_converted_type($logical_type)
- .with_length($len)
- .with_precision($prec)
- .with_scale($scale)
- .build()
- .unwrap();
- Arc::new(ColumnDescriptor::new(
- Arc::new(tpe),
- 0,
- 0,
- ColumnPath::from("col"),
- ))
- }};
- }
-
- #[test]
- fn test_row_convert_bool() {
- // BOOLEAN value does not depend on logical type
- let descr = make_column_descr![PhysicalType::BOOLEAN, ConvertedType::NONE];
-
- let row = Field::convert_bool(&descr, true);
- assert_eq!(row, Field::Bool(true));
-
- let row = Field::convert_bool(&descr, false);
- assert_eq!(row, Field::Bool(false));
- }
-
- #[test]
- fn test_row_convert_int32() {
- let descr = make_column_descr![PhysicalType::INT32, ConvertedType::INT_8];
- let row = Field::convert_int32(&descr, 111);
- assert_eq!(row, Field::Byte(111));
-
- let descr = make_column_descr![PhysicalType::INT32, ConvertedType::INT_16];
- let row = Field::convert_int32(&descr, 222);
- assert_eq!(row, Field::Short(222));
-
- let descr = make_column_descr![PhysicalType::INT32, ConvertedType::INT_32];
- let row = Field::convert_int32(&descr, 333);
- assert_eq!(row, Field::Int(333));
-
- let descr = make_column_descr![PhysicalType::INT32, ConvertedType::UINT_8];
- let row = Field::convert_int32(&descr, -1);
- assert_eq!(row, Field::UByte(255));
-
- let descr = make_column_descr![PhysicalType::INT32, ConvertedType::UINT_16];
- let row = Field::convert_int32(&descr, 256);
- assert_eq!(row, Field::UShort(256));
-
- let descr = make_column_descr![PhysicalType::INT32, ConvertedType::UINT_32];
- let row = Field::convert_int32(&descr, 1234);
- assert_eq!(row, Field::UInt(1234));
-
- let descr = make_column_descr![PhysicalType::INT32, ConvertedType::NONE];
- let row = Field::convert_int32(&descr, 444);
- assert_eq!(row, Field::Int(444));
-
- let descr = make_column_descr![PhysicalType::INT32, ConvertedType::DATE];
- let row = Field::convert_int32(&descr, 14611);
- assert_eq!(row, Field::Date(14611));
-
- let descr =
- make_column_descr![PhysicalType::INT32, ConvertedType::DECIMAL, 0, 8, 2];
- let row = Field::convert_int32(&descr, 444);
- assert_eq!(row, Field::Decimal(Decimal::from_i32(444, 8, 2)));
- }
-
- #[test]
- fn test_row_convert_int64() {
- let descr = make_column_descr![PhysicalType::INT64, ConvertedType::INT_64];
- let row = Field::convert_int64(&descr, 1111);
- assert_eq!(row, Field::Long(1111));
-
- let descr = make_column_descr![PhysicalType::INT64, ConvertedType::UINT_64];
- let row = Field::convert_int64(&descr, 78239823);
- assert_eq!(row, Field::ULong(78239823));
-
- let descr =
- make_column_descr![PhysicalType::INT64, ConvertedType::TIMESTAMP_MILLIS];
- let row = Field::convert_int64(&descr, 1541186529153);
- assert_eq!(row, Field::TimestampMillis(1541186529153));
-
- let descr =
- make_column_descr![PhysicalType::INT64, ConvertedType::TIMESTAMP_MICROS];
- let row = Field::convert_int64(&descr, 1541186529153123);
- assert_eq!(row, Field::TimestampMicros(1541186529153123));
-
- let descr = make_column_descr![PhysicalType::INT64, ConvertedType::NONE];
- let row = Field::convert_int64(&descr, 2222);
- assert_eq!(row, Field::Long(2222));
-
- let descr =
- make_column_descr![PhysicalType::INT64, ConvertedType::DECIMAL, 0, 8, 2];
- let row = Field::convert_int64(&descr, 3333);
- assert_eq!(row, Field::Decimal(Decimal::from_i64(3333, 8, 2)));
- }
-
- #[test]
- fn test_row_convert_int96() {
- // INT96 value does not depend on logical type
- let descr = make_column_descr![PhysicalType::INT96, ConvertedType::NONE];
-
- let value = Int96::from(vec![0, 0, 2454923]);
- let row = Field::convert_int96(&descr, value);
- assert_eq!(row, Field::TimestampMillis(1238544000000));
-
- let value = Int96::from(vec![4165425152, 13, 2454923]);
- let row = Field::convert_int96(&descr, value);
- assert_eq!(row, Field::TimestampMillis(1238544060000));
- }
-
- #[test]
- fn test_row_convert_float() {
- // FLOAT value does not depend on logical type
- let descr = make_column_descr![PhysicalType::FLOAT, ConvertedType::NONE];
- let row = Field::convert_float(&descr, 2.31);
- assert_eq!(row, Field::Float(2.31));
- }
-
- #[test]
- fn test_row_convert_double() {
- // DOUBLE value does not depend on logical type
- let descr = make_column_descr![PhysicalType::DOUBLE, ConvertedType::NONE];
- let row = Field::convert_double(&descr, 1.56);
- assert_eq!(row, Field::Double(1.56));
- }
-
- #[test]
- fn test_row_convert_byte_array() {
- // UTF8
- let descr = make_column_descr![PhysicalType::BYTE_ARRAY, ConvertedType::UTF8];
- let value = ByteArray::from(vec![b'A', b'B', b'C', b'D']);
- let row = Field::convert_byte_array(&descr, value);
- assert_eq!(row, Field::Str("ABCD".to_string()));
-
- // ENUM
- let descr = make_column_descr![PhysicalType::BYTE_ARRAY, ConvertedType::ENUM];
- let value = ByteArray::from(vec![b'1', b'2', b'3']);
- let row = Field::convert_byte_array(&descr, value);
- assert_eq!(row, Field::Str("123".to_string()));
-
- // JSON
- let descr = make_column_descr![PhysicalType::BYTE_ARRAY, ConvertedType::JSON];
- let value = ByteArray::from(vec![b'{', b'"', b'a', b'"', b':', b'1', b'}']);
- let row = Field::convert_byte_array(&descr, value);
- assert_eq!(row, Field::Str("{\"a\":1}".to_string()));
-
- // NONE
- let descr = make_column_descr![PhysicalType::BYTE_ARRAY, ConvertedType::NONE];
- let value = ByteArray::from(vec![1, 2, 3, 4, 5]);
- let row = Field::convert_byte_array(&descr, value.clone());
- assert_eq!(row, Field::Bytes(value));
-
- // BSON
- let descr = make_column_descr![PhysicalType::BYTE_ARRAY, ConvertedType::BSON];
- let value = ByteArray::from(vec![1, 2, 3, 4, 5]);
- let row = Field::convert_byte_array(&descr, value.clone());
- assert_eq!(row, Field::Bytes(value));
-
- // DECIMAL
- let descr =
- make_column_descr![PhysicalType::BYTE_ARRAY, ConvertedType::DECIMAL, 0, 8, 2];
- let value = ByteArray::from(vec![207, 200]);
- let row = Field::convert_byte_array(&descr, value.clone());
- assert_eq!(row, Field::Decimal(Decimal::from_bytes(value, 8, 2)));
-
- // DECIMAL (FIXED_LEN_BYTE_ARRAY)
- let descr = make_column_descr![
- PhysicalType::FIXED_LEN_BYTE_ARRAY,
- ConvertedType::DECIMAL,
- 8,
- 17,
- 5
- ];
- let value = ByteArray::from(vec![0, 0, 0, 0, 0, 4, 147, 224]);
- let row = Field::convert_byte_array(&descr, value.clone());
- assert_eq!(row, Field::Decimal(Decimal::from_bytes(value, 17, 5)));
-
- // NONE (FIXED_LEN_BYTE_ARRAY)
- let descr = make_column_descr![
- PhysicalType::FIXED_LEN_BYTE_ARRAY,
- ConvertedType::NONE,
- 6,
- 0,
- 0
- ];
- let value = ByteArray::from(vec![1, 2, 3, 4, 5, 6]);
- let row = Field::convert_byte_array(&descr, value.clone());
- assert_eq!(row, Field::Bytes(value));
- }
-
- #[test]
- fn test_convert_date_to_string() {
- fn check_date_conversion(y: u32, m: u32, d: u32) {
- let datetime = chrono::NaiveDate::from_ymd(y as i32, m, d).and_hms(0, 0, 0);
- let dt = Utc.from_utc_datetime(&datetime);
- let res = convert_date_to_string((dt.timestamp() / 60 / 60 / 24) as u32);
- let exp = format!("{}", dt.format("%Y-%m-%d %:z"));
- assert_eq!(res, exp);
- }
-
- check_date_conversion(2010, 1, 2);
- check_date_conversion(2014, 5, 1);
- check_date_conversion(2016, 2, 29);
- check_date_conversion(2017, 9, 12);
- check_date_conversion(2018, 3, 31);
- }
-
- #[test]
- fn test_convert_timestamp_to_string() {
- fn check_datetime_conversion(y: u32, m: u32, d: u32, h: u32, mi: u32, s: u32) {
- let datetime = chrono::NaiveDate::from_ymd(y as i32, m, d).and_hms(h, mi, s);
- let dt = Utc.from_utc_datetime(&datetime);
- let res = convert_timestamp_millis_to_string(dt.timestamp_millis() as u64);
- let exp = format!("{}", dt.format("%Y-%m-%d %H:%M:%S %:z"));
- assert_eq!(res, exp);
- }
-
- check_datetime_conversion(2010, 1, 2, 13, 12, 54);
- check_datetime_conversion(2011, 1, 3, 8, 23, 1);
- check_datetime_conversion(2012, 4, 5, 11, 6, 32);
- check_datetime_conversion(2013, 5, 12, 16, 38, 0);
- check_datetime_conversion(2014, 11, 28, 21, 15, 12);
- }
-
- #[test]
- fn test_convert_float_to_string() {
- assert_eq!(format!("{}", Field::Float(1.0)), "1.0");
- assert_eq!(format!("{}", Field::Float(9.63)), "9.63");
- assert_eq!(format!("{}", Field::Float(1e-15)), "0.000000000000001");
- assert_eq!(format!("{}", Field::Float(1e-16)), "1E-16");
- assert_eq!(format!("{}", Field::Float(1e19)), "10000000000000000000.0");
- assert_eq!(format!("{}", Field::Float(1e20)), "1E20");
- assert_eq!(format!("{}", Field::Float(1.7976931E30)), "1.7976931E30");
- assert_eq!(format!("{}", Field::Float(-1.7976931E30)), "-1.7976931E30");
- }
-
- #[test]
- fn test_convert_double_to_string() {
- assert_eq!(format!("{}", Field::Double(1.0)), "1.0");
- assert_eq!(format!("{}", Field::Double(9.63)), "9.63");
- assert_eq!(format!("{}", Field::Double(1e-15)), "0.000000000000001");
- assert_eq!(format!("{}", Field::Double(1e-16)), "1E-16");
- assert_eq!(format!("{}", Field::Double(1e19)), "10000000000000000000.0");
- assert_eq!(format!("{}", Field::Double(1e20)), "1E20");
- assert_eq!(
- format!("{}", Field::Double(1.79769313486E308)),
- "1.79769313486E308"
- );
- assert_eq!(
- format!("{}", Field::Double(-1.79769313486E308)),
- "-1.79769313486E308"
- );
- }
-
- #[test]
- fn test_convert_decimal_to_string() {
- // Helper method to compare decimal
- fn check_decimal(bytes: Vec<u8>, precision: i32, scale: i32, res: &str) {
- let decimal = Decimal::from_bytes(ByteArray::from(bytes), precision, scale);
- assert_eq!(convert_decimal_to_string(&decimal), res);
- }
-
- // This example previously used to fail in some engines
- check_decimal(
- vec![0, 0, 0, 0, 0, 0, 0, 0, 13, 224, 182, 179, 167, 100, 0, 0],
- 38,
- 18,
- "1.000000000000000000",
- );
- check_decimal(
- vec![
- 249, 233, 247, 16, 185, 192, 202, 223, 215, 165, 192, 166, 67, 72,
- ],
- 36,
- 28,
- "-12344.0242342304923409234234293432",
- );
- check_decimal(vec![0, 0, 0, 0, 0, 4, 147, 224], 17, 5, "3.00000");
- check_decimal(vec![0, 0, 0, 0, 1, 201, 195, 140], 18, 2, "300000.12");
- check_decimal(vec![207, 200], 10, 2, "-123.44");
- check_decimal(vec![207, 200], 10, 8, "-0.00012344");
- }
-
- #[test]
- fn test_row_display() {
- // Primitive types
- assert_eq!(format!("{}", Field::Null), "null");
- assert_eq!(format!("{}", Field::Bool(true)), "true");
- assert_eq!(format!("{}", Field::Bool(false)), "false");
- assert_eq!(format!("{}", Field::Byte(1)), "1");
- assert_eq!(format!("{}", Field::Short(2)), "2");
- assert_eq!(format!("{}", Field::Int(3)), "3");
- assert_eq!(format!("{}", Field::Long(4)), "4");
- assert_eq!(format!("{}", Field::UByte(1)), "1");
- assert_eq!(format!("{}", Field::UShort(2)), "2");
- assert_eq!(format!("{}", Field::UInt(3)), "3");
- assert_eq!(format!("{}", Field::ULong(4)), "4");
- assert_eq!(format!("{}", Field::Float(5.0)), "5.0");
- assert_eq!(format!("{}", Field::Float(5.1234)), "5.1234");
- assert_eq!(format!("{}", Field::Double(6.0)), "6.0");
- assert_eq!(format!("{}", Field::Double(6.1234)), "6.1234");
- assert_eq!(format!("{}", Field::Str("abc".to_string())), "\"abc\"");
- assert_eq!(
- format!("{}", Field::Bytes(ByteArray::from(vec![1, 2, 3]))),
- "[1, 2, 3]"
- );
- assert_eq!(
- format!("{}", Field::Date(14611)),
- convert_date_to_string(14611)
- );
- assert_eq!(
- format!("{}", Field::TimestampMillis(1262391174000)),
- convert_timestamp_millis_to_string(1262391174000)
- );
- assert_eq!(
- format!("{}", Field::TimestampMicros(1262391174000000)),
- convert_timestamp_micros_to_string(1262391174000000)
- );
- assert_eq!(
- format!("{}", Field::Decimal(Decimal::from_i32(4, 8, 2))),
- convert_decimal_to_string(&Decimal::from_i32(4, 8, 2))
- );
-
- // Complex types
- let fields = vec![
- ("x".to_string(), Field::Null),
- ("Y".to_string(), Field::Int(2)),
- ("z".to_string(), Field::Float(3.1)),
- ("a".to_string(), Field::Str("abc".to_string())),
- ];
- let row = Field::Group(make_row(fields));
- assert_eq!(format!("{}", row), "{x: null, Y: 2, z: 3.1, a: \"abc\"}");
-
- let row = Field::ListInternal(make_list(vec![
- Field::Int(2),
- Field::Int(1),
- Field::Null,
- Field::Int(12),
- ]));
- assert_eq!(format!("{}", row), "[2, 1, null, 12]");
-
- let row = Field::MapInternal(make_map(vec![
- (Field::Int(1), Field::Float(1.2)),
- (Field::Int(2), Field::Float(4.5)),
- (Field::Int(3), Field::Float(2.3)),
- ]));
- assert_eq!(format!("{}", row), "{1 -> 1.2, 2 -> 4.5, 3 -> 2.3}");
- }
-
- #[test]
- fn test_is_primitive() {
- // primitives
- assert!(Field::Null.is_primitive());
- assert!(Field::Bool(true).is_primitive());
- assert!(Field::Bool(false).is_primitive());
- assert!(Field::Byte(1).is_primitive());
- assert!(Field::Short(2).is_primitive());
- assert!(Field::Int(3).is_primitive());
- assert!(Field::Long(4).is_primitive());
- assert!(Field::UByte(1).is_primitive());
- assert!(Field::UShort(2).is_primitive());
- assert!(Field::UInt(3).is_primitive());
- assert!(Field::ULong(4).is_primitive());
- assert!(Field::Float(5.0).is_primitive());
- assert!(Field::Float(5.1234).is_primitive());
- assert!(Field::Double(6.0).is_primitive());
- assert!(Field::Double(6.1234).is_primitive());
- assert!(Field::Str("abc".to_string()).is_primitive());
- assert!(Field::Bytes(ByteArray::from(vec![1, 2, 3])).is_primitive());
- assert!(Field::TimestampMillis(12345678).is_primitive());
- assert!(Field::TimestampMicros(12345678901).is_primitive());
- assert!(Field::Decimal(Decimal::from_i32(4, 8, 2)).is_primitive());
-
- // complex types
- assert_eq!(
- false,
- Field::Group(make_row(vec![
- ("x".to_string(), Field::Null),
- ("Y".to_string(), Field::Int(2)),
- ("z".to_string(), Field::Float(3.1)),
- ("a".to_string(), Field::Str("abc".to_string()))
- ]))
- .is_primitive()
- );
-
- assert_eq!(
- false,
- Field::ListInternal(make_list(vec![
- Field::Int(2),
- Field::Int(1),
- Field::Null,
- Field::Int(12)
- ]))
- .is_primitive()
- );
-
- assert_eq!(
- false,
- Field::MapInternal(make_map(vec![
- (Field::Int(1), Field::Float(1.2)),
- (Field::Int(2), Field::Float(4.5)),
- (Field::Int(3), Field::Float(2.3))
- ]))
- .is_primitive()
- );
- }
-
- #[test]
- fn test_row_primitive_field_fmt() {
- // Primitives types
- let row = make_row(vec![
- ("00".to_string(), Field::Null),
- ("01".to_string(), Field::Bool(false)),
- ("02".to_string(), Field::Byte(3)),
- ("03".to_string(), Field::Short(4)),
- ("04".to_string(), Field::Int(5)),
- ("05".to_string(), Field::Long(6)),
- ("06".to_string(), Field::UByte(7)),
- ("07".to_string(), Field::UShort(8)),
- ("08".to_string(), Field::UInt(9)),
- ("09".to_string(), Field::ULong(10)),
- ("10".to_string(), Field::Float(11.1)),
- ("11".to_string(), Field::Double(12.1)),
- ("12".to_string(), Field::Str("abc".to_string())),
- (
- "13".to_string(),
- Field::Bytes(ByteArray::from(vec![1, 2, 3, 4, 5])),
- ),
- ("14".to_string(), Field::Date(14611)),
- ("15".to_string(), Field::TimestampMillis(1262391174000)),
- ("16".to_string(), Field::TimestampMicros(1262391174000000)),
- ("17".to_string(), Field::Decimal(Decimal::from_i32(4, 7, 2))),
- ]);
-
- assert_eq!("null", format!("{}", row.fmt(0)));
- assert_eq!("false", format!("{}", row.fmt(1)));
- assert_eq!("3", format!("{}", row.fmt(2)));
- assert_eq!("4", format!("{}", row.fmt(3)));
- assert_eq!("5", format!("{}", row.fmt(4)));
- assert_eq!("6", format!("{}", row.fmt(5)));
- assert_eq!("7", format!("{}", row.fmt(6)));
- assert_eq!("8", format!("{}", row.fmt(7)));
- assert_eq!("9", format!("{}", row.fmt(8)));
- assert_eq!("10", format!("{}", row.fmt(9)));
- assert_eq!("11.1", format!("{}", row.fmt(10)));
- assert_eq!("12.1", format!("{}", row.fmt(11)));
- assert_eq!("\"abc\"", format!("{}", row.fmt(12)));
- assert_eq!("[1, 2, 3, 4, 5]", format!("{}", row.fmt(13)));
- assert_eq!(convert_date_to_string(14611), format!("{}", row.fmt(14)));
- assert_eq!(
- convert_timestamp_millis_to_string(1262391174000),
- format!("{}", row.fmt(15))
- );
- assert_eq!(
- convert_timestamp_micros_to_string(1262391174000000),
- format!("{}", row.fmt(16))
- );
- assert_eq!("0.04", format!("{}", row.fmt(17)));
- }
-
- #[test]
- fn test_row_complex_field_fmt() {
- // Complex types
- let row = make_row(vec![
- (
- "00".to_string(),
- Field::Group(make_row(vec![
- ("x".to_string(), Field::Null),
- ("Y".to_string(), Field::Int(2)),
- ])),
- ),
- (
- "01".to_string(),
- Field::ListInternal(make_list(vec![
- Field::Int(2),
- Field::Int(1),
- Field::Null,
- Field::Int(12),
- ])),
- ),
- (
- "02".to_string(),
- Field::MapInternal(make_map(vec![
- (Field::Int(1), Field::Float(1.2)),
- (Field::Int(2), Field::Float(4.5)),
- (Field::Int(3), Field::Float(2.3)),
- ])),
- ),
- ]);
-
- assert_eq!("{x: null, Y: 2}", format!("{}", row.fmt(0)));
- assert_eq!("[2, 1, null, 12]", format!("{}", row.fmt(1)));
- assert_eq!("{1 -> 1.2, 2 -> 4.5, 3 -> 2.3}", format!("{}", row.fmt(2)));
- }
-
- #[test]
- fn test_row_primitive_accessors() {
- // primitives
- let row = make_row(vec![
- ("a".to_string(), Field::Null),
- ("b".to_string(), Field::Bool(false)),
- ("c".to_string(), Field::Byte(3)),
- ("d".to_string(), Field::Short(4)),
- ("e".to_string(), Field::Int(5)),
- ("f".to_string(), Field::Long(6)),
- ("g".to_string(), Field::UByte(3)),
- ("h".to_string(), Field::UShort(4)),
- ("i".to_string(), Field::UInt(5)),
- ("j".to_string(), Field::ULong(6)),
- ("k".to_string(), Field::Float(7.1)),
- ("l".to_string(), Field::Double(8.1)),
- ("m".to_string(), Field::Str("abc".to_string())),
- (
- "n".to_string(),
- Field::Bytes(ByteArray::from(vec![1, 2, 3, 4, 5])),
- ),
- ("o".to_string(), Field::Decimal(Decimal::from_i32(4, 7, 2))),
- ]);
-
- assert_eq!(false, row.get_bool(1).unwrap());
- assert_eq!(3, row.get_byte(2).unwrap());
- assert_eq!(4, row.get_short(3).unwrap());
- assert_eq!(5, row.get_int(4).unwrap());
- assert_eq!(6, row.get_long(5).unwrap());
- assert_eq!(3, row.get_ubyte(6).unwrap());
- assert_eq!(4, row.get_ushort(7).unwrap());
- assert_eq!(5, row.get_uint(8).unwrap());
- assert_eq!(6, row.get_ulong(9).unwrap());
- assert!(7.1 - row.get_float(10).unwrap() < f32::EPSILON);
- assert!(8.1 - row.get_double(11).unwrap() < f64::EPSILON);
- assert_eq!("abc", row.get_string(12).unwrap());
- assert_eq!(5, row.get_bytes(13).unwrap().len());
- assert_eq!(7, row.get_decimal(14).unwrap().precision());
- }
-
- #[test]
- fn test_row_primitive_invalid_accessors() {
- // primitives
- let row = make_row(vec![
- ("a".to_string(), Field::Null),
- ("b".to_string(), Field::Bool(false)),
- ("c".to_string(), Field::Byte(3)),
- ("d".to_string(), Field::Short(4)),
- ("e".to_string(), Field::Int(5)),
- ("f".to_string(), Field::Long(6)),
- ("g".to_string(), Field::UByte(3)),
- ("h".to_string(), Field::UShort(4)),
- ("i".to_string(), Field::UInt(5)),
- ("j".to_string(), Field::ULong(6)),
- ("k".to_string(), Field::Float(7.1)),
- ("l".to_string(), Field::Double(8.1)),
- ("m".to_string(), Field::Str("abc".to_string())),
- (
- "n".to_string(),
- Field::Bytes(ByteArray::from(vec![1, 2, 3, 4, 5])),
- ),
- ("o".to_string(), Field::Decimal(Decimal::from_i32(4, 7, 2))),
- ]);
-
- for i in 0..row.len() {
- assert!(row.get_group(i).is_err());
- }
- }
-
- #[test]
- fn test_row_complex_accessors() {
- let row = make_row(vec![
- (
- "a".to_string(),
- Field::Group(make_row(vec![
- ("x".to_string(), Field::Null),
- ("Y".to_string(), Field::Int(2)),
- ])),
- ),
- (
- "b".to_string(),
- Field::ListInternal(make_list(vec![
- Field::Int(2),
- Field::Int(1),
- Field::Null,
- Field::Int(12),
- ])),
- ),
- (
- "c".to_string(),
- Field::MapInternal(make_map(vec![
- (Field::Int(1), Field::Float(1.2)),
- (Field::Int(2), Field::Float(4.5)),
- (Field::Int(3), Field::Float(2.3)),
- ])),
- ),
- ]);
-
- assert_eq!(2, row.get_group(0).unwrap().len());
- assert_eq!(4, row.get_list(1).unwrap().len());
- assert_eq!(3, row.get_map(2).unwrap().len());
- }
-
- #[test]
- fn test_row_complex_invalid_accessors() {
- let row = make_row(vec![
- (
- "a".to_string(),
- Field::Group(make_row(vec![
- ("x".to_string(), Field::Null),
- ("Y".to_string(), Field::Int(2)),
- ])),
- ),
- (
- "b".to_string(),
- Field::ListInternal(make_list(vec![
- Field::Int(2),
- Field::Int(1),
- Field::Null,
- Field::Int(12),
- ])),
- ),
- (
- "c".to_string(),
- Field::MapInternal(make_map(vec![
- (Field::Int(1), Field::Float(1.2)),
- (Field::Int(2), Field::Float(4.5)),
- (Field::Int(3), Field::Float(2.3)),
- ])),
- ),
- ]);
-
- assert_eq!(
- ParquetError::General("Cannot access Group as Float".to_string()),
- row.get_float(0).unwrap_err()
- );
- assert_eq!(
- ParquetError::General("Cannot access ListInternal as Float".to_string()),
- row.get_float(1).unwrap_err()
- );
- assert_eq!(
- ParquetError::General("Cannot access MapInternal as Float".to_string()),
- row.get_float(2).unwrap_err()
- );
- }
-
- #[test]
- fn test_list_primitive_accessors() {
- // primitives
- let list = make_list(vec![Field::Bool(false)]);
- assert_eq!(false, list.get_bool(0).unwrap());
-
- let list = make_list(vec![Field::Byte(3), Field::Byte(4)]);
- assert_eq!(4, list.get_byte(1).unwrap());
-
- let list = make_list(vec![Field::Short(4), Field::Short(5), Field::Short(6)]);
- assert_eq!(6, list.get_short(2).unwrap());
-
- let list = make_list(vec![Field::Int(5)]);
- assert_eq!(5, list.get_int(0).unwrap());
-
- let list = make_list(vec![Field::Long(6), Field::Long(7)]);
- assert_eq!(7, list.get_long(1).unwrap());
-
- let list = make_list(vec![Field::UByte(3), Field::UByte(4)]);
- assert_eq!(4, list.get_ubyte(1).unwrap());
-
- let list = make_list(vec![Field::UShort(4), Field::UShort(5), Field::UShort(6)]);
- assert_eq!(6, list.get_ushort(2).unwrap());
-
- let list = make_list(vec![Field::UInt(5)]);
- assert_eq!(5, list.get_uint(0).unwrap());
-
- let list = make_list(vec![Field::ULong(6), Field::ULong(7)]);
- assert_eq!(7, list.get_ulong(1).unwrap());
-
- let list = make_list(vec![
- Field::Float(8.1),
- Field::Float(9.2),
- Field::Float(10.3),
- ]);
- assert!(10.3 - list.get_float(2).unwrap() < f32::EPSILON);
-
- let list = make_list(vec![Field::Double(3.1415)]);
- assert!(3.1415 - list.get_double(0).unwrap() < f64::EPSILON);
-
- let list = make_list(vec![Field::Str("abc".to_string())]);
- assert_eq!(&"abc".to_string(), list.get_string(0).unwrap());
-
- let list = make_list(vec![Field::Bytes(ByteArray::from(vec![1, 2, 3, 4, 5]))]);
- assert_eq!(&[1, 2, 3, 4, 5], list.get_bytes(0).unwrap().data());
-
- let list = make_list(vec![Field::Decimal(Decimal::from_i32(4, 5, 2))]);
- assert_eq!(&[0, 0, 0, 4], list.get_decimal(0).unwrap().data());
- }
-
- #[test]
- fn test_list_primitive_invalid_accessors() {
- // primitives
- let list = make_list(vec![Field::Bool(false)]);
- assert!(list.get_byte(0).is_err());
-
- let list = make_list(vec![Field::Byte(3), Field::Byte(4)]);
- assert!(list.get_short(1).is_err());
-
- let list = make_list(vec![Field::Short(4), Field::Short(5), Field::Short(6)]);
- assert!(list.get_int(2).is_err());
-
- let list = make_list(vec![Field::Int(5)]);
- assert!(list.get_long(0).is_err());
-
- let list = make_list(vec![Field::Long(6), Field::Long(7)]);
- assert!(list.get_float(1).is_err());
-
- let list = make_list(vec![Field::UByte(3), Field::UByte(4)]);
- assert!(list.get_short(1).is_err());
-
- let list = make_list(vec![Field::UShort(4), Field::UShort(5), Field::UShort(6)]);
- assert!(list.get_int(2).is_err());
-
- let list = make_list(vec![Field::UInt(5)]);
- assert!(list.get_long(0).is_err());
-
- let list = make_list(vec![Field::ULong(6), Field::ULong(7)]);
- assert!(list.get_float(1).is_err());
-
- let list = make_list(vec![
- Field::Float(8.1),
- Field::Float(9.2),
- Field::Float(10.3),
- ]);
- assert!(list.get_double(2).is_err());
-
- let list = make_list(vec![Field::Double(3.1415)]);
- assert!(list.get_string(0).is_err());
-
- let list = make_list(vec![Field::Str("abc".to_string())]);
- assert!(list.get_bytes(0).is_err());
-
- let list = make_list(vec![Field::Bytes(ByteArray::from(vec![1, 2, 3, 4, 5]))]);
- assert!(list.get_bool(0).is_err());
-
- let list = make_list(vec![Field::Decimal(Decimal::from_i32(4, 5, 2))]);
- assert!(list.get_bool(0).is_err());
- }
-
- #[test]
- fn test_list_complex_accessors() {
- let list = make_list(vec![Field::Group(make_row(vec![
- ("x".to_string(), Field::Null),
- ("Y".to_string(), Field::Int(2)),
- ]))]);
- assert_eq!(2, list.get_group(0).unwrap().len());
-
- let list = make_list(vec![Field::ListInternal(make_list(vec![
- Field::Int(2),
- Field::Int(1),
- Field::Null,
- Field::Int(12),
- ]))]);
- assert_eq!(4, list.get_list(0).unwrap().len());
-
- let list = make_list(vec![Field::MapInternal(make_map(vec![
- (Field::Int(1), Field::Float(1.2)),
- (Field::Int(2), Field::Float(4.5)),
- (Field::Int(3), Field::Float(2.3)),
- ]))]);
- assert_eq!(3, list.get_map(0).unwrap().len());
- }
-
- #[test]
- fn test_list_complex_invalid_accessors() {
- let list = make_list(vec![Field::Group(make_row(vec![
- ("x".to_string(), Field::Null),
- ("Y".to_string(), Field::Int(2)),
- ]))]);
- assert_eq!(
- general_err!("Cannot access Group as Float".to_string()),
- list.get_float(0).unwrap_err()
- );
-
- let list = make_list(vec![Field::ListInternal(make_list(vec![
- Field::Int(2),
- Field::Int(1),
- Field::Null,
- Field::Int(12),
- ]))]);
- assert_eq!(
- general_err!("Cannot access ListInternal as Float".to_string()),
- list.get_float(0).unwrap_err()
- );
-
- let list = make_list(vec![Field::MapInternal(make_map(vec![
- (Field::Int(1), Field::Float(1.2)),
- (Field::Int(2), Field::Float(4.5)),
- (Field::Int(3), Field::Float(2.3)),
- ]))]);
- assert_eq!(
- general_err!("Cannot access MapInternal as Float".to_string()),
- list.get_float(0).unwrap_err()
- );
- }
-
- #[test]
- fn test_map_accessors() {
- // a map from int to string
- let map = make_map(vec![
- (Field::Int(1), Field::Str("a".to_string())),
- (Field::Int(2), Field::Str("b".to_string())),
- (Field::Int(3), Field::Str("c".to_string())),
- (Field::Int(4), Field::Str("d".to_string())),
- (Field::Int(5), Field::Str("e".to_string())),
- ]);
-
- assert_eq!(5, map.len());
- for i in 0..5 {
- assert_eq!((i + 1) as i32, map.get_keys().get_int(i).unwrap());
- assert_eq!(
- &((i as u8 + b'a') as char).to_string(),
- map.get_values().get_string(i).unwrap()
- );
- }
- }
-
- #[test]
- #[cfg(feature = "cli")]
- fn test_to_json_value() {
- assert_eq!(Field::Null.to_json_value(), Value::Null);
- assert_eq!(Field::Bool(true).to_json_value(), Value::Bool(true));
- assert_eq!(Field::Bool(false).to_json_value(), Value::Bool(false));
- assert_eq!(
- Field::Byte(1).to_json_value(),
- Value::Number(serde_json::Number::from(1))
- );
- assert_eq!(
- Field::Short(2).to_json_value(),
- Value::Number(serde_json::Number::from(2))
- );
- assert_eq!(
- Field::Int(3).to_json_value(),
- Value::Number(serde_json::Number::from(3))
- );
- assert_eq!(
- Field::Long(4).to_json_value(),
- Value::Number(serde_json::Number::from(4))
- );
- assert_eq!(
- Field::UByte(1).to_json_value(),
- Value::Number(serde_json::Number::from(1))
- );
- assert_eq!(
- Field::UShort(2).to_json_value(),
- Value::Number(serde_json::Number::from(2))
- );
- assert_eq!(
- Field::UInt(3).to_json_value(),
- Value::Number(serde_json::Number::from(3))
- );
- assert_eq!(
- Field::ULong(4).to_json_value(),
- Value::Number(serde_json::Number::from(4))
- );
- assert_eq!(
- Field::Float(5.0).to_json_value(),
- Value::Number(serde_json::Number::from_f64(f64::from(5.0 as f32)).unwrap())
- );
- assert_eq!(
- Field::Float(5.1234).to_json_value(),
- Value::Number(
- serde_json::Number::from_f64(f64::from(5.1234 as f32)).unwrap()
- )
- );
- assert_eq!(
- Field::Double(6.0).to_json_value(),
- Value::Number(serde_json::Number::from_f64(6.0 as f64).unwrap())
- );
- assert_eq!(
- Field::Double(6.1234).to_json_value(),
- Value::Number(serde_json::Number::from_f64(6.1234 as f64).unwrap())
- );
- assert_eq!(
- Field::Str("abc".to_string()).to_json_value(),
- Value::String(String::from("abc"))
- );
- assert_eq!(
- Field::Decimal(Decimal::from_i32(4, 8, 2)).to_json_value(),
- Value::String(String::from("0.04"))
- );
- assert_eq!(
- Field::Bytes(ByteArray::from(vec![1, 2, 3])).to_json_value(),
- Value::String(String::from("AQID"))
- );
- assert_eq!(
- Field::TimestampMillis(12345678).to_json_value(),
- Value::String("1970-01-01 03:25:45 +00:00".to_string())
- );
- assert_eq!(
- Field::TimestampMicros(12345678901).to_json_value(),
- Value::String(convert_timestamp_micros_to_string(12345678901))
- );
-
- let fields = vec![
- ("X".to_string(), Field::Int(1)),
- ("Y".to_string(), Field::Double(2.2)),
- ("Z".to_string(), Field::Str("abc".to_string())),
- ];
- let row = Field::Group(make_row(fields));
- assert_eq!(
- row.to_json_value(),
- serde_json::json!({"X": 1, "Y": 2.2, "Z": "abc"})
- );
-
- let row = Field::ListInternal(make_list(vec![
- Field::Int(1),
- Field::Int(12),
- Field::Null,
- ]));
- let array = vec![
- Value::Number(serde_json::Number::from(1)),
- Value::Number(serde_json::Number::from(12)),
- Value::Null,
- ];
- assert_eq!(row.to_json_value(), Value::Array(array));
-
- let row = Field::MapInternal(make_map(vec![
- (Field::Str("k1".to_string()), Field::Double(1.2)),
- (Field::Str("k2".to_string()), Field::Double(3.4)),
- (Field::Str("k3".to_string()), Field::Double(4.5)),
- ]));
- assert_eq!(
- row.to_json_value(),
- serde_json::json!({"k1": 1.2, "k2": 3.4, "k3": 4.5})
- );
- }
-}
-
-#[cfg(test)]
-#[allow(clippy::approx_constant, clippy::many_single_char_names)]
-mod api_tests {
- use super::{make_list, make_map, make_row};
- use crate::record::Field;
-
- #[test]
- fn test_field_visibility() {
- let row = make_row(vec![(
- "a".to_string(),
- Field::Group(make_row(vec![
- ("x".to_string(), Field::Null),
- ("Y".to_string(), Field::Int(2)),
- ])),
- )]);
-
- match row.get_column_iter().next() {
- Some(column) => {
- assert_eq!("a", column.0);
- match column.1 {
- Field::Group(r) => {
- assert_eq!(
- &make_row(vec![
- ("x".to_string(), Field::Null),
- ("Y".to_string(), Field::Int(2)),
- ]),
- r
- );
- }
- _ => panic!("Expected the first column to be Field::Group"),
- }
- }
- None => panic!("Expected at least one column"),
- }
- }
-
- #[test]
- fn test_list_element_access() {
- let expected = vec![
- Field::Int(1),
- Field::Group(make_row(vec![
- ("x".to_string(), Field::Null),
- ("Y".to_string(), Field::Int(2)),
- ])),
- ];
-
- let list = make_list(expected.clone());
- assert_eq!(expected.as_slice(), list.elements());
- }
-
- #[test]
- fn test_map_entry_access() {
- let expected = vec![
- (Field::Str("one".to_owned()), Field::Int(1)),
- (Field::Str("two".to_owned()), Field::Int(2)),
- ];
-
- let map = make_map(expected.clone());
- assert_eq!(expected.as_slice(), map.entries());
- }
-}
diff --git a/parquet/src/record/mod.rs b/parquet/src/record/mod.rs
deleted file mode 100644
index fb4abb5..0000000
--- a/parquet/src/record/mod.rs
+++ /dev/null
@@ -1,28 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-//! Contains record-based API for reading Parquet files.
-
-mod api;
-pub mod reader;
-mod record_writer;
-mod triplet;
-
-pub use self::{
- api::{Field, List, ListAccessor, Map, MapAccessor, Row, RowAccessor},
- record_writer::RecordWriter,
-};
diff --git a/parquet/src/record/reader.rs b/parquet/src/record/reader.rs
deleted file mode 100644
index 691afe8..0000000
--- a/parquet/src/record/reader.rs
+++ /dev/null
@@ -1,1667 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-//! Contains implementation of record assembly and converting Parquet types into
-//! [`Row`](crate::record::Row)s.
-
-use std::{collections::HashMap, fmt, sync::Arc};
-
-use crate::basic::{ConvertedType, Repetition};
-use crate::errors::{ParquetError, Result};
-use crate::file::reader::{FileReader, RowGroupReader};
-use crate::record::{
- api::{make_list, make_map, make_row, Field, Row},
- triplet::TripletIter,
-};
-use crate::schema::types::{ColumnPath, SchemaDescPtr, SchemaDescriptor, Type, TypePtr};
-
-/// Default batch size for a reader
-const DEFAULT_BATCH_SIZE: usize = 1024;
-
-/// Tree builder for `Reader` enum.
-/// Serves as a container of options for building a reader tree and a builder, and
-/// accessing a records iterator [`RowIter`].
-pub struct TreeBuilder {
- // Batch size (>= 1) for triplet iterators
- batch_size: usize,
-}
-
-impl TreeBuilder {
- /// Creates new tree builder with default parameters.
- pub fn new() -> Self {
- Self {
- batch_size: DEFAULT_BATCH_SIZE,
- }
- }
-
- /// Sets batch size for this tree builder.
- pub fn with_batch_size(mut self, batch_size: usize) -> Self {
- self.batch_size = batch_size;
- self
- }
-
- /// Creates new root reader for provided schema and row group.
- pub fn build(
- &self,
- descr: SchemaDescPtr,
- row_group_reader: &dyn RowGroupReader,
- ) -> Reader {
- // Prepare lookup table of column path -> original column index
- // This allows to prune columns and map schema leaf nodes to the column readers
- let mut paths: HashMap<ColumnPath, usize> = HashMap::new();
- let row_group_metadata = row_group_reader.metadata();
-
- for col_index in 0..row_group_reader.num_columns() {
- let col_meta = row_group_metadata.column(col_index);
- let col_path = col_meta.column_path().clone();
- paths.insert(col_path, col_index);
- }
-
- // Build child readers for the message type
- let mut readers = Vec::new();
- let mut path = Vec::new();
-
- for field in descr.root_schema().get_fields() {
- let reader = self.reader_tree(
- field.clone(),
- &mut path,
- 0,
- 0,
- &paths,
- row_group_reader,
- );
- readers.push(reader);
- }
-
- // Return group reader for message type,
- // it is always required with definition level 0
- Reader::GroupReader(None, 0, readers)
- }
-
- /// Creates iterator of `Row`s directly from schema descriptor and row group.
- pub fn as_iter(
- &self,
- descr: SchemaDescPtr,
- row_group_reader: &dyn RowGroupReader,
- ) -> ReaderIter {
- let num_records = row_group_reader.metadata().num_rows() as usize;
- ReaderIter::new(self.build(descr, row_group_reader), num_records)
- }
-
- /// Builds tree of readers for the current schema recursively.
- fn reader_tree(
- &self,
- field: TypePtr,
- mut path: &mut Vec<String>,
- mut curr_def_level: i16,
- mut curr_rep_level: i16,
- paths: &HashMap<ColumnPath, usize>,
- row_group_reader: &dyn RowGroupReader,
- ) -> Reader {
- assert!(field.get_basic_info().has_repetition());
- // Update current definition and repetition levels for this type
- let repetition = field.get_basic_info().repetition();
- match repetition {
- Repetition::OPTIONAL => {
- curr_def_level += 1;
- }
- Repetition::REPEATED => {
- curr_def_level += 1;
- curr_rep_level += 1;
- }
- _ => {}
- }
-
- path.push(String::from(field.name()));
- let reader = if field.is_primitive() {
- let col_path = ColumnPath::new(path.to_vec());
- let orig_index = *paths.get(&col_path).unwrap();
- let col_descr = row_group_reader
- .metadata()
- .column(orig_index)
- .column_descr_ptr();
- let col_reader = row_group_reader.get_column_reader(orig_index).unwrap();
- let column = TripletIter::new(col_descr, col_reader, self.batch_size);
- Reader::PrimitiveReader(field, column)
- } else {
- match field.get_basic_info().converted_type() {
- // List types
- ConvertedType::LIST => {
- assert_eq!(
- field.get_fields().len(),
- 1,
- "Invalid list type {:?}",
- field
- );
-
- let repeated_field = field.get_fields()[0].clone();
- assert_eq!(
- repeated_field.get_basic_info().repetition(),
- Repetition::REPEATED,
- "Invalid list type {:?}",
- field
- );
-
- if Reader::is_element_type(&repeated_field) {
- // Support for backward compatible lists
- let reader = self.reader_tree(
- repeated_field,
- &mut path,
- curr_def_level,
- curr_rep_level,
- paths,
- row_group_reader,
- );
-
- Reader::RepeatedReader(
- field,
- curr_def_level,
- curr_rep_level,
- Box::new(reader),
- )
- } else {
- let child_field = repeated_field.get_fields()[0].clone();
-
- path.push(String::from(repeated_field.name()));
-
- let reader = self.reader_tree(
- child_field,
- &mut path,
- curr_def_level + 1,
- curr_rep_level + 1,
- paths,
- row_group_reader,
- );
-
- path.pop();
-
- Reader::RepeatedReader(
- field,
- curr_def_level,
- curr_rep_level,
- Box::new(reader),
- )
- }
- }
- // Map types (key-value pairs)
- ConvertedType::MAP | ConvertedType::MAP_KEY_VALUE => {
- assert_eq!(
- field.get_fields().len(),
- 1,
- "Invalid map type: {:?}",
- field
- );
- assert!(
- !field.get_fields()[0].is_primitive(),
- "Invalid map type: {:?}",
- field
- );
-
- let key_value_type = field.get_fields()[0].clone();
- assert_eq!(
- key_value_type.get_basic_info().repetition(),
- Repetition::REPEATED,
- "Invalid map type: {:?}",
- field
- );
- assert_eq!(
- key_value_type.get_fields().len(),
- 2,
- "Invalid map type: {:?}",
- field
- );
-
- path.push(String::from(key_value_type.name()));
-
- let key_type = &key_value_type.get_fields()[0];
- assert!(
- key_type.is_primitive(),
- "Map key type is expected to be a primitive type, but found {:?}",
- key_type
- );
- let key_reader = self.reader_tree(
- key_type.clone(),
- &mut path,
- curr_def_level + 1,
- curr_rep_level + 1,
- paths,
- row_group_reader,
- );
-
- let value_type = &key_value_type.get_fields()[1];
- let value_reader = self.reader_tree(
- value_type.clone(),
- &mut path,
- curr_def_level + 1,
- curr_rep_level + 1,
- paths,
- row_group_reader,
- );
-
- path.pop();
-
- Reader::KeyValueReader(
- field,
- curr_def_level,
- curr_rep_level,
- Box::new(key_reader),
- Box::new(value_reader),
- )
- }
- // A repeated field that is neither contained by a `LIST`- or
- // `MAP`-annotated group nor annotated by `LIST` or `MAP`
- // should be interpreted as a required list of required
- // elements where the element type is the type of the field.
- _ if repetition == Repetition::REPEATED => {
- let required_field = Type::group_type_builder(field.name())
- .with_repetition(Repetition::REQUIRED)
- .with_converted_type(field.get_basic_info().converted_type())
- .with_fields(&mut Vec::from(field.get_fields()))
- .build()
- .unwrap();
-
- path.pop();
-
- let reader = self.reader_tree(
- Arc::new(required_field),
- &mut path,
- curr_def_level,
- curr_rep_level,
- paths,
- row_group_reader,
- );
-
- Reader::RepeatedReader(
- field,
- curr_def_level - 1,
- curr_rep_level - 1,
- Box::new(reader),
- )
- }
- // Group types (structs)
- _ => {
- let mut readers = Vec::new();
- for child in field.get_fields() {
- let reader = self.reader_tree(
- child.clone(),
- &mut path,
- curr_def_level,
- curr_rep_level,
- paths,
- row_group_reader,
- );
- readers.push(reader);
- }
- Reader::GroupReader(Some(field), curr_def_level, readers)
- }
- }
- };
- path.pop();
-
- Reader::option(repetition, curr_def_level, reader)
- }
-}
-
-/// Reader tree for record assembly
-pub enum Reader {
- // Primitive reader with type information and triplet iterator
- PrimitiveReader(TypePtr, TripletIter),
- // Optional reader with definition level of a parent and a reader
- OptionReader(i16, Box<Reader>),
- // Group (struct) reader with type information, definition level and list of child
- // readers. When it represents message type, type information is None
- GroupReader(Option<TypePtr>, i16, Vec<Reader>),
- // Reader for repeated values, e.g. lists, contains type information, definition
- // level, repetition level and a child reader
- RepeatedReader(TypePtr, i16, i16, Box<Reader>),
- // Reader of key-value pairs, e.g. maps, contains type information, definition
- // level, repetition level, child reader for keys and child reader for values
- KeyValueReader(TypePtr, i16, i16, Box<Reader>, Box<Reader>),
-}
-
-impl Reader {
- /// Wraps reader in option reader based on repetition.
- fn option(repetition: Repetition, def_level: i16, reader: Reader) -> Self {
- if repetition == Repetition::OPTIONAL {
- Reader::OptionReader(def_level - 1, Box::new(reader))
- } else {
- reader
- }
- }
-
- /// Returns true if repeated type is an element type for the list.
- /// Used to determine legacy list types.
- /// This method is copied from Spark Parquet reader and is based on the reference:
- /// <https://github.com/apache/parquet-format/blob/master/LogicalTypes.md>
- /// #backward-compatibility-rules
- fn is_element_type(repeated_type: &Type) -> bool {
- // For legacy 2-level list types with primitive element type, e.g.:
- //
- // // ARRAY<INT> (nullable list, non-null elements)
- // optional group my_list (LIST) {
- // repeated int32 element;
- // }
- //
- repeated_type.is_primitive() ||
- // For legacy 2-level list types whose element type is a group type with 2 or more
- // fields, e.g.:
- //
- // // ARRAY<STRUCT<str: STRING, num: INT>> (nullable list, non-null elements)
- // optional group my_list (LIST) {
- // repeated group element {
- // required binary str (UTF8);
- // required int32 num;
- // };
- // }
- //
- repeated_type.is_group() && repeated_type.get_fields().len() > 1 ||
- // For legacy 2-level list types generated by parquet-avro (Parquet version < 1.6.0),
- // e.g.:
- //
- // // ARRAY<STRUCT<str: STRING>> (nullable list, non-null elements)
- // optional group my_list (LIST) {
- // repeated group array {
- // required binary str (UTF8);
- // };
- // }
- //
- repeated_type.name() == "array" ||
- // For Parquet data generated by parquet-thrift, e.g.:
- //
- // // ARRAY<STRUCT<str: STRING>> (nullable list, non-null elements)
- // optional group my_list (LIST) {
- // repeated group my_list_tuple {
- // required binary str (UTF8);
- // };
- // }
- //
- repeated_type.name().ends_with("_tuple")
- }
-
- /// Reads current record as `Row` from the reader tree.
- /// Automatically advances all necessary readers.
- /// This must be called on the root level reader (i.e., for Message type).
- /// Otherwise, it will panic.
- fn read(&mut self) -> Row {
- match *self {
- Reader::GroupReader(_, _, ref mut readers) => {
- let mut fields = Vec::new();
- for reader in readers {
- fields.push((String::from(reader.field_name()), reader.read_field()));
- }
- make_row(fields)
- }
- _ => panic!("Cannot call read() on {}", self),
- }
- }
-
- /// Reads current record as `Field` from the reader tree.
- /// Automatically advances all necessary readers.
- fn read_field(&mut self) -> Field {
- match *self {
- Reader::PrimitiveReader(_, ref mut column) => {
- let value = column.current_value();
- column.read_next().unwrap();
- value
- }
- Reader::OptionReader(def_level, ref mut reader) => {
- if reader.current_def_level() > def_level {
- reader.read_field()
- } else {
- reader.advance_columns();
- Field::Null
- }
- }
- Reader::GroupReader(_, def_level, ref mut readers) => {
- let mut fields = Vec::new();
- for reader in readers {
- if reader.repetition() != Repetition::OPTIONAL
- || reader.current_def_level() > def_level
- {
- fields.push((
- String::from(reader.field_name()),
- reader.read_field(),
- ));
- } else {
- reader.advance_columns();
- fields.push((String::from(reader.field_name()), Field::Null));
- }
- }
- let row = make_row(fields);
- Field::Group(row)
- }
- Reader::RepeatedReader(_, def_level, rep_level, ref mut reader) => {
- let mut elements = Vec::new();
- loop {
- if reader.current_def_level() > def_level {
- elements.push(reader.read_field());
- } else {
- reader.advance_columns();
- // If the current definition level is equal to the definition
- // level of this repeated type, then the
- // result is an empty list and the repetition level
- // will always be <= rl.
- break;
- }
-
- // This covers case when we are out of repetition levels and should
- // close the group, or there are no values left to
- // buffer.
- if !reader.has_next() || reader.current_rep_level() <= rep_level {
- break;
- }
- }
- Field::ListInternal(make_list(elements))
- }
- Reader::KeyValueReader(
- _,
- def_level,
- rep_level,
- ref mut keys,
- ref mut values,
- ) => {
- let mut pairs = Vec::new();
- loop {
- if keys.current_def_level() > def_level {
- pairs.push((keys.read_field(), values.read_field()));
- } else {
- keys.advance_columns();
- values.advance_columns();
- // If the current definition level is equal to the definition
- // level of this repeated type, then the
- // result is an empty list and the repetition level
- // will always be <= rl.
- break;
- }
-
- // This covers case when we are out of repetition levels and should
- // close the group, or there are no values left to
- // buffer.
- if !keys.has_next() || keys.current_rep_level() <= rep_level {
- break;
- }
- }
-
- Field::MapInternal(make_map(pairs))
- }
- }
- }
-
- /// Returns field name for the current reader.
- fn field_name(&self) -> &str {
- match *self {
- Reader::PrimitiveReader(ref field, _) => field.name(),
- Reader::OptionReader(_, ref reader) => reader.field_name(),
- Reader::GroupReader(ref opt, ..) => match opt {
- Some(ref field) => field.name(),
- None => panic!("Field is None for group reader"),
- },
- Reader::RepeatedReader(ref field, ..) => field.name(),
- Reader::KeyValueReader(ref field, ..) => field.name(),
- }
- }
-
- /// Returns repetition for the current reader.
- fn repetition(&self) -> Repetition {
- match *self {
- Reader::PrimitiveReader(ref field, _) => field.get_basic_info().repetition(),
- Reader::OptionReader(_, ref reader) => reader.repetition(),
- Reader::GroupReader(ref opt, ..) => match opt {
- Some(ref field) => field.get_basic_info().repetition(),
- None => panic!("Field is None for group reader"),
- },
- Reader::RepeatedReader(ref field, ..) => field.get_basic_info().repetition(),
- Reader::KeyValueReader(ref field, ..) => field.get_basic_info().repetition(),
- }
- }
-
- /// Returns true, if current reader has more values, false otherwise.
- /// Method does not advance internal iterator.
- fn has_next(&self) -> bool {
- match *self {
- Reader::PrimitiveReader(_, ref column) => column.has_next(),
- Reader::OptionReader(_, ref reader) => reader.has_next(),
- Reader::GroupReader(_, _, ref readers) => readers.first().unwrap().has_next(),
- Reader::RepeatedReader(_, _, _, ref reader) => reader.has_next(),
- Reader::KeyValueReader(_, _, _, ref keys, _) => keys.has_next(),
- }
- }
-
- /// Returns current definition level,
- /// Method does not advance internal iterator.
- fn current_def_level(&self) -> i16 {
- match *self {
- Reader::PrimitiveReader(_, ref column) => column.current_def_level(),
- Reader::OptionReader(_, ref reader) => reader.current_def_level(),
- Reader::GroupReader(_, _, ref readers) => match readers.first() {
- Some(reader) => reader.current_def_level(),
- None => panic!("Current definition level: empty group reader"),
- },
- Reader::RepeatedReader(_, _, _, ref reader) => reader.current_def_level(),
- Reader::KeyValueReader(_, _, _, ref keys, _) => keys.current_def_level(),
- }
- }
-
- /// Returns current repetition level.
- /// Method does not advance internal iterator.
- fn current_rep_level(&self) -> i16 {
- match *self {
- Reader::PrimitiveReader(_, ref column) => column.current_rep_level(),
- Reader::OptionReader(_, ref reader) => reader.current_rep_level(),
- Reader::GroupReader(_, _, ref readers) => match readers.first() {
- Some(reader) => reader.current_rep_level(),
- None => panic!("Current repetition level: empty group reader"),
- },
- Reader::RepeatedReader(_, _, _, ref reader) => reader.current_rep_level(),
- Reader::KeyValueReader(_, _, _, ref keys, _) => keys.current_rep_level(),
- }
- }
-
- /// Advances leaf columns for the current reader.
- fn advance_columns(&mut self) {
- match *self {
- Reader::PrimitiveReader(_, ref mut column) => {
- column.read_next().unwrap();
- }
- Reader::OptionReader(_, ref mut reader) => {
- reader.advance_columns();
- }
- Reader::GroupReader(_, _, ref mut readers) => {
- for reader in readers {
- reader.advance_columns();
- }
- }
- Reader::RepeatedReader(_, _, _, ref mut reader) => {
- reader.advance_columns();
- }
- Reader::KeyValueReader(_, _, _, ref mut keys, ref mut values) => {
- keys.advance_columns();
- values.advance_columns();
- }
- }
- }
-}
-
-impl fmt::Display for Reader {
- fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
- let s = match self {
- Reader::PrimitiveReader(..) => "PrimitiveReader",
- Reader::OptionReader(..) => "OptionReader",
- Reader::GroupReader(..) => "GroupReader",
- Reader::RepeatedReader(..) => "RepeatedReader",
- Reader::KeyValueReader(..) => "KeyValueReader",
- };
- write!(f, "{}", s)
- }
-}
-
-// ----------------------------------------------------------------------
-// Row iterators
-
-/// The enum Either with variants That represet a reference and a box of
-/// [`FileReader`](crate::file::reader::FileReader).
-enum Either<'a> {
- Left(&'a dyn FileReader),
- Right(Box<dyn FileReader>),
-}
-
-impl<'a> Either<'a> {
- fn reader(&self) -> &dyn FileReader {
- match *self {
- Either::Left(r) => r,
- Either::Right(ref r) => &**r,
- }
- }
-}
-
-/// Iterator of [`Row`](crate::record::Row)s.
-/// It is used either for a single row group to iterate over data in that row group, or
-/// an entire file with auto buffering of all row groups.
-pub struct RowIter<'a> {
- descr: SchemaDescPtr,
- tree_builder: TreeBuilder,
- file_reader: Option<Either<'a>>,
- current_row_group: usize,
- num_row_groups: usize,
- row_iter: Option<ReaderIter>,
-}
-
-impl<'a> RowIter<'a> {
- /// Creates a new iterator of [`Row`](crate::record::Row)s.
- fn new(
- file_reader: Option<Either<'a>>,
- row_iter: Option<ReaderIter>,
- descr: SchemaDescPtr,
- ) -> Self {
- let tree_builder = Self::tree_builder();
- let num_row_groups = match file_reader {
- Some(ref r) => r.reader().num_row_groups(),
- None => 0,
- };
-
- Self {
- descr,
- file_reader,
- tree_builder,
- num_row_groups,
- row_iter,
- current_row_group: 0,
- }
- }
-
- /// Creates iterator of [`Row`](crate::record::Row)s for all row groups in a
- /// file.
- pub fn from_file(proj: Option<Type>, reader: &'a dyn FileReader) -> Result<Self> {
- let either = Either::Left(reader);
- let descr = Self::get_proj_descr(
- proj,
- reader.metadata().file_metadata().schema_descr_ptr(),
- )?;
-
- Ok(Self::new(Some(either), None, descr))
- }
-
- /// Creates iterator of [`Row`](crate::record::Row)s for a specific row group.
- pub fn from_row_group(
- proj: Option<Type>,
- reader: &'a dyn RowGroupReader,
- ) -> Result<Self> {
- let descr = Self::get_proj_descr(proj, reader.metadata().schema_descr_ptr())?;
- let tree_builder = Self::tree_builder();
- let row_iter = tree_builder.as_iter(descr.clone(), reader);
-
- // For row group we need to set `current_row_group` >= `num_row_groups`, because
- // we only have one row group and can't buffer more.
- Ok(Self::new(None, Some(row_iter), descr))
- }
-
- /// Creates a iterator of [`Row`](crate::record::Row)s from a
- /// [`FileReader`](crate::file::reader::FileReader) using the full file schema.
- pub fn from_file_into(reader: Box<dyn FileReader>) -> Self {
- let either = Either::Right(reader);
- let descr = either
- .reader()
- .metadata()
- .file_metadata()
- .schema_descr_ptr();
-
- Self::new(Some(either), None, descr)
- }
-
- /// Tries to create a iterator of [`Row`](crate::record::Row)s using projections.
- /// Returns a error if a file reader is not the source of this iterator.
- ///
- /// The Projected schema can be a subset of or equal to the file schema,
- /// when it is None, full file schema is assumed.
- pub fn project(self, proj: Option<Type>) -> Result<Self> {
- match self.file_reader {
- Some(ref either) => {
- let schema = either
- .reader()
- .metadata()
- .file_metadata()
- .schema_descr_ptr();
- let descr = Self::get_proj_descr(proj, schema)?;
-
- Ok(Self::new(self.file_reader, None, descr))
- }
- None => Err(general_err!("File reader is required to use projections")),
- }
- }
-
- /// Helper method to get schema descriptor for projected schema.
- /// If projection is None, then full schema is returned.
- #[inline]
- fn get_proj_descr(
- proj: Option<Type>,
- root_descr: SchemaDescPtr,
- ) -> Result<SchemaDescPtr> {
- match proj {
- Some(projection) => {
- // check if projection is part of file schema
- let root_schema = root_descr.root_schema();
- if !root_schema.check_contains(&projection) {
- return Err(general_err!("Root schema does not contain projection"));
- }
- Ok(Arc::new(SchemaDescriptor::new(Arc::new(projection))))
- }
- None => Ok(root_descr),
- }
- }
-
- /// Returns common tree builder, so the same settings are applied to both iterators
- /// from file reader and row group.
- #[inline]
- fn tree_builder() -> TreeBuilder {
- TreeBuilder::new()
- }
-}
-
-impl<'a> Iterator for RowIter<'a> {
- type Item = Row;
-
- fn next(&mut self) -> Option<Row> {
- let mut row = None;
- if let Some(ref mut iter) = self.row_iter {
- row = iter.next();
- }
-
- while row.is_none() && self.current_row_group < self.num_row_groups {
- // We do not expect any failures when accessing a row group, and file reader
- // must be set for selecting next row group.
- if let Some(ref either) = self.file_reader {
- let file_reader = either.reader();
- let row_group_reader = &*file_reader
- .get_row_group(self.current_row_group)
- .expect("Row group is required to advance");
-
- let mut iter = self
- .tree_builder
- .as_iter(self.descr.clone(), row_group_reader);
-
- row = iter.next();
-
- self.current_row_group += 1;
- self.row_iter = Some(iter);
- }
- }
-
- row
- }
-}
-
-/// Internal iterator of [`Row`](crate::record::Row)s for a reader.
-pub struct ReaderIter {
- root_reader: Reader,
- records_left: usize,
-}
-
-impl ReaderIter {
- fn new(mut root_reader: Reader, num_records: usize) -> Self {
- // Prepare root reader by advancing all column vectors
- root_reader.advance_columns();
- Self {
- root_reader,
- records_left: num_records,
- }
- }
-}
-
-impl Iterator for ReaderIter {
- type Item = Row;
-
- fn next(&mut self) -> Option<Row> {
- if self.records_left > 0 {
- self.records_left -= 1;
- Some(self.root_reader.read())
- } else {
- None
- }
- }
-}
-
-#[cfg(test)]
-mod tests {
- use super::*;
-
- use crate::errors::{ParquetError, Result};
- use crate::file::reader::{FileReader, SerializedFileReader};
- use crate::record::api::{Field, Row, RowAccessor, RowFormatter};
- use crate::schema::parser::parse_message_type;
- use crate::util::test_common::{get_test_file, get_test_path};
- use std::convert::TryFrom;
-
- // Convenient macros to assemble row, list, map, and group.
-
- macro_rules! row {
- () => {
- {
- let result = Vec::new();
- make_row(result)
- }
- };
- ( $( $e:expr ), + ) => {
- {
- let mut result = Vec::new();
- $(
- result.push($e);
- )*
- make_row(result)
- }
- }
- }
-
- macro_rules! list {
- () => {
- {
- let result = Vec::new();
- Field::ListInternal(make_list(result))
- }
- };
- ( $( $e:expr ), + ) => {
- {
- let mut result = Vec::new();
- $(
- result.push($e);
- )*
- Field::ListInternal(make_list(result))
- }
- }
- }
-
- macro_rules! map {
- () => {
- {
- let result = Vec::new();
- Field::MapInternal(make_map(result))
- }
- };
- ( $( $e:expr ), + ) => {
- {
- let mut result = Vec::new();
- $(
- result.push($e);
- )*
- Field::MapInternal(make_map(result))
- }
- }
- }
-
- macro_rules! group {
- ( $( $e:expr ), * ) => {
- {
- Field::Group(row!($( $e ), *))
- }
- }
- }
-
- #[test]
- fn test_file_reader_rows_nulls() {
- let rows = test_file_reader_rows("nulls.snappy.parquet", None).unwrap();
- let expected_rows = vec![
- row![(
- "b_struct".to_string(),
- group![("b_c_int".to_string(), Field::Null)]
- )],
- row![(
- "b_struct".to_string(),
- group![("b_c_int".to_string(), Field::Null)]
- )],
- row![(
- "b_struct".to_string(),
- group![("b_c_int".to_string(), Field::Null)]
- )],
- row![(
- "b_struct".to_string(),
- group![("b_c_int".to_string(), Field::Null)]
- )],
- row![(
- "b_struct".to_string(),
- group![("b_c_int".to_string(), Field::Null)]
- )],
- row![(
- "b_struct".to_string(),
- group![("b_c_int".to_string(), Field::Null)]
- )],
- row![(
- "b_struct".to_string(),
- group![("b_c_int".to_string(), Field::Null)]
- )],
- row![(
- "b_struct".to_string(),
- group![("b_c_int".to_string(), Field::Null)]
- )],
- ];
- assert_eq!(rows, expected_rows);
- }
-
- #[test]
- fn test_file_reader_rows_nonnullable() {
- let rows = test_file_reader_rows("nonnullable.impala.parquet", None).unwrap();
- let expected_rows = vec![row![
- ("ID".to_string(), Field::Long(8)),
- ("Int_Array".to_string(), list![Field::Int(-1)]),
- (
- "int_array_array".to_string(),
- list![list![Field::Int(-1), Field::Int(-2)], list![]]
- ),
- (
- "Int_Map".to_string(),
- map![(Field::Str("k1".to_string()), Field::Int(-1))]
- ),
- (
- "int_map_array".to_string(),
- list![
- map![],
- map![(Field::Str("k1".to_string()), Field::Int(1))],
- map![],
- map![]
- ]
- ),
- (
- "nested_Struct".to_string(),
- group![
- ("a".to_string(), Field::Int(-1)),
- ("B".to_string(), list![Field::Int(-1)]),
- (
- "c".to_string(),
- group![(
- "D".to_string(),
- list![list![group![
- ("e".to_string(), Field::Int(-1)),
- ("f".to_string(), Field::Str("nonnullable".to_string()))
- ]]]
- )]
- ),
- ("G".to_string(), map![])
- ]
- )
- ]];
- assert_eq!(rows, expected_rows);
- }
-
- #[test]
- fn test_file_reader_rows_nullable() {
- let rows = test_file_reader_rows("nullable.impala.parquet", None).unwrap();
- let expected_rows = vec![
- row![
- ("id".to_string(), Field::Long(1)),
- (
- "int_array".to_string(),
- list![Field::Int(1), Field::Int(2), Field::Int(3)]
- ),
- (
- "int_array_Array".to_string(),
- list![
- list![Field::Int(1), Field::Int(2)],
- list![Field::Int(3), Field::Int(4)]
- ]
- ),
- (
- "int_map".to_string(),
- map![
- (Field::Str("k1".to_string()), Field::Int(1)),
- (Field::Str("k2".to_string()), Field::Int(100))
- ]
- ),
- (
- "int_Map_Array".to_string(),
- list![map![(Field::Str("k1".to_string()), Field::Int(1))]]
- ),
- (
- "nested_struct".to_string(),
- group![
- ("A".to_string(), Field::Int(1)),
- ("b".to_string(), list![Field::Int(1)]),
- (
- "C".to_string(),
- group![(
- "d".to_string(),
- list![
- list![
- group![
- ("E".to_string(), Field::Int(10)),
- (
- "F".to_string(),
- Field::Str("aaa".to_string())
- )
- ],
- group![
- ("E".to_string(), Field::Int(-10)),
- (
- "F".to_string(),
- Field::Str("bbb".to_string())
- )
- ]
- ],
- list![group![
- ("E".to_string(), Field::Int(11)),
- ("F".to_string(), Field::Str("c".to_string()))
- ]]
- ]
- )]
- ),
- (
- "g".to_string(),
- map![(
- Field::Str("foo".to_string()),
- group![(
- "H".to_string(),
- group![("i".to_string(), list![Field::Double(1.1)])]
- )]
- )]
- )
- ]
- )
- ],
- row![
- ("id".to_string(), Field::Long(2)),
- (
- "int_array".to_string(),
- list![
- Field::Null,
- Field::Int(1),
- Field::Int(2),
- Field::Null,
- Field::Int(3),
- Field::Null
- ]
- ),
- (
- "int_array_Array".to_string(),
- list![
- list![Field::Null, Field::Int(1), Field::Int(2), Field::Null],
- list![Field::Int(3), Field::Null, Field::Int(4)],
- list![],
- Field::Null
- ]
- ),
- (
- "int_map".to_string(),
- map![
- (Field::Str("k1".to_string()), Field::Int(2)),
- (Field::Str("k2".to_string()), Field::Null)
- ]
- ),
- (
- "int_Map_Array".to_string(),
- list![
- map![
- (Field::Str("k3".to_string()), Field::Null),
- (Field::Str("k1".to_string()), Field::Int(1))
- ],
- Field::Null,
- map![]
- ]
- ),
- (
- "nested_struct".to_string(),
- group![
- ("A".to_string(), Field::Null),
- ("b".to_string(), list![Field::Null]),
- (
- "C".to_string(),
- group![(
- "d".to_string(),
- list![
- list![
- group![
- ("E".to_string(), Field::Null),
- ("F".to_string(), Field::Null)
- ],
- group![
- ("E".to_string(), Field::Int(10)),
- (
- "F".to_string(),
- Field::Str("aaa".to_string())
- )
- ],
- group![
- ("E".to_string(), Field::Null),
- ("F".to_string(), Field::Null)
- ],
- group![
- ("E".to_string(), Field::Int(-10)),
- (
- "F".to_string(),
- Field::Str("bbb".to_string())
- )
- ],
- group![
- ("E".to_string(), Field::Null),
- ("F".to_string(), Field::Null)
- ]
- ],
- list![
- group![
- ("E".to_string(), Field::Int(11)),
- (
- "F".to_string(),
- Field::Str("c".to_string())
- )
- ],
- Field::Null
- ],
- list![],
- Field::Null
- ]
- )]
- ),
- (
- "g".to_string(),
- map![
- (
- Field::Str("g1".to_string()),
- group![(
- "H".to_string(),
- group![(
- "i".to_string(),
- list![Field::Double(2.2), Field::Null]
- )]
- )]
- ),
- (
- Field::Str("g2".to_string()),
- group![(
- "H".to_string(),
- group![("i".to_string(), list![])]
- )]
- ),
- (Field::Str("g3".to_string()), Field::Null),
- (
- Field::Str("g4".to_string()),
- group![(
- "H".to_string(),
- group![("i".to_string(), Field::Null)]
- )]
- ),
- (
- Field::Str("g5".to_string()),
- group![("H".to_string(), Field::Null)]
- )
- ]
- )
- ]
- )
- ],
- row![
- ("id".to_string(), Field::Long(3)),
- ("int_array".to_string(), list![]),
- ("int_array_Array".to_string(), list![Field::Null]),
- ("int_map".to_string(), map![]),
- ("int_Map_Array".to_string(), list![Field::Null, Field::Null]),
- (
- "nested_struct".to_string(),
- group![
- ("A".to_string(), Field::Null),
- ("b".to_string(), Field::Null),
- ("C".to_string(), group![("d".to_string(), list![])]),
- ("g".to_string(), map![])
- ]
- )
- ],
- row![
- ("id".to_string(), Field::Long(4)),
- ("int_array".to_string(), Field::Null),
- ("int_array_Array".to_string(), list![]),
- ("int_map".to_string(), map![]),
- ("int_Map_Array".to_string(), list![]),
- (
- "nested_struct".to_string(),
- group![
- ("A".to_string(), Field::Null),
- ("b".to_string(), Field::Null),
- ("C".to_string(), group![("d".to_string(), Field::Null)]),
- ("g".to_string(), Field::Null)
- ]
- )
- ],
- row![
- ("id".to_string(), Field::Long(5)),
- ("int_array".to_string(), Field::Null),
- ("int_array_Array".to_string(), Field::Null),
- ("int_map".to_string(), map![]),
- ("int_Map_Array".to_string(), Field::Null),
- (
- "nested_struct".to_string(),
- group![
- ("A".to_string(), Field::Null),
- ("b".to_string(), Field::Null),
- ("C".to_string(), Field::Null),
- (
- "g".to_string(),
- map![(
- Field::Str("foo".to_string()),
- group![(
- "H".to_string(),
- group![(
- "i".to_string(),
- list![Field::Double(2.2), Field::Double(3.3)]
- )]
- )]
- )]
- )
- ]
- )
- ],
- row![
- ("id".to_string(), Field::Long(6)),
- ("int_array".to_string(), Field::Null),
- ("int_array_Array".to_string(), Field::Null),
- ("int_map".to_string(), Field::Null),
- ("int_Map_Array".to_string(), Field::Null),
- ("nested_struct".to_string(), Field::Null)
- ],
- row![
- ("id".to_string(), Field::Long(7)),
- ("int_array".to_string(), Field::Null),
- (
- "int_array_Array".to_string(),
- list![Field::Null, list![Field::Int(5), Field::Int(6)]]
- ),
- (
- "int_map".to_string(),
- map![
- (Field::Str("k1".to_string()), Field::Null),
- (Field::Str("k3".to_string()), Field::Null)
- ]
- ),
- ("int_Map_Array".to_string(), Field::Null),
- (
- "nested_struct".to_string(),
- group![
- ("A".to_string(), Field::Int(7)),
- (
- "b".to_string(),
- list![Field::Int(2), Field::Int(3), Field::Null]
- ),
- (
- "C".to_string(),
- group![(
- "d".to_string(),
- list![list![], list![Field::Null], Field::Null]
- )]
- ),
- ("g".to_string(), Field::Null)
- ]
- )
- ],
- ];
- assert_eq!(rows, expected_rows);
- }
-
- #[test]
- fn test_file_reader_rows_projection() {
- let schema = "
- message spark_schema {
- REQUIRED DOUBLE c;
- REQUIRED INT32 b;
- }
- ";
- let schema = parse_message_type(&schema).unwrap();
- let rows =
- test_file_reader_rows("nested_maps.snappy.parquet", Some(schema)).unwrap();
- let expected_rows = vec![
- row![
- ("c".to_string(), Field::Double(1.0)),
- ("b".to_string(), Field::Int(1))
- ],
- row![
- ("c".to_string(), Field::Double(1.0)),
- ("b".to_string(), Field::Int(1))
- ],
- row![
- ("c".to_string(), Field::Double(1.0)),
- ("b".to_string(), Field::Int(1))
- ],
- row![
- ("c".to_string(), Field::Double(1.0)),
- ("b".to_string(), Field::Int(1))
- ],
- row![
- ("c".to_string(), Field::Double(1.0)),
- ("b".to_string(), Field::Int(1))
- ],
- row![
- ("c".to_string(), Field::Double(1.0)),
- ("b".to_string(), Field::Int(1))
- ],
- ];
- assert_eq!(rows, expected_rows);
- }
-
- #[test]
- fn test_iter_columns_in_row() {
- let r = row![
- ("c".to_string(), Field::Double(1.0)),
- ("b".to_string(), Field::Int(1))
- ];
- let mut result = Vec::new();
- for (name, record) in r.get_column_iter() {
- result.push((name, record));
- }
- assert_eq!(
- vec![
- (&"c".to_string(), &Field::Double(1.0)),
- (&"b".to_string(), &Field::Int(1))
- ],
- result
- );
- }
-
- #[test]
- fn test_file_reader_rows_projection_map() {
- let schema = "
- message spark_schema {
- OPTIONAL group a (MAP) {
- REPEATED group key_value {
- REQUIRED BYTE_ARRAY key (UTF8);
- OPTIONAL group value (MAP) {
- REPEATED group key_value {
- REQUIRED INT32 key;
- REQUIRED BOOLEAN value;
- }
- }
- }
- }
- }
- ";
- let schema = parse_message_type(&schema).unwrap();
- let rows =
- test_file_reader_rows("nested_maps.snappy.parquet", Some(schema)).unwrap();
- let expected_rows = vec![
- row![(
- "a".to_string(),
- map![(
- Field::Str("a".to_string()),
- map![
- (Field::Int(1), Field::Bool(true)),
- (Field::Int(2), Field::Bool(false))
- ]
- )]
- )],
- row![(
- "a".to_string(),
- map![(
- Field::Str("b".to_string()),
- map![(Field::Int(1), Field::Bool(true))]
- )]
- )],
- row![(
- "a".to_string(),
- map![(Field::Str("c".to_string()), Field::Null)]
- )],
- row![("a".to_string(), map![(Field::Str("d".to_string()), map![])])],
- row![(
- "a".to_string(),
- map![(
- Field::Str("e".to_string()),
- map![(Field::Int(1), Field::Bool(true))]
- )]
- )],
- row![(
- "a".to_string(),
- map![(
- Field::Str("f".to_string()),
- map![
- (Field::Int(3), Field::Bool(true)),
- (Field::Int(4), Field::Bool(false)),
- (Field::Int(5), Field::Bool(true))
- ]
- )]
- )],
- ];
- assert_eq!(rows, expected_rows);
- }
-
- #[test]
- fn test_file_reader_rows_projection_list() {
- let schema = "
- message spark_schema {
- OPTIONAL group a (LIST) {
- REPEATED group list {
- OPTIONAL group element (LIST) {
- REPEATED group list {
- OPTIONAL group element (LIST) {
- REPEATED group list {
- OPTIONAL BYTE_ARRAY element (UTF8);
- }
- }
- }
- }
- }
- }
- }
- ";
- let schema = parse_message_type(&schema).unwrap();
- let rows =
- test_file_reader_rows("nested_lists.snappy.parquet", Some(schema)).unwrap();
- let expected_rows = vec![
- row![(
- "a".to_string(),
- list![
- list![
- list![Field::Str("a".to_string()), Field::Str("b".to_string())],
- list![Field::Str("c".to_string())]
- ],
- list![Field::Null, list![Field::Str("d".to_string())]]
- ]
- )],
- row![(
- "a".to_string(),
- list![
- list![
- list![Field::Str("a".to_string()), Field::Str("b".to_string())],
- list![Field::Str("c".to_string()), Field::Str("d".to_string())]
- ],
- list![Field::Null, list![Field::Str("e".to_string())]]
- ]
- )],
- row![(
- "a".to_string(),
- list![
- list![
- list![Field::Str("a".to_string()), Field::Str("b".to_string())],
- list![Field::Str("c".to_string()), Field::Str("d".to_string())],
- list![Field::Str("e".to_string())]
- ],
- list![Field::Null, list![Field::Str("f".to_string())]]
- ]
- )],
- ];
- assert_eq!(rows, expected_rows);
- }
-
- #[test]
- fn test_file_reader_rows_invalid_projection() {
- let schema = "
- message spark_schema {
- REQUIRED INT32 key;
- REQUIRED BOOLEAN value;
- }
- ";
- let schema = parse_message_type(&schema).unwrap();
- let res = test_file_reader_rows("nested_maps.snappy.parquet", Some(schema));
- assert!(res.is_err());
- assert_eq!(
- res.unwrap_err(),
- general_err!("Root schema does not contain projection")
- );
- }
-
- #[test]
- fn test_row_group_rows_invalid_projection() {
- let schema = "
- message spark_schema {
- REQUIRED INT32 key;
- REQUIRED BOOLEAN value;
- }
- ";
- let schema = parse_message_type(&schema).unwrap();
- let res = test_row_group_rows("nested_maps.snappy.parquet", Some(schema));
- assert!(res.is_err());
- assert_eq!(
- res.unwrap_err(),
- general_err!("Root schema does not contain projection")
- );
- }
-
- #[test]
- #[should_panic(expected = "Invalid map type")]
- fn test_file_reader_rows_invalid_map_type() {
- let schema = "
- message spark_schema {
- OPTIONAL group a (MAP) {
- REPEATED group key_value {
- REQUIRED BYTE_ARRAY key (UTF8);
- OPTIONAL group value (MAP) {
- REPEATED group key_value {
- REQUIRED INT32 key;
- }
- }
- }
- }
- }
- ";
- let schema = parse_message_type(&schema).unwrap();
- test_file_reader_rows("nested_maps.snappy.parquet", Some(schema)).unwrap();
- }
-
- #[test]
- fn test_file_reader_iter() {
- let path = get_test_path("alltypes_plain.parquet");
- let vec = vec![path]
- .iter()
- .map(|p| SerializedFileReader::try_from(p.as_path()).unwrap())
- .flat_map(|r| RowIter::from_file_into(Box::new(r)))
- .flat_map(|r| r.get_int(0))
- .collect::<Vec<_>>();
-
- assert_eq!(vec, vec![4, 5, 6, 7, 2, 3, 0, 1]);
- }
-
- #[test]
- fn test_file_reader_iter_projection() {
- let path = get_test_path("alltypes_plain.parquet");
- let values = vec![path]
- .iter()
- .map(|p| SerializedFileReader::try_from(p.as_path()).unwrap())
- .flat_map(|r| {
- let schema = "message schema { OPTIONAL INT32 id; }";
- let proj = parse_message_type(&schema).ok();
-
- RowIter::from_file_into(Box::new(r)).project(proj).unwrap()
- })
- .map(|r| format!("id:{}", r.fmt(0)))
- .collect::<Vec<_>>()
- .join(", ");
-
- assert_eq!(values, "id:4, id:5, id:6, id:7, id:2, id:3, id:0, id:1");
- }
-
- #[test]
- fn test_file_reader_iter_projection_err() {
- let schema = "
- message spark_schema {
- REQUIRED INT32 key;
- REQUIRED BOOLEAN value;
- }
- ";
- let proj = parse_message_type(&schema).ok();
- let path = get_test_path("nested_maps.snappy.parquet");
- let reader = SerializedFileReader::try_from(path.as_path()).unwrap();
- let res = RowIter::from_file_into(Box::new(reader)).project(proj);
-
- assert!(res.is_err());
- assert_eq!(
- res.err().unwrap(),
- general_err!("Root schema does not contain projection")
- );
- }
-
- #[test]
- fn test_tree_reader_handle_repeated_fields_with_no_annotation() {
- // Array field `phoneNumbers` does not contain LIST annotation.
- // We parse it as struct with `phone` repeated field as array.
- let rows = test_file_reader_rows("repeated_no_annotation.parquet", None).unwrap();
- let expected_rows = vec![
- row![
- ("id".to_string(), Field::Int(1)),
- ("phoneNumbers".to_string(), Field::Null)
- ],
- row![
- ("id".to_string(), Field::Int(2)),
- ("phoneNumbers".to_string(), Field::Null)
- ],
- row![
- ("id".to_string(), Field::Int(3)),
- (
- "phoneNumbers".to_string(),
- group![("phone".to_string(), list![])]
- )
- ],
- row![
- ("id".to_string(), Field::Int(4)),
- (
- "phoneNumbers".to_string(),
- group![(
- "phone".to_string(),
- list![group![
- ("number".to_string(), Field::Long(5555555555)),
- ("kind".to_string(), Field::Null)
- ]]
- )]
- )
- ],
- row![
- ("id".to_string(), Field::Int(5)),
- (
- "phoneNumbers".to_string(),
- group![(
- "phone".to_string(),
- list![group![
- ("number".to_string(), Field::Long(1111111111)),
- ("kind".to_string(), Field::Str("home".to_string()))
- ]]
- )]
- )
- ],
- row![
- ("id".to_string(), Field::Int(6)),
- (
- "phoneNumbers".to_string(),
- group![(
- "phone".to_string(),
- list![
- group![
- ("number".to_string(), Field::Long(1111111111)),
- ("kind".to_string(), Field::Str("home".to_string()))
- ],
- group![
- ("number".to_string(), Field::Long(2222222222)),
- ("kind".to_string(), Field::Null)
- ],
- group![
- ("number".to_string(), Field::Long(3333333333)),
- ("kind".to_string(), Field::Str("mobile".to_string()))
- ]
- ]
- )]
- )
- ],
- ];
-
- assert_eq!(rows, expected_rows);
- }
-
- fn test_file_reader_rows(file_name: &str, schema: Option<Type>) -> Result<Vec<Row>> {
- let file = get_test_file(file_name);
- let file_reader: Box<dyn FileReader> = Box::new(SerializedFileReader::new(file)?);
- let iter = file_reader.get_row_iter(schema)?;
- Ok(iter.collect())
- }
-
- fn test_row_group_rows(file_name: &str, schema: Option<Type>) -> Result<Vec<Row>> {
- let file = get_test_file(file_name);
- let file_reader: Box<dyn FileReader> = Box::new(SerializedFileReader::new(file)?);
- // Check the first row group only, because files will contain only single row
- // group
- let row_group_reader = file_reader.get_row_group(0).unwrap();
- let iter = row_group_reader.get_row_iter(schema)?;
- Ok(iter.collect())
- }
-}
diff --git a/parquet/src/record/record_writer.rs b/parquet/src/record/record_writer.rs
deleted file mode 100644
index 56817eb..0000000
--- a/parquet/src/record/record_writer.rs
+++ /dev/null
@@ -1,26 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-use super::super::errors::ParquetError;
-use super::super::file::writer::RowGroupWriter;
-
-pub trait RecordWriter<T> {
- fn write_to_row_group(
- &self,
- row_group_writer: &mut Box<dyn RowGroupWriter>,
- ) -> Result<(), ParquetError>;
-}
diff --git a/parquet/src/record/triplet.rs b/parquet/src/record/triplet.rs
deleted file mode 100644
index bb4f942..0000000
--- a/parquet/src/record/triplet.rs
+++ /dev/null
@@ -1,561 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-use crate::basic::Type as PhysicalType;
-use crate::column::reader::{get_typed_column_reader, ColumnReader, ColumnReaderImpl};
-use crate::data_type::*;
-use crate::errors::{ParquetError, Result};
-use crate::record::api::Field;
-use crate::schema::types::ColumnDescPtr;
-
-/// Macro to generate simple functions that cover all types of triplet iterator.
-/// $func is a function of a typed triplet iterator and $token is a either {`ref`} or
-/// {`ref`, `mut`}
-macro_rules! triplet_enum_func {
- ($self:ident, $func:ident, $( $token:tt ),*) => ({
- match *$self {
- TripletIter::BoolTripletIter($($token)* typed) => typed.$func(),
- TripletIter::Int32TripletIter($($token)* typed) => typed.$func(),
- TripletIter::Int64TripletIter($($token)* typed) => typed.$func(),
- TripletIter::Int96TripletIter($($token)* typed) => typed.$func(),
- TripletIter::FloatTripletIter($($token)* typed) => typed.$func(),
- TripletIter::DoubleTripletIter($($token)* typed) => typed.$func(),
- TripletIter::ByteArrayTripletIter($($token)* typed) => typed.$func(),
- TripletIter::FixedLenByteArrayTripletIter($($token)* typed) => typed.$func()
- }
- });
-}
-
-/// High level API wrapper on column reader.
-/// Provides per-element access for each primitive column.
-pub enum TripletIter {
- BoolTripletIter(TypedTripletIter<BoolType>),
- Int32TripletIter(TypedTripletIter<Int32Type>),
- Int64TripletIter(TypedTripletIter<Int64Type>),
- Int96TripletIter(TypedTripletIter<Int96Type>),
- FloatTripletIter(TypedTripletIter<FloatType>),
- DoubleTripletIter(TypedTripletIter<DoubleType>),
- ByteArrayTripletIter(TypedTripletIter<ByteArrayType>),
- FixedLenByteArrayTripletIter(TypedTripletIter<FixedLenByteArrayType>),
-}
-
-impl TripletIter {
- /// Creates new triplet for column reader
- pub fn new(descr: ColumnDescPtr, reader: ColumnReader, batch_size: usize) -> Self {
- match descr.physical_type() {
- PhysicalType::BOOLEAN => TripletIter::BoolTripletIter(TypedTripletIter::new(
- descr, batch_size, reader,
- )),
- PhysicalType::INT32 => TripletIter::Int32TripletIter(TypedTripletIter::new(
- descr, batch_size, reader,
- )),
- PhysicalType::INT64 => TripletIter::Int64TripletIter(TypedTripletIter::new(
- descr, batch_size, reader,
- )),
- PhysicalType::INT96 => TripletIter::Int96TripletIter(TypedTripletIter::new(
- descr, batch_size, reader,
- )),
- PhysicalType::FLOAT => TripletIter::FloatTripletIter(TypedTripletIter::new(
- descr, batch_size, reader,
- )),
- PhysicalType::DOUBLE => TripletIter::DoubleTripletIter(
- TypedTripletIter::new(descr, batch_size, reader),
- ),
- PhysicalType::BYTE_ARRAY => TripletIter::ByteArrayTripletIter(
- TypedTripletIter::new(descr, batch_size, reader),
- ),
- PhysicalType::FIXED_LEN_BYTE_ARRAY => {
- TripletIter::FixedLenByteArrayTripletIter(TypedTripletIter::new(
- descr, batch_size, reader,
- ))
- }
- }
- }
-
- /// Invokes underlying typed triplet iterator to buffer current value.
- /// Should be called once - either before `is_null` or `current_value`.
- #[inline]
- pub fn read_next(&mut self) -> Result<bool> {
- triplet_enum_func!(self, read_next, ref, mut)
- }
-
- /// Provides check on values/levels left without invoking the underlying typed triplet
- /// iterator.
- /// Returns true if more values/levels exist, false otherwise.
- /// It is always in sync with `read_next` method.
- #[inline]
- pub fn has_next(&self) -> bool {
- triplet_enum_func!(self, has_next, ref)
- }
-
- /// Returns current definition level for a leaf triplet iterator
- #[inline]
- pub fn current_def_level(&self) -> i16 {
- triplet_enum_func!(self, current_def_level, ref)
- }
-
- /// Returns max definition level for a leaf triplet iterator
- #[inline]
- pub fn max_def_level(&self) -> i16 {
- triplet_enum_func!(self, max_def_level, ref)
- }
-
- /// Returns current repetition level for a leaf triplet iterator
- #[inline]
- pub fn current_rep_level(&self) -> i16 {
- triplet_enum_func!(self, current_rep_level, ref)
- }
-
- /// Returns max repetition level for a leaf triplet iterator
- #[inline]
- pub fn max_rep_level(&self) -> i16 {
- triplet_enum_func!(self, max_rep_level, ref)
- }
-
- /// Returns true, if current value is null.
- /// Based on the fact that for non-null value current definition level
- /// equals to max definition level.
- #[inline]
- pub fn is_null(&self) -> bool {
- self.current_def_level() < self.max_def_level()
- }
-
- /// Updates non-null value for current row.
- pub fn current_value(&self) -> Field {
- assert!(!self.is_null(), "Value is null");
- match *self {
- TripletIter::BoolTripletIter(ref typed) => {
- Field::convert_bool(typed.column_descr(), *typed.current_value())
- }
- TripletIter::Int32TripletIter(ref typed) => {
- Field::convert_int32(typed.column_descr(), *typed.current_value())
- }
- TripletIter::Int64TripletIter(ref typed) => {
- Field::convert_int64(typed.column_descr(), *typed.current_value())
- }
- TripletIter::Int96TripletIter(ref typed) => {
- Field::convert_int96(typed.column_descr(), typed.current_value().clone())
- }
- TripletIter::FloatTripletIter(ref typed) => {
- Field::convert_float(typed.column_descr(), *typed.current_value())
- }
- TripletIter::DoubleTripletIter(ref typed) => {
- Field::convert_double(typed.column_descr(), *typed.current_value())
- }
- TripletIter::ByteArrayTripletIter(ref typed) => Field::convert_byte_array(
- typed.column_descr(),
- typed.current_value().clone(),
- ),
- TripletIter::FixedLenByteArrayTripletIter(ref typed) => {
- Field::convert_byte_array(
- typed.column_descr(),
- typed.current_value().clone().into(),
- )
- }
- }
- }
-}
-
-/// Internal typed triplet iterator as a wrapper for column reader
-/// (primitive leaf column), provides per-element access.
-pub struct TypedTripletIter<T: DataType> {
- reader: ColumnReaderImpl<T>,
- column_descr: ColumnDescPtr,
- batch_size: usize,
- // type properties
- max_def_level: i16,
- max_rep_level: i16,
- // values and levels
- values: Vec<T::T>,
- def_levels: Option<Vec<i16>>,
- rep_levels: Option<Vec<i16>>,
- // current index for the triplet (value, def, rep)
- curr_triplet_index: usize,
- // how many triplets are left before we need to buffer
- triplets_left: usize,
- // helper flag to quickly check if we have more values/levels to read
- has_next: bool,
-}
-
-impl<T: DataType> TypedTripletIter<T> {
- /// Creates new typed triplet iterator based on provided column reader.
- /// Use batch size to specify the amount of values to buffer from column reader.
- fn new(descr: ColumnDescPtr, batch_size: usize, column_reader: ColumnReader) -> Self {
- assert!(
- batch_size > 0,
- "Expected positive batch size, found: {}",
- batch_size
- );
-
- let max_def_level = descr.max_def_level();
- let max_rep_level = descr.max_rep_level();
-
- let def_levels = if max_def_level == 0 {
- None
- } else {
- Some(vec![0; batch_size])
- };
- let rep_levels = if max_rep_level == 0 {
- None
- } else {
- Some(vec![0; batch_size])
- };
-
- Self {
- reader: get_typed_column_reader(column_reader),
- column_descr: descr,
- batch_size,
- max_def_level,
- max_rep_level,
- values: vec![T::T::default(); batch_size],
- def_levels,
- rep_levels,
- curr_triplet_index: 0,
- triplets_left: 0,
- has_next: false,
- }
- }
-
- /// Returns column descriptor reference for the current typed triplet iterator.
- #[inline]
- pub fn column_descr(&self) -> &ColumnDescPtr {
- &self.column_descr
- }
-
- /// Returns maximum definition level for the triplet iterator (leaf column).
- #[inline]
- fn max_def_level(&self) -> i16 {
- self.max_def_level
- }
-
- /// Returns maximum repetition level for the triplet iterator (leaf column).
- #[inline]
- fn max_rep_level(&self) -> i16 {
- self.max_rep_level
- }
-
- /// Returns current value.
- /// Method does not advance the iterator, therefore can be called multiple times.
- #[inline]
- fn current_value(&self) -> &T::T {
- assert!(
- self.current_def_level() == self.max_def_level(),
- "Cannot extract value, max definition level: {}, current level: {}",
- self.max_def_level(),
- self.current_def_level()
- );
- &self.values[self.curr_triplet_index]
- }
-
- /// Returns current definition level.
- /// If field is required, then maximum definition level is returned.
- #[inline]
- fn current_def_level(&self) -> i16 {
- match self.def_levels {
- Some(ref vec) => vec[self.curr_triplet_index],
- None => self.max_def_level,
- }
- }
-
- /// Returns current repetition level.
- /// If field is required, then maximum repetition level is returned.
- #[inline]
- fn current_rep_level(&self) -> i16 {
- match self.rep_levels {
- Some(ref vec) => vec[self.curr_triplet_index],
- None => self.max_rep_level,
- }
- }
-
- /// Quick check if iterator has more values/levels to read.
- /// It is updated as a result of `read_next` method, so they are synchronized.
- #[inline]
- fn has_next(&self) -> bool {
- self.has_next
- }
-
- /// Advances to the next triplet.
- /// Returns true, if there are more records to read, false there are no records left.
- fn read_next(&mut self) -> Result<bool> {
- self.curr_triplet_index += 1;
-
- if self.curr_triplet_index >= self.triplets_left {
- let (values_read, levels_read) = {
- // Get slice of definition levels, if available
- let def_levels = self.def_levels.as_mut().map(|vec| &mut vec[..]);
-
- // Get slice of repetition levels, if available
- let rep_levels = self.rep_levels.as_mut().map(|vec| &mut vec[..]);
-
- // Buffer triplets
- self.reader.read_batch(
- self.batch_size,
- def_levels,
- rep_levels,
- &mut self.values,
- )?
- };
-
- // No more values or levels to read
- if values_read == 0 && levels_read == 0 {
- self.has_next = false;
- return Ok(false);
- }
-
- // We never read values more than levels
- if levels_read == 0 || values_read == levels_read {
- // There are no definition levels to read, column is required
- // or definition levels match values, so it does not require spacing
- self.curr_triplet_index = 0;
- self.triplets_left = values_read;
- } else if values_read < levels_read {
- // Add spacing for triplets.
- // The idea is setting values for positions in def_levels when current
- // definition level equals to maximum definition level.
- // Values and levels are guaranteed to line up, because of
- // the column reader method.
-
- // Note: if values_read == 0, then spacing will not be triggered
- let mut idx = values_read;
- let def_levels = self.def_levels.as_ref().unwrap();
- for i in 0..levels_read {
- if def_levels[levels_read - i - 1] == self.max_def_level {
- idx -= 1; // This is done to avoid usize becoming a negative value
- self.values.swap(levels_read - i - 1, idx);
- }
- }
- self.curr_triplet_index = 0;
- self.triplets_left = levels_read;
- } else {
- return Err(general_err!(
- "Spacing of values/levels is wrong, values_read: {}, levels_read: {}",
- values_read,
- levels_read
- ));
- }
- }
-
- self.has_next = true;
- Ok(true)
- }
-}
-
-#[cfg(test)]
-mod tests {
- use super::*;
-
- use crate::file::reader::{FileReader, SerializedFileReader};
- use crate::schema::types::ColumnPath;
- use crate::util::test_common::get_test_file;
-
- #[test]
- #[should_panic(expected = "Expected positive batch size, found: 0")]
- fn test_triplet_zero_batch_size() {
- let column_path =
- ColumnPath::from(vec!["b_struct".to_string(), "b_c_int".to_string()]);
- test_column_in_file("nulls.snappy.parquet", 0, &column_path, &[], &[], &[]);
- }
-
- #[test]
- fn test_triplet_null_column() {
- let path = vec!["b_struct", "b_c_int"];
- let values = vec![];
- let def_levels = vec![1, 1, 1, 1, 1, 1, 1, 1];
- let rep_levels = vec![0, 0, 0, 0, 0, 0, 0, 0];
- test_triplet_iter(
- "nulls.snappy.parquet",
- path,
- &values,
- &def_levels,
- &rep_levels,
- );
- }
-
- #[test]
- fn test_triplet_required_column() {
- let path = vec!["ID"];
- let values = vec![Field::Long(8)];
- let def_levels = vec![0];
- let rep_levels = vec![0];
- test_triplet_iter(
- "nonnullable.impala.parquet",
- path,
- &values,
- &def_levels,
- &rep_levels,
- );
- }
-
- #[test]
- fn test_triplet_optional_column() {
- let path = vec!["nested_struct", "A"];
- let values = vec![Field::Int(1), Field::Int(7)];
- let def_levels = vec![2, 1, 1, 1, 1, 0, 2];
- let rep_levels = vec![0, 0, 0, 0, 0, 0, 0];
- test_triplet_iter(
- "nullable.impala.parquet",
- path,
- &values,
- &def_levels,
- &rep_levels,
- );
- }
-
- #[test]
- fn test_triplet_optional_list_column() {
- let path = vec!["a", "list", "element", "list", "element", "list", "element"];
- let values = vec![
- Field::Str("a".to_string()),
- Field::Str("b".to_string()),
- Field::Str("c".to_string()),
- Field::Str("d".to_string()),
- Field::Str("a".to_string()),
- Field::Str("b".to_string()),
- Field::Str("c".to_string()),
- Field::Str("d".to_string()),
- Field::Str("e".to_string()),
- Field::Str("a".to_string()),
- Field::Str("b".to_string()),
- Field::Str("c".to_string()),
- Field::Str("d".to_string()),
- Field::Str("e".to_string()),
- Field::Str("f".to_string()),
- ];
- let def_levels = vec![7, 7, 7, 4, 7, 7, 7, 7, 7, 4, 7, 7, 7, 7, 7, 7, 4, 7];
- let rep_levels = vec![0, 3, 2, 1, 2, 0, 3, 2, 3, 1, 2, 0, 3, 2, 3, 2, 1, 2];
- test_triplet_iter(
- "nested_lists.snappy.parquet",
- path,
- &values,
- &def_levels,
- &rep_levels,
- );
- }
-
- #[test]
- fn test_triplet_optional_map_column() {
- let path = vec!["a", "key_value", "value", "key_value", "key"];
- let values = vec![
- Field::Int(1),
- Field::Int(2),
- Field::Int(1),
- Field::Int(1),
- Field::Int(3),
- Field::Int(4),
- Field::Int(5),
- ];
- let def_levels = vec![4, 4, 4, 2, 3, 4, 4, 4, 4];
- let rep_levels = vec![0, 2, 0, 0, 0, 0, 0, 2, 2];
- test_triplet_iter(
- "nested_maps.snappy.parquet",
- path,
- &values,
- &def_levels,
- &rep_levels,
- );
- }
-
- // Check triplet iterator across different batch sizes
- fn test_triplet_iter(
- file_name: &str,
- column_path: Vec<&str>,
- expected_values: &[Field],
- expected_def_levels: &[i16],
- expected_rep_levels: &[i16],
- ) {
- // Convert path into column path
- let path: Vec<String> = column_path.iter().map(|x| x.to_string()).collect();
- let column_path = ColumnPath::from(path);
-
- let batch_sizes = vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 128, 256];
- for batch_size in batch_sizes {
- test_column_in_file(
- file_name,
- batch_size,
- &column_path,
- expected_values,
- expected_def_levels,
- expected_rep_levels,
- );
- }
- }
-
- // Check values of a selectd column in a file
- fn test_column_in_file(
- file_name: &str,
- batch_size: usize,
- column_path: &ColumnPath,
- expected_values: &[Field],
- expected_def_levels: &[i16],
- expected_rep_levels: &[i16],
- ) {
- let file = get_test_file(file_name);
- let file_reader = SerializedFileReader::new(file).unwrap();
- let metadata = file_reader.metadata();
- // Get schema descriptor
- let file_metadata = metadata.file_metadata();
- let schema = file_metadata.schema_descr();
- // Get first row group
- let row_group_reader = file_reader.get_row_group(0).unwrap();
-
- for i in 0..schema.num_columns() {
- let descr = schema.column(i);
- if descr.path() == column_path {
- let reader = row_group_reader.get_column_reader(i).unwrap();
- test_triplet_column(
- descr,
- reader,
- batch_size,
- expected_values,
- expected_def_levels,
- expected_rep_levels,
- );
- }
- }
- }
-
- // Check values for individual triplet iterator
- fn test_triplet_column(
- descr: ColumnDescPtr,
- reader: ColumnReader,
- batch_size: usize,
- expected_values: &[Field],
- expected_def_levels: &[i16],
- expected_rep_levels: &[i16],
- ) {
- let mut iter = TripletIter::new(descr.clone(), reader, batch_size);
- let mut values: Vec<Field> = Vec::new();
- let mut def_levels: Vec<i16> = Vec::new();
- let mut rep_levels: Vec<i16> = Vec::new();
-
- assert_eq!(iter.max_def_level(), descr.max_def_level());
- assert_eq!(iter.max_rep_level(), descr.max_rep_level());
-
- while let Ok(true) = iter.read_next() {
- assert!(iter.has_next());
- if !iter.is_null() {
- values.push(iter.current_value());
- }
- def_levels.push(iter.current_def_level());
- rep_levels.push(iter.current_rep_level());
- }
-
- assert_eq!(values, expected_values);
- assert_eq!(def_levels, expected_def_levels);
- assert_eq!(rep_levels, expected_rep_levels);
- }
-}
diff --git a/parquet/src/schema/mod.rs b/parquet/src/schema/mod.rs
deleted file mode 100644
index 1ebee2e..0000000
--- a/parquet/src/schema/mod.rs
+++ /dev/null
@@ -1,67 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-//! Parquet schema definitions and methods to print and parse schema.
-//!
-//! # Example
-//!
-//! ```rust
-//! use parquet::{
-//! basic::{ConvertedType, Repetition, Type as PhysicalType},
-//! schema::{parser, printer, types::Type},
-//! };
-//! use std::sync::Arc;
-//!
-//! // Create the following schema:
-//! //
-//! // message schema {
-//! // OPTIONAL BYTE_ARRAY a (UTF8);
-//! // REQUIRED INT32 b;
-//! // }
-//!
-//! let field_a = Type::primitive_type_builder("a", PhysicalType::BYTE_ARRAY)
-//! .with_converted_type(ConvertedType::UTF8)
-//! .with_repetition(Repetition::OPTIONAL)
-//! .build()
-//! .unwrap();
-//!
-//! let field_b = Type::primitive_type_builder("b", PhysicalType::INT32)
-//! .with_repetition(Repetition::REQUIRED)
-//! .build()
-//! .unwrap();
-//!
-//! let schema = Type::group_type_builder("schema")
-//! .with_fields(&mut vec![Arc::new(field_a), Arc::new(field_b)])
-//! .build()
-//! .unwrap();
-//!
-//! let mut buf = Vec::new();
-//!
-//! // Print schema into buffer
-//! printer::print_schema(&mut buf, &schema);
-//!
-//! // Parse schema from the string
-//! let string_schema = String::from_utf8(buf).unwrap();
-//! let parsed_schema = parser::parse_message_type(&string_schema).unwrap();
-//!
-//! assert_eq!(schema, parsed_schema);
-//! ```
-
-pub mod parser;
-pub mod printer;
-pub mod types;
-pub mod visitor;
diff --git a/parquet/src/schema/parser.rs b/parquet/src/schema/parser.rs
deleted file mode 100644
index 41f6290..0000000
--- a/parquet/src/schema/parser.rs
+++ /dev/null
@@ -1,1241 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-//! Parquet schema parser.
-//! Provides methods to parse and validate string message type into Parquet
-//! [`Type`](crate::schema::types::Type).
-//!
-//! # Example
-//!
-//! ```rust
-//! use parquet::schema::parser::parse_message_type;
-//!
-//! let message_type = "
-//! message spark_schema {
-//! OPTIONAL BYTE_ARRAY a (UTF8);
-//! REQUIRED INT32 b;
-//! REQUIRED DOUBLE c;
-//! REQUIRED BOOLEAN d;
-//! OPTIONAL group e (LIST) {
-//! REPEATED group list {
-//! REQUIRED INT32 element;
-//! }
-//! }
-//! }
-//! ";
-//!
-//! let schema = parse_message_type(message_type).expect("Expected valid schema");
-//! println!("{:?}", schema);
-//! ```
-
-use std::sync::Arc;
-
-use crate::basic::{
- ConvertedType, DecimalType, IntType, LogicalType, Repetition, TimeType, TimeUnit,
- TimestampType, Type as PhysicalType,
-};
-use crate::errors::{ParquetError, Result};
-use crate::schema::types::{Type, TypePtr};
-
-/// Parses message type as string into a Parquet [`Type`](crate::schema::types::Type)
-/// which, for example, could be used to extract individual columns. Returns Parquet
-/// general error when parsing or validation fails.
-pub fn parse_message_type(message_type: &str) -> Result<Type> {
- let mut parser = Parser {
- tokenizer: &mut Tokenizer::from_str(message_type),
- };
- parser.parse_message_type()
-}
-
-/// Tokenizer to split message type string into tokens that are separated using characters
-/// defined in `is_schema_delim` method. Tokenizer also preserves delimiters as tokens.
-/// Tokenizer provides Iterator interface to process tokens; it also allows to step back
-/// to reprocess previous tokens.
-struct Tokenizer<'a> {
- // List of all tokens for a string
- tokens: Vec<&'a str>,
- // Current index of vector
- index: usize,
-}
-
-impl<'a> Tokenizer<'a> {
- // Create tokenizer from message type string
- pub fn from_str(string: &'a str) -> Self {
- let vec = string
- .split_whitespace()
- .flat_map(|t| Self::split_token(t))
- .collect();
- Tokenizer {
- tokens: vec,
- index: 0,
- }
- }
-
- // List of all special characters in schema
- fn is_schema_delim(c: char) -> bool {
- c == ';' || c == '{' || c == '}' || c == '(' || c == ')' || c == '=' || c == ','
- }
-
- /// Splits string into tokens; input string can already be token or can contain
- /// delimiters, e.g. required" -> Vec("required") and
- /// "(UTF8);" -> Vec("(", "UTF8", ")", ";")
- fn split_token(string: &str) -> Vec<&str> {
- let mut buffer: Vec<&str> = Vec::new();
- let mut tail = string;
- while let Some(index) = tail.find(Self::is_schema_delim) {
- let (h, t) = tail.split_at(index);
- if !h.is_empty() {
- buffer.push(h);
- }
- buffer.push(&t[0..1]);
- tail = &t[1..];
- }
- if !tail.is_empty() {
- buffer.push(tail);
- }
- buffer
- }
-
- // Move pointer to a previous element
- fn backtrack(&mut self) {
- self.index -= 1;
- }
-}
-
-impl<'a> Iterator for Tokenizer<'a> {
- type Item = &'a str;
-
- fn next(&mut self) -> Option<&'a str> {
- if self.index < self.tokens.len() {
- self.index += 1;
- Some(self.tokens[self.index - 1])
- } else {
- None
- }
- }
-}
-
-/// Internal Schema parser.
-/// Traverses message type using tokenizer and parses each group/primitive type
-/// recursively.
-struct Parser<'a> {
- tokenizer: &'a mut Tokenizer<'a>,
-}
-
-// Utility function to assert token on validity.
-fn assert_token(token: Option<&str>, expected: &str) -> Result<()> {
- match token {
- Some(value) if value == expected => Ok(()),
- Some(other) => Err(general_err!(
- "Expected '{}', found token '{}'",
- expected,
- other
- )),
- None => Err(general_err!(
- "Expected '{}', but no token found (None)",
- expected
- )),
- }
-}
-
-// Utility function to parse i32 or return general error.
-#[inline]
-fn parse_i32(
- value: Option<&str>,
- not_found_msg: &str,
- parse_fail_msg: &str,
-) -> Result<i32> {
- value
- .ok_or_else(|| general_err!(not_found_msg))
- .and_then(|v| v.parse::<i32>().map_err(|_| general_err!(parse_fail_msg)))
-}
-
-// Utility function to parse boolean or return general error.
-#[inline]
-fn parse_bool(
- value: Option<&str>,
- not_found_msg: &str,
- parse_fail_msg: &str,
-) -> Result<bool> {
- value
- .ok_or_else(|| general_err!(not_found_msg))
- .and_then(|v| {
- v.to_lowercase()
- .parse::<bool>()
- .map_err(|_| general_err!(parse_fail_msg))
- })
-}
-
-// Utility function to parse TimeUnit or return general error.
-fn parse_timeunit(
- value: Option<&str>,
- not_found_msg: &str,
- parse_fail_msg: &str,
-) -> Result<TimeUnit> {
- value
- .ok_or_else(|| general_err!(not_found_msg))
- .and_then(|v| match v.to_uppercase().as_str() {
- "MILLIS" => Ok(TimeUnit::MILLIS(Default::default())),
- "MICROS" => Ok(TimeUnit::MICROS(Default::default())),
- "NANOS" => Ok(TimeUnit::NANOS(Default::default())),
- _ => Err(general_err!(parse_fail_msg)),
- })
-}
-
-impl<'a> Parser<'a> {
- // Entry function to parse message type, uses internal tokenizer.
- fn parse_message_type(&mut self) -> Result<Type> {
- // Check that message type starts with "message".
- match self.tokenizer.next() {
- Some("message") => {
- let name = self
- .tokenizer
- .next()
- .ok_or_else(|| general_err!("Expected name, found None"))?;
- let mut fields = self.parse_child_types()?;
- Type::group_type_builder(name)
- .with_fields(&mut fields)
- .build()
- }
- _ => Err(general_err!("Message type does not start with 'message'")),
- }
- }
-
- // Parses child types for a current group type.
- // This is only invoked on root and group types.
- fn parse_child_types(&mut self) -> Result<Vec<TypePtr>> {
- assert_token(self.tokenizer.next(), "{")?;
- let mut vec = Vec::new();
- while let Some(value) = self.tokenizer.next() {
- if value == "}" {
- break;
- } else {
- self.tokenizer.backtrack();
- vec.push(Arc::new(self.add_type()?));
- }
- }
- Ok(vec)
- }
-
- fn add_type(&mut self) -> Result<Type> {
- // Parse repetition
- let repetition = self
- .tokenizer
- .next()
- .ok_or_else(|| general_err!("Expected repetition, found None"))
- .and_then(|v| v.to_uppercase().parse::<Repetition>())?;
-
- match self.tokenizer.next() {
- Some(group) if group.to_uppercase() == "GROUP" => {
- self.add_group_type(Some(repetition))
- }
- Some(type_string) => {
- let physical_type = type_string.to_uppercase().parse::<PhysicalType>()?;
- self.add_primitive_type(repetition, physical_type)
- }
- None => Err(general_err!("Invalid type, could not extract next token")),
- }
- }
-
- fn add_group_type(&mut self, repetition: Option<Repetition>) -> Result<Type> {
- // Parse name of the group type
- let name = self
- .tokenizer
- .next()
- .ok_or_else(|| general_err!("Expected name, found None"))?;
-
- // Parse logical or converted type if exists
- let (logical_type, converted_type) = if let Some("(") = self.tokenizer.next() {
- let tpe = self
- .tokenizer
- .next()
- .ok_or_else(|| general_err!("Expected converted type, found None"))
- .and_then(|v| {
- // Try logical type first
- let upper = v.to_uppercase();
- let logical = upper.parse::<LogicalType>();
- match logical {
- Ok(logical) => Ok((
- Some(logical.clone()),
- ConvertedType::from(Some(logical)),
- )),
- Err(_) => Ok((None, upper.parse::<ConvertedType>()?)),
- }
- })?;
- assert_token(self.tokenizer.next(), ")")?;
- tpe
- } else {
- self.tokenizer.backtrack();
- (None, ConvertedType::NONE)
- };
-
- // Parse optional id
- let id = if let Some("=") = self.tokenizer.next() {
- self.tokenizer.next().and_then(|v| v.parse::<i32>().ok())
- } else {
- self.tokenizer.backtrack();
- None
- };
-
- let mut fields = self.parse_child_types()?;
- let mut builder = Type::group_type_builder(name)
- .with_logical_type(logical_type)
- .with_converted_type(converted_type)
- .with_fields(&mut fields);
- if let Some(rep) = repetition {
- builder = builder.with_repetition(rep);
- }
- if let Some(id) = id {
- builder = builder.with_id(id);
- }
- builder.build()
- }
-
- fn add_primitive_type(
- &mut self,
- repetition: Repetition,
- physical_type: PhysicalType,
- ) -> Result<Type> {
- // Read type length if the type is FIXED_LEN_BYTE_ARRAY.
- let mut length: i32 = -1;
- if physical_type == PhysicalType::FIXED_LEN_BYTE_ARRAY {
- assert_token(self.tokenizer.next(), "(")?;
- length = parse_i32(
- self.tokenizer.next(),
- "Expected length for FIXED_LEN_BYTE_ARRAY, found None",
- "Failed to parse length for FIXED_LEN_BYTE_ARRAY",
- )?;
- assert_token(self.tokenizer.next(), ")")?;
- }
-
- // Parse name of the primitive type
- let name = self
- .tokenizer
- .next()
- .ok_or_else(|| general_err!("Expected name, found None"))?;
-
- // Parse converted type
- let (logical_type, converted_type, precision, scale) = if let Some("(") =
- self.tokenizer.next()
- {
- let (mut logical, mut converted) = self
- .tokenizer
- .next()
- .ok_or_else(|| {
- general_err!("Expected logical or converted type, found None")
- })
- .and_then(|v| {
- let upper = v.to_uppercase();
- let logical = upper.parse::<LogicalType>();
- match logical {
- Ok(logical) => Ok((
- Some(logical.clone()),
- ConvertedType::from(Some(logical)),
- )),
- Err(_) => Ok((None, upper.parse::<ConvertedType>()?)),
- }
- })?;
-
- // Parse precision and scale for decimals
- let mut precision: i32 = -1;
- let mut scale: i32 = -1;
-
- // Parse the concrete logical type
- if let Some(tpe) = &logical {
- match tpe {
- LogicalType::DECIMAL(_) => {
- if let Some("(") = self.tokenizer.next() {
- precision = parse_i32(
- self.tokenizer.next(),
- "Expected precision, found None",
- "Failed to parse precision for DECIMAL type",
- )?;
- if let Some(",") = self.tokenizer.next() {
- scale = parse_i32(
- self.tokenizer.next(),
- "Expected scale, found None",
- "Failed to parse scale for DECIMAL type",
- )?;
- assert_token(self.tokenizer.next(), ")")?;
- logical = Some(LogicalType::DECIMAL(DecimalType {
- scale,
- precision,
- }));
- converted = ConvertedType::from(logical.clone());
- } else {
- scale = 0;
- logical = Some(LogicalType::DECIMAL(DecimalType {
- scale,
- precision,
- }));
- converted = ConvertedType::from(logical.clone());
- }
- }
- }
- LogicalType::TIME(_) => {
- if let Some("(") = self.tokenizer.next() {
- let unit = parse_timeunit(
- self.tokenizer.next(),
- "Invalid timeunit found",
- "Failed to parse timeunit for TIME type",
- )?;
- if let Some(",") = self.tokenizer.next() {
- let is_adjusted_to_u_t_c = parse_bool(
- self.tokenizer.next(),
- "Invalid boolean found",
- "Failed to parse timezone info for TIME type",
- )?;
- assert_token(self.tokenizer.next(), ")")?;
- logical = Some(LogicalType::TIME(TimeType {
- is_adjusted_to_u_t_c,
- unit,
- }));
- converted = ConvertedType::from(logical.clone());
- } else {
- // Invalid token for unit
- self.tokenizer.backtrack();
- }
- }
- }
- LogicalType::TIMESTAMP(_) => {
- if let Some("(") = self.tokenizer.next() {
- let unit = parse_timeunit(
- self.tokenizer.next(),
- "Invalid timeunit found",
- "Failed to parse timeunit for TIMESTAMP type",
- )?;
- if let Some(",") = self.tokenizer.next() {
- let is_adjusted_to_u_t_c = parse_bool(
- self.tokenizer.next(),
- "Invalid boolean found",
- "Failed to parse timezone info for TIMESTAMP type",
- )?;
- assert_token(self.tokenizer.next(), ")")?;
- logical = Some(LogicalType::TIMESTAMP(TimestampType {
- is_adjusted_to_u_t_c,
- unit,
- }));
- converted = ConvertedType::from(logical.clone());
- } else {
- // Invalid token for unit
- self.tokenizer.backtrack();
- }
- }
- }
- LogicalType::INTEGER(_) => {
- if let Some("(") = self.tokenizer.next() {
- let bit_width = parse_i32(
- self.tokenizer.next(),
- "Invalid bit_width found",
- "Failed to parse bit_width for INTEGER type",
- )? as i8;
- match physical_type {
- PhysicalType::INT32 => {
- match bit_width {
- 8 | 16 | 32 => {}
- _ => {
- return Err(general_err!("Incorrect bit width {} for INT32", bit_width))
- }
- }
- }
- PhysicalType::INT64 => {
- if bit_width != 64 {
- return Err(general_err!("Incorrect bit width {} for INT64", bit_width))
- }
- }
- _ => {
- return Err(general_err!("Logical type INTEGER cannot be used with physical type {}", physical_type))
- }
- }
- if let Some(",") = self.tokenizer.next() {
- let is_signed = parse_bool(
- self.tokenizer.next(),
- "Invalid boolean found",
- "Failed to parse is_signed for INTEGER type",
- )?;
- assert_token(self.tokenizer.next(), ")")?;
- logical = Some(LogicalType::INTEGER(IntType {
- bit_width,
- is_signed,
- }));
- converted = ConvertedType::from(logical.clone());
- } else {
- // Invalid token for unit
- self.tokenizer.backtrack();
- }
- }
- }
- _ => {}
- }
- } else if converted == ConvertedType::DECIMAL {
- if let Some("(") = self.tokenizer.next() {
- // Parse precision
- precision = parse_i32(
- self.tokenizer.next(),
- "Expected precision, found None",
- "Failed to parse precision for DECIMAL type",
- )?;
-
- // Parse scale
- scale = if let Some(",") = self.tokenizer.next() {
- parse_i32(
- self.tokenizer.next(),
- "Expected scale, found None",
- "Failed to parse scale for DECIMAL type",
- )?
- } else {
- // Scale is not provided, set it to 0.
- self.tokenizer.backtrack();
- 0
- };
-
- assert_token(self.tokenizer.next(), ")")?;
- } else {
- self.tokenizer.backtrack();
- }
- }
-
- assert_token(self.tokenizer.next(), ")")?;
- (logical, converted, precision, scale)
- } else {
- self.tokenizer.backtrack();
- (None, ConvertedType::NONE, -1, -1)
- };
-
- // Parse optional id
- let id = if let Some("=") = self.tokenizer.next() {
- self.tokenizer.next().and_then(|v| v.parse::<i32>().ok())
- } else {
- self.tokenizer.backtrack();
- None
- };
- assert_token(self.tokenizer.next(), ";")?;
-
- let mut builder = Type::primitive_type_builder(name, physical_type)
- .with_repetition(repetition)
- .with_logical_type(logical_type)
- .with_converted_type(converted_type)
- .with_length(length)
- .with_precision(precision)
- .with_scale(scale);
- if let Some(id) = id {
- builder = builder.with_id(id);
- }
- builder.build()
- }
-}
-
-#[cfg(test)]
-mod tests {
- use super::*;
-
- #[test]
- fn test_tokenize_empty_string() {
- assert_eq!(Tokenizer::from_str("").next(), None);
- }
-
- #[test]
- fn test_tokenize_delimiters() {
- let mut iter = Tokenizer::from_str(",;{}()=");
- assert_eq!(iter.next(), Some(","));
- assert_eq!(iter.next(), Some(";"));
- assert_eq!(iter.next(), Some("{"));
- assert_eq!(iter.next(), Some("}"));
- assert_eq!(iter.next(), Some("("));
- assert_eq!(iter.next(), Some(")"));
- assert_eq!(iter.next(), Some("="));
- assert_eq!(iter.next(), None);
- }
-
- #[test]
- fn test_tokenize_delimiters_with_whitespaces() {
- let mut iter = Tokenizer::from_str(" , ; { } ( ) = ");
- assert_eq!(iter.next(), Some(","));
- assert_eq!(iter.next(), Some(";"));
- assert_eq!(iter.next(), Some("{"));
- assert_eq!(iter.next(), Some("}"));
- assert_eq!(iter.next(), Some("("));
- assert_eq!(iter.next(), Some(")"));
- assert_eq!(iter.next(), Some("="));
- assert_eq!(iter.next(), None);
- }
-
- #[test]
- fn test_tokenize_words() {
- let mut iter = Tokenizer::from_str("abc def ghi jkl mno");
- assert_eq!(iter.next(), Some("abc"));
- assert_eq!(iter.next(), Some("def"));
- assert_eq!(iter.next(), Some("ghi"));
- assert_eq!(iter.next(), Some("jkl"));
- assert_eq!(iter.next(), Some("mno"));
- assert_eq!(iter.next(), None);
- }
-
- #[test]
- fn test_tokenize_backtrack() {
- let mut iter = Tokenizer::from_str("abc;");
- assert_eq!(iter.next(), Some("abc"));
- assert_eq!(iter.next(), Some(";"));
- iter.backtrack();
- assert_eq!(iter.next(), Some(";"));
- assert_eq!(iter.next(), None);
- }
-
- #[test]
- fn test_tokenize_message_type() {
- let schema = "
- message schema {
- required int32 a;
- optional binary c (UTF8);
- required group d {
- required int32 a;
- optional binary c (UTF8);
- }
- required group e (LIST) {
- repeated group list {
- required int32 element;
- }
- }
- }
- ";
- let iter = Tokenizer::from_str(schema);
- let mut res = Vec::new();
- for token in iter {
- res.push(token);
- }
- assert_eq!(
- res,
- vec![
- "message", "schema", "{", "required", "int32", "a", ";", "optional",
- "binary", "c", "(", "UTF8", ")", ";", "required", "group", "d", "{",
- "required", "int32", "a", ";", "optional", "binary", "c", "(", "UTF8",
- ")", ";", "}", "required", "group", "e", "(", "LIST", ")", "{",
- "repeated", "group", "list", "{", "required", "int32", "element", ";",
- "}", "}", "}"
- ]
- );
- }
-
- #[test]
- fn test_assert_token() {
- assert!(assert_token(Some("a"), "a").is_ok());
- assert!(assert_token(Some("a"), "b").is_err());
- assert!(assert_token(None, "b").is_err());
- }
-
- #[test]
- fn test_parse_message_type_invalid() {
- let mut iter = Tokenizer::from_str("test");
- let result = Parser {
- tokenizer: &mut iter,
- }
- .parse_message_type();
- assert!(result.is_err());
- assert_eq!(
- result.unwrap_err().to_string(),
- "Parquet error: Message type does not start with 'message'"
- );
- }
-
- #[test]
- fn test_parse_message_type_no_name() {
- let mut iter = Tokenizer::from_str("message");
- let result = Parser {
- tokenizer: &mut iter,
- }
- .parse_message_type();
- assert!(result.is_err());
- assert_eq!(
- result.unwrap_err().to_string(),
- "Parquet error: Expected name, found None"
- );
- }
-
- #[test]
- fn test_parse_message_type_fixed_byte_array() {
- let schema = "
- message schema {
- REQUIRED FIXED_LEN_BYTE_ARRAY col;
- }
- ";
- let mut iter = Tokenizer::from_str(schema);
- let result = Parser {
- tokenizer: &mut iter,
- }
- .parse_message_type();
- assert!(result.is_err());
-
- let schema = "
- message schema {
- REQUIRED FIXED_LEN_BYTE_ARRAY(16) col;
- }
- ";
- let mut iter = Tokenizer::from_str(schema);
- let result = Parser {
- tokenizer: &mut iter,
- }
- .parse_message_type();
- assert!(result.is_ok());
- }
-
- #[test]
- fn test_parse_message_type_integer() {
- // Invalid integer syntax
- let schema = "
- message root {
- optional int64 f1 (INTEGER());
- }
- ";
- let mut iter = Tokenizer::from_str(schema);
- let result = Parser {
- tokenizer: &mut iter,
- }
- .parse_message_type();
- assert_eq!(
- result,
- Err(general_err!("Failed to parse bit_width for INTEGER type"))
- );
-
- // Invalid integer syntax, needs both bit-width and UTC sign
- let schema = "
- message root {
- optional int64 f1 (INTEGER(32,));
- }
- ";
- let mut iter = Tokenizer::from_str(schema);
- let result = Parser {
- tokenizer: &mut iter,
- }
- .parse_message_type();
- assert_eq!(
- result,
- Err(general_err!("Incorrect bit width 32 for INT64"))
- );
-
- // Invalid integer because of non-numeric bit width
- let schema = "
- message root {
- optional int32 f1 (INTEGER(eight,true));
- }
- ";
- let mut iter = Tokenizer::from_str(schema);
- let result = Parser {
- tokenizer: &mut iter,
- }
- .parse_message_type();
- assert_eq!(
- result,
- Err(general_err!("Failed to parse bit_width for INTEGER type"))
- );
-
- // Valid types
- let schema = "
- message root {
- optional int32 f1 (INTEGER(8,false));
- optional int32 f2 (INTEGER(8,true));
- optional int32 f3 (INTEGER(16,false));
- optional int32 f4 (INTEGER(16,true));
- optional int32 f5 (INTEGER(32,false));
- optional int32 f6 (INTEGER(32,true));
- optional int64 f7 (INTEGER(64,false));
- optional int64 f7 (INTEGER(64,true));
- }
- ";
- let mut iter = Tokenizer::from_str(schema);
- let result = Parser {
- tokenizer: &mut iter,
- }
- .parse_message_type();
- assert!(result.is_ok());
- }
-
- #[test]
- fn test_parse_message_type_temporal() {
- // Invalid timestamp syntax
- let schema = "
- message root {
- optional int64 f1 (TIMESTAMP();
- }
- ";
- let mut iter = Tokenizer::from_str(schema);
- let result = Parser {
- tokenizer: &mut iter,
- }
- .parse_message_type();
- assert_eq!(
- result,
- Err(general_err!("Failed to parse timeunit for TIMESTAMP type"))
- );
-
- // Invalid timestamp syntax, needs both unit and UTC adjustment
- let schema = "
- message root {
- optional int64 f1 (TIMESTAMP(MILLIS,));
- }
- ";
- let mut iter = Tokenizer::from_str(schema);
- let result = Parser {
- tokenizer: &mut iter,
- }
- .parse_message_type();
- assert_eq!(
- result,
- Err(general_err!(
- "Failed to parse timezone info for TIMESTAMP type"
- ))
- );
-
- // Invalid timestamp because of unknown unit
- let schema = "
- message root {
- optional int64 f1 (TIMESTAMP(YOCTOS,));
- }
- ";
- let mut iter = Tokenizer::from_str(schema);
- let result = Parser {
- tokenizer: &mut iter,
- }
- .parse_message_type();
- assert_eq!(
- result,
- Err(general_err!("Failed to parse timeunit for TIMESTAMP type"))
- );
-
- // Valid types
- let schema = "
- message root {
- optional int32 f1 (DATE);
- optional int32 f2 (TIME(MILLIS,true));
- optional int64 f3 (TIME(MICROS,false));
- optional int64 f4 (TIME(NANOS,true));
- optional int64 f5 (TIMESTAMP(MILLIS,true));
- optional int64 f6 (TIMESTAMP(MICROS,true));
- optional int64 f7 (TIMESTAMP(NANOS,false));
- }
- ";
- let mut iter = Tokenizer::from_str(schema);
- let result = Parser {
- tokenizer: &mut iter,
- }
- .parse_message_type();
- assert!(result.is_ok());
- }
-
- #[test]
- fn test_parse_message_type_decimal() {
- // It is okay for decimal to omit precision and scale with right syntax.
- // Here we test wrong syntax of decimal type
-
- // Invalid decimal syntax
- let schema = "
- message root {
- optional int32 f1 (DECIMAL();
- }
- ";
- let mut iter = Tokenizer::from_str(schema);
- let result = Parser {
- tokenizer: &mut iter,
- }
- .parse_message_type();
- assert!(result.is_err());
-
- // Invalid decimal, need precision and scale
- let schema = "
- message root {
- optional int32 f1 (DECIMAL());
- }
- ";
- let mut iter = Tokenizer::from_str(schema);
- let result = Parser {
- tokenizer: &mut iter,
- }
- .parse_message_type();
- assert!(result.is_err());
-
- // Invalid decimal because of `,` - has precision, needs scale
- let schema = "
- message root {
- optional int32 f1 (DECIMAL(8,));
- }
- ";
- let mut iter = Tokenizer::from_str(schema);
- let result = Parser {
- tokenizer: &mut iter,
- }
- .parse_message_type();
- assert!(result.is_err());
-
- // Invalid decimal because, we always require either precision or scale to be
- // specified as part of converted type
- let schema = "
- message root {
- optional int32 f3 (DECIMAL);
- }
- ";
- let mut iter = Tokenizer::from_str(schema);
- let result = Parser {
- tokenizer: &mut iter,
- }
- .parse_message_type();
- assert!(result.is_err());
-
- // Valid decimal (precision, scale)
- let schema = "
- message root {
- optional int32 f1 (DECIMAL(8, 3));
- optional int32 f2 (DECIMAL(8));
- }
- ";
- let mut iter = Tokenizer::from_str(schema);
- let result = Parser {
- tokenizer: &mut iter,
- }
- .parse_message_type();
- assert!(result.is_ok());
- }
-
- #[test]
- fn test_parse_message_type_compare_1() {
- let schema = "
- message root {
- optional fixed_len_byte_array(5) f1 (DECIMAL(9, 3));
- optional fixed_len_byte_array (16) f2 (DECIMAL (38, 18));
- }
- ";
- let mut iter = Tokenizer::from_str(schema);
- let message = Parser {
- tokenizer: &mut iter,
- }
- .parse_message_type()
- .unwrap();
-
- let expected = Type::group_type_builder("root")
- .with_fields(&mut vec![
- Arc::new(
- Type::primitive_type_builder(
- "f1",
- PhysicalType::FIXED_LEN_BYTE_ARRAY,
- )
- .with_logical_type(Some(LogicalType::DECIMAL(DecimalType {
- precision: 9,
- scale: 3,
- })))
- .with_converted_type(ConvertedType::DECIMAL)
- .with_length(5)
- .with_precision(9)
- .with_scale(3)
- .build()
- .unwrap(),
- ),
- Arc::new(
- Type::primitive_type_builder(
- "f2",
- PhysicalType::FIXED_LEN_BYTE_ARRAY,
- )
- .with_logical_type(Some(LogicalType::DECIMAL(DecimalType {
- precision: 38,
- scale: 18,
- })))
- .with_converted_type(ConvertedType::DECIMAL)
- .with_length(16)
- .with_precision(38)
- .with_scale(18)
- .build()
- .unwrap(),
- ),
- ])
- .build()
- .unwrap();
-
- assert_eq!(message, expected);
- }
-
- #[test]
- fn test_parse_message_type_compare_2() {
- let schema = "
- message root {
- required group a0 {
- optional group a1 (LIST) {
- repeated binary a2 (UTF8);
- }
-
- optional group b1 (LIST) {
- repeated group b2 {
- optional int32 b3;
- optional double b4;
- }
- }
- }
- }
- ";
- let mut iter = Tokenizer::from_str(schema);
- let message = Parser {
- tokenizer: &mut iter,
- }
- .parse_message_type()
- .unwrap();
-
- let expected = Type::group_type_builder("root")
- .with_fields(&mut vec![Arc::new(
- Type::group_type_builder("a0")
- .with_repetition(Repetition::REQUIRED)
- .with_fields(&mut vec![
- Arc::new(
- Type::group_type_builder("a1")
- .with_repetition(Repetition::OPTIONAL)
- .with_logical_type(Some(LogicalType::LIST(
- Default::default(),
- )))
- .with_converted_type(ConvertedType::LIST)
- .with_fields(&mut vec![Arc::new(
- Type::primitive_type_builder(
- "a2",
- PhysicalType::BYTE_ARRAY,
- )
- .with_repetition(Repetition::REPEATED)
- .with_converted_type(ConvertedType::UTF8)
- .build()
- .unwrap(),
- )])
- .build()
- .unwrap(),
- ),
- Arc::new(
- Type::group_type_builder("b1")
- .with_repetition(Repetition::OPTIONAL)
- .with_logical_type(Some(LogicalType::LIST(
- Default::default(),
- )))
- .with_converted_type(ConvertedType::LIST)
- .with_fields(&mut vec![Arc::new(
- Type::group_type_builder("b2")
- .with_repetition(Repetition::REPEATED)
- .with_fields(&mut vec![
- Arc::new(
- Type::primitive_type_builder(
- "b3",
- PhysicalType::INT32,
- )
- .build()
- .unwrap(),
- ),
- Arc::new(
- Type::primitive_type_builder(
- "b4",
- PhysicalType::DOUBLE,
- )
- .build()
- .unwrap(),
- ),
- ])
- .build()
- .unwrap(),
- )])
- .build()
- .unwrap(),
- ),
- ])
- .build()
- .unwrap(),
- )])
- .build()
- .unwrap();
-
- assert_eq!(message, expected);
- }
-
- #[test]
- fn test_parse_message_type_compare_3() {
- let schema = "
- message root {
- required int32 _1 (INT_8);
- required int32 _2 (INT_16);
- required float _3;
- required double _4;
- optional int32 _5 (DATE);
- optional binary _6 (UTF8);
- }
- ";
- let mut iter = Tokenizer::from_str(schema);
- let message = Parser {
- tokenizer: &mut iter,
- }
- .parse_message_type()
- .unwrap();
-
- let mut fields = vec![
- Arc::new(
- Type::primitive_type_builder("_1", PhysicalType::INT32)
- .with_repetition(Repetition::REQUIRED)
- .with_converted_type(ConvertedType::INT_8)
- .build()
- .unwrap(),
- ),
- Arc::new(
- Type::primitive_type_builder("_2", PhysicalType::INT32)
- .with_repetition(Repetition::REQUIRED)
- .with_converted_type(ConvertedType::INT_16)
- .build()
- .unwrap(),
- ),
- Arc::new(
- Type::primitive_type_builder("_3", PhysicalType::FLOAT)
- .with_repetition(Repetition::REQUIRED)
- .build()
- .unwrap(),
- ),
- Arc::new(
- Type::primitive_type_builder("_4", PhysicalType::DOUBLE)
- .with_repetition(Repetition::REQUIRED)
- .build()
- .unwrap(),
- ),
- Arc::new(
- Type::primitive_type_builder("_5", PhysicalType::INT32)
- .with_logical_type(Some(LogicalType::DATE(Default::default())))
- .with_converted_type(ConvertedType::DATE)
- .build()
- .unwrap(),
- ),
- Arc::new(
- Type::primitive_type_builder("_6", PhysicalType::BYTE_ARRAY)
- .with_converted_type(ConvertedType::UTF8)
- .build()
- .unwrap(),
- ),
- ];
-
- let expected = Type::group_type_builder("root")
- .with_fields(&mut fields)
- .build()
- .unwrap();
- assert_eq!(message, expected);
- }
-
- #[test]
- fn test_parse_message_type_compare_4() {
- let schema = "
- message root {
- required int32 _1 (INTEGER(8,true));
- required int32 _2 (INTEGER(16,false));
- required float _3;
- required double _4;
- optional int32 _5 (DATE);
- optional int32 _6 (TIME(MILLIS,false));
- optional int64 _7 (TIME(MICROS,true));
- optional int64 _8 (TIMESTAMP(MILLIS,true));
- optional int64 _9 (TIMESTAMP(NANOS,false));
- optional binary _10 (STRING);
- }
- ";
- let mut iter = Tokenizer::from_str(schema);
- let message = Parser {
- tokenizer: &mut iter,
- }
- .parse_message_type()
- .unwrap();
-
- let mut fields = vec![
- Arc::new(
- Type::primitive_type_builder("_1", PhysicalType::INT32)
- .with_repetition(Repetition::REQUIRED)
- .with_logical_type(Some(LogicalType::INTEGER(IntType {
- bit_width: 8,
- is_signed: true,
- })))
- .build()
- .unwrap(),
- ),
- Arc::new(
- Type::primitive_type_builder("_2", PhysicalType::INT32)
- .with_repetition(Repetition::REQUIRED)
- .with_logical_type(Some(LogicalType::INTEGER(IntType {
- bit_width: 16,
- is_signed: false,
- })))
- .build()
- .unwrap(),
- ),
- Arc::new(
- Type::primitive_type_builder("_3", PhysicalType::FLOAT)
- .with_repetition(Repetition::REQUIRED)
- .build()
- .unwrap(),
- ),
- Arc::new(
- Type::primitive_type_builder("_4", PhysicalType::DOUBLE)
- .with_repetition(Repetition::REQUIRED)
- .build()
- .unwrap(),
- ),
- Arc::new(
- Type::primitive_type_builder("_5", PhysicalType::INT32)
- .with_logical_type(Some(LogicalType::DATE(Default::default())))
- .build()
- .unwrap(),
- ),
- Arc::new(
- Type::primitive_type_builder("_6", PhysicalType::INT32)
- .with_logical_type(Some(LogicalType::TIME(TimeType {
- unit: TimeUnit::MILLIS(Default::default()),
- is_adjusted_to_u_t_c: false,
- })))
- .build()
- .unwrap(),
- ),
- Arc::new(
- Type::primitive_type_builder("_7", PhysicalType::INT64)
- .with_logical_type(Some(LogicalType::TIME(TimeType {
- unit: TimeUnit::MICROS(Default::default()),
- is_adjusted_to_u_t_c: true,
- })))
- .build()
- .unwrap(),
- ),
- Arc::new(
- Type::primitive_type_builder("_8", PhysicalType::INT64)
- .with_logical_type(Some(LogicalType::TIMESTAMP(TimestampType {
- unit: TimeUnit::MILLIS(Default::default()),
- is_adjusted_to_u_t_c: true,
- })))
- .build()
- .unwrap(),
- ),
- Arc::new(
- Type::primitive_type_builder("_9", PhysicalType::INT64)
- .with_logical_type(Some(LogicalType::TIMESTAMP(TimestampType {
- unit: TimeUnit::NANOS(Default::default()),
- is_adjusted_to_u_t_c: false,
- })))
- .build()
- .unwrap(),
- ),
- Arc::new(
- Type::primitive_type_builder("_10", PhysicalType::BYTE_ARRAY)
- .with_logical_type(Some(LogicalType::STRING(Default::default())))
- .build()
- .unwrap(),
- ),
- ];
-
- let expected = Type::group_type_builder("root")
- .with_fields(&mut fields)
- .build()
- .unwrap();
- assert_eq!(message, expected);
- }
-}
diff --git a/parquet/src/schema/printer.rs b/parquet/src/schema/printer.rs
deleted file mode 100644
index b1e739f..0000000
--- a/parquet/src/schema/printer.rs
+++ /dev/null
@@ -1,827 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-//! Parquet schema printer.
-//! Provides methods to print Parquet file schema and list file metadata.
-//!
-//! # Example
-//!
-//! ```rust
-//! use parquet::{
-//! file::reader::{FileReader, SerializedFileReader},
-//! schema::printer::{print_file_metadata, print_parquet_metadata, print_schema},
-//! };
-//! use std::{fs::File, path::Path};
-//!
-//! // Open a file
-//! let path = Path::new("test.parquet");
-//! if let Ok(file) = File::open(&path) {
-//! let reader = SerializedFileReader::new(file).unwrap();
-//! let parquet_metadata = reader.metadata();
-//!
-//! print_parquet_metadata(&mut std::io::stdout(), &parquet_metadata);
-//! print_file_metadata(&mut std::io::stdout(), &parquet_metadata.file_metadata());
-//!
-//! print_schema(
-//! &mut std::io::stdout(),
-//! &parquet_metadata.file_metadata().schema(),
-//! );
-//! }
-//! ```
-
-use std::{fmt, io};
-
-use crate::basic::{ConvertedType, LogicalType, TimeUnit, Type as PhysicalType};
-use crate::file::metadata::{
- ColumnChunkMetaData, FileMetaData, ParquetMetaData, RowGroupMetaData,
-};
-use crate::schema::types::Type;
-
-/// Prints Parquet metadata [`ParquetMetaData`](crate::file::metadata::ParquetMetaData)
-/// information.
-#[allow(unused_must_use)]
-pub fn print_parquet_metadata(out: &mut dyn io::Write, metadata: &ParquetMetaData) {
- print_file_metadata(out, &metadata.file_metadata());
- writeln!(out);
- writeln!(out);
- writeln!(out, "num of row groups: {}", metadata.num_row_groups());
- writeln!(out, "row groups:");
- writeln!(out);
- for (i, rg) in metadata.row_groups().iter().enumerate() {
- writeln!(out, "row group {}:", i);
- print_dashes(out, 80);
- print_row_group_metadata(out, rg);
- }
-}
-
-/// Prints file metadata [`FileMetaData`](crate::file::metadata::FileMetaData)
-/// information.
-#[allow(unused_must_use)]
-pub fn print_file_metadata(out: &mut dyn io::Write, file_metadata: &FileMetaData) {
- writeln!(out, "version: {}", file_metadata.version());
- writeln!(out, "num of rows: {}", file_metadata.num_rows());
- if let Some(created_by) = file_metadata.created_by().as_ref() {
- writeln!(out, "created by: {}", created_by);
- }
- if let Some(metadata) = file_metadata.key_value_metadata() {
- writeln!(out, "metadata:");
- for kv in metadata.iter() {
- writeln!(
- out,
- " {}: {}",
- &kv.key,
- kv.value.as_ref().unwrap_or(&"".to_owned())
- );
- }
- }
- let schema = file_metadata.schema();
- print_schema(out, schema);
-}
-
-/// Prints Parquet [`Type`](crate::schema::types::Type) information.
-#[allow(unused_must_use)]
-pub fn print_schema(out: &mut dyn io::Write, tp: &Type) {
- // TODO: better if we can pass fmt::Write to Printer.
- // But how can we make it to accept both io::Write & fmt::Write?
- let mut s = String::new();
- {
- let mut printer = Printer::new(&mut s);
- printer.print(tp);
- }
- writeln!(out, "{}", s);
-}
-
-#[allow(unused_must_use)]
-fn print_row_group_metadata(out: &mut dyn io::Write, rg_metadata: &RowGroupMetaData) {
- writeln!(out, "total byte size: {}", rg_metadata.total_byte_size());
- writeln!(out, "num of rows: {}", rg_metadata.num_rows());
- writeln!(out);
- writeln!(out, "num of columns: {}", rg_metadata.num_columns());
- writeln!(out, "columns: ");
- for (i, cc) in rg_metadata.columns().iter().enumerate() {
- writeln!(out);
- writeln!(out, "column {}:", i);
- print_dashes(out, 80);
- print_column_chunk_metadata(out, cc);
- }
-}
-
-#[allow(unused_must_use)]
-fn print_column_chunk_metadata(
- out: &mut dyn io::Write,
- cc_metadata: &ColumnChunkMetaData,
-) {
- writeln!(out, "column type: {}", cc_metadata.column_type());
- writeln!(out, "column path: {}", cc_metadata.column_path());
- let encoding_strs: Vec<_> = cc_metadata
- .encodings()
- .iter()
- .map(|e| format!("{}", e))
- .collect();
- writeln!(out, "encodings: {}", encoding_strs.join(" "));
- let file_path_str = match cc_metadata.file_path() {
- None => "N/A",
- Some(ref fp) => *fp,
- };
- writeln!(out, "file path: {}", file_path_str);
- writeln!(out, "file offset: {}", cc_metadata.file_offset());
- writeln!(out, "num of values: {}", cc_metadata.num_values());
- writeln!(
- out,
- "total compressed size (in bytes): {}",
- cc_metadata.compressed_size()
- );
- writeln!(
- out,
- "total uncompressed size (in bytes): {}",
- cc_metadata.uncompressed_size()
- );
- writeln!(out, "data page offset: {}", cc_metadata.data_page_offset());
- let index_page_offset_str = match cc_metadata.index_page_offset() {
- None => "N/A".to_owned(),
- Some(ipo) => ipo.to_string(),
- };
- writeln!(out, "index page offset: {}", index_page_offset_str);
- let dict_page_offset_str = match cc_metadata.dictionary_page_offset() {
- None => "N/A".to_owned(),
- Some(dpo) => dpo.to_string(),
- };
- writeln!(out, "dictionary page offset: {}", dict_page_offset_str);
- let statistics_str = match cc_metadata.statistics() {
- None => "N/A".to_owned(),
- Some(stats) => stats.to_string(),
- };
- writeln!(out, "statistics: {}", statistics_str);
- writeln!(out);
-}
-
-#[allow(unused_must_use)]
-fn print_dashes(out: &mut dyn io::Write, num: i32) {
- for _ in 0..num {
- write!(out, "-");
- }
- writeln!(out);
-}
-
-const INDENT_WIDTH: i32 = 2;
-
-/// Struct for printing Parquet message type.
-struct Printer<'a> {
- output: &'a mut dyn fmt::Write,
- indent: i32,
-}
-
-#[allow(unused_must_use)]
-impl<'a> Printer<'a> {
- fn new(output: &'a mut dyn fmt::Write) -> Self {
- Printer { output, indent: 0 }
- }
-
- fn print_indent(&mut self) {
- for _ in 0..self.indent {
- write!(self.output, " ");
- }
- }
-}
-
-#[inline]
-fn print_timeunit(unit: &TimeUnit) -> &str {
- match unit {
- TimeUnit::MILLIS(_) => "MILLIS",
- TimeUnit::MICROS(_) => "MICROS",
- TimeUnit::NANOS(_) => "NANOS",
- }
-}
-
-#[inline]
-fn print_logical_and_converted(
- logical_type: &Option<LogicalType>,
- converted_type: ConvertedType,
- precision: i32,
- scale: i32,
-) -> String {
- match logical_type {
- Some(logical_type) => match logical_type {
- LogicalType::INTEGER(t) => {
- format!("INTEGER({},{})", t.bit_width, t.is_signed)
- }
- LogicalType::DECIMAL(t) => {
- format!("DECIMAL({},{})", t.precision, t.scale)
- }
- LogicalType::TIMESTAMP(t) => {
- format!(
- "TIMESTAMP({},{})",
- print_timeunit(&t.unit),
- t.is_adjusted_to_u_t_c
- )
- }
- LogicalType::TIME(t) => {
- format!(
- "TIME({},{})",
- print_timeunit(&t.unit),
- t.is_adjusted_to_u_t_c
- )
- }
- LogicalType::DATE(_) => "DATE".to_string(),
- LogicalType::BSON(_) => "BSON".to_string(),
- LogicalType::JSON(_) => "JSON".to_string(),
- LogicalType::STRING(_) => "STRING".to_string(),
- LogicalType::UUID(_) => "UUID".to_string(),
- LogicalType::ENUM(_) => "ENUM".to_string(),
- LogicalType::LIST(_) => "LIST".to_string(),
- LogicalType::MAP(_) => "MAP".to_string(),
- LogicalType::UNKNOWN(_) => "UNKNOWN".to_string(),
- },
- None => {
- // Also print converted type if it is available
- match converted_type {
- ConvertedType::NONE => format!(""),
- decimal @ ConvertedType::DECIMAL => {
- // For decimal type we should print precision and scale if they
- // are > 0, e.g. DECIMAL(9, 2) -
- // DECIMAL(9) - DECIMAL
- let precision_scale = match (precision, scale) {
- (p, s) if p > 0 && s > 0 => {
- format!("{}, {}", p, s)
- }
- (p, 0) if p > 0 => format!("{}", p),
- _ => format!(""),
- };
- format!("{}{}", decimal, precision_scale)
- }
- other_converted_type => {
- format!("{}", other_converted_type)
- }
- }
- }
- }
-}
-
-#[allow(unused_must_use)]
-impl<'a> Printer<'a> {
- pub fn print(&mut self, tp: &Type) {
- self.print_indent();
- match *tp {
- Type::PrimitiveType {
- ref basic_info,
- physical_type,
- type_length,
- scale,
- precision,
- } => {
- let phys_type_str = match physical_type {
- PhysicalType::FIXED_LEN_BYTE_ARRAY => {
- // We need to include length for fixed byte array
- format!("{} ({})", physical_type, type_length)
- }
- _ => format!("{}", physical_type),
- };
- // Also print logical type if it is available
- // If there is a logical type, do not print converted type
- let logical_type_str = print_logical_and_converted(
- &basic_info.logical_type(),
- basic_info.converted_type(),
- scale,
- precision,
- );
- if logical_type_str.is_empty() {
- write!(
- self.output,
- "{} {} {};",
- basic_info.repetition(),
- phys_type_str,
- basic_info.name()
- );
- } else {
- write!(
- self.output,
- "{} {} {} ({});",
- basic_info.repetition(),
- phys_type_str,
- basic_info.name(),
- logical_type_str
- );
- }
- }
- Type::GroupType {
- ref basic_info,
- ref fields,
- } => {
- if basic_info.has_repetition() {
- let r = basic_info.repetition();
- write!(self.output, "{} group {} ", r, basic_info.name());
- let logical_str = print_logical_and_converted(
- &basic_info.logical_type(),
- basic_info.converted_type(),
- 0,
- 0,
- );
- if !logical_str.is_empty() {
- write!(self.output, "({}) ", logical_str);
- }
- writeln!(self.output, "{{");
- } else {
- writeln!(self.output, "message {} {{", basic_info.name());
- }
-
- self.indent += INDENT_WIDTH;
- for c in fields {
- self.print(&c);
- writeln!(self.output);
- }
- self.indent -= INDENT_WIDTH;
- self.print_indent();
- write!(self.output, "}}");
- }
- }
- }
-}
-
-#[cfg(test)]
-mod tests {
- use super::*;
-
- use std::sync::Arc;
-
- use crate::basic::{
- DateType, DecimalType, IntType, LogicalType, Repetition, TimeType, TimestampType,
- Type as PhysicalType,
- };
- use crate::errors::Result;
- use crate::schema::{parser::parse_message_type, types::Type};
-
- fn assert_print_parse_message(message: Type) {
- let mut s = String::new();
- {
- let mut p = Printer::new(&mut s);
- p.print(&message);
- }
- println!("{}", &s);
- let parsed = parse_message_type(&s).unwrap();
- assert_eq!(message, parsed);
- }
-
- #[test]
- fn test_print_primitive_type() {
- let mut s = String::new();
- {
- let mut p = Printer::new(&mut s);
- let field = Type::primitive_type_builder("field", PhysicalType::INT32)
- .with_repetition(Repetition::REQUIRED)
- .with_converted_type(ConvertedType::INT_32)
- .build()
- .unwrap();
- p.print(&field);
- }
- assert_eq!(&mut s, "REQUIRED INT32 field (INT_32);");
- }
-
- #[inline]
- fn build_primitive_type(
- name: &str,
- physical_type: PhysicalType,
- logical_type: Option<LogicalType>,
- converted_type: ConvertedType,
- repetition: Repetition,
- ) -> Result<Type> {
- Type::primitive_type_builder(name, physical_type)
- .with_repetition(repetition)
- .with_logical_type(logical_type)
- .with_converted_type(converted_type)
- .build()
- }
-
- #[test]
- fn test_print_logical_types() {
- let types_and_strings = vec![
- (
- build_primitive_type(
- "field",
- PhysicalType::INT32,
- Some(LogicalType::INTEGER(IntType {
- bit_width: 32,
- is_signed: true,
- })),
- ConvertedType::NONE,
- Repetition::REQUIRED,
- )
- .unwrap(),
- "REQUIRED INT32 field (INTEGER(32,true));",
- ),
- (
- build_primitive_type(
- "field",
- PhysicalType::INT32,
- Some(LogicalType::INTEGER(IntType {
- bit_width: 8,
- is_signed: false,
- })),
- ConvertedType::NONE,
- Repetition::OPTIONAL,
- )
- .unwrap(),
- "OPTIONAL INT32 field (INTEGER(8,false));",
- ),
- (
- build_primitive_type(
- "field",
- PhysicalType::INT32,
- Some(LogicalType::INTEGER(IntType {
- bit_width: 16,
- is_signed: true,
- })),
- ConvertedType::INT_16,
- Repetition::REPEATED,
- )
- .unwrap(),
- "REPEATED INT32 field (INTEGER(16,true));",
- ),
- (
- build_primitive_type(
- "field",
- PhysicalType::INT64,
- None,
- ConvertedType::NONE,
- Repetition::REPEATED,
- )
- .unwrap(),
- "REPEATED INT64 field;",
- ),
- (
- build_primitive_type(
- "field",
- PhysicalType::FLOAT,
- None,
- ConvertedType::NONE,
- Repetition::REQUIRED,
- )
- .unwrap(),
- "REQUIRED FLOAT field;",
- ),
- (
- build_primitive_type(
- "booleans",
- PhysicalType::BOOLEAN,
- None,
- ConvertedType::NONE,
- Repetition::OPTIONAL,
- )
- .unwrap(),
- "OPTIONAL BOOLEAN booleans;",
- ),
- (
- build_primitive_type(
- "field",
- PhysicalType::INT64,
- Some(LogicalType::TIMESTAMP(TimestampType {
- is_adjusted_to_u_t_c: true,
- unit: TimeUnit::MILLIS(Default::default()),
- })),
- ConvertedType::NONE,
- Repetition::REQUIRED,
- )
- .unwrap(),
- "REQUIRED INT64 field (TIMESTAMP(MILLIS,true));",
- ),
- (
- build_primitive_type(
- "field",
- PhysicalType::INT32,
- Some(LogicalType::DATE(DateType {})),
- ConvertedType::NONE,
- Repetition::OPTIONAL,
- )
- .unwrap(),
- "OPTIONAL INT32 field (DATE);",
- ),
- (
- build_primitive_type(
- "field",
- PhysicalType::INT32,
- Some(LogicalType::TIME(TimeType {
- unit: TimeUnit::MILLIS(Default::default()),
- is_adjusted_to_u_t_c: false,
- })),
- ConvertedType::TIME_MILLIS,
- Repetition::REQUIRED,
- )
- .unwrap(),
- "REQUIRED INT32 field (TIME(MILLIS,false));",
- ),
- (
- build_primitive_type(
- "field",
- PhysicalType::BYTE_ARRAY,
- None,
- ConvertedType::NONE,
- Repetition::REQUIRED,
- )
- .unwrap(),
- "REQUIRED BYTE_ARRAY field;",
- ),
- (
- build_primitive_type(
- "field",
- PhysicalType::BYTE_ARRAY,
- None,
- ConvertedType::UTF8,
- Repetition::REQUIRED,
- )
- .unwrap(),
- "REQUIRED BYTE_ARRAY field (UTF8);",
- ),
- (
- build_primitive_type(
- "field",
- PhysicalType::BYTE_ARRAY,
- Some(LogicalType::JSON(Default::default())),
- ConvertedType::JSON,
- Repetition::REQUIRED,
- )
- .unwrap(),
- "REQUIRED BYTE_ARRAY field (JSON);",
- ),
- (
- build_primitive_type(
- "field",
- PhysicalType::BYTE_ARRAY,
- Some(LogicalType::BSON(Default::default())),
- ConvertedType::BSON,
- Repetition::REQUIRED,
- )
- .unwrap(),
- "REQUIRED BYTE_ARRAY field (BSON);",
- ),
- (
- build_primitive_type(
- "field",
- PhysicalType::BYTE_ARRAY,
- Some(LogicalType::STRING(Default::default())),
- ConvertedType::NONE,
- Repetition::REQUIRED,
- )
- .unwrap(),
- "REQUIRED BYTE_ARRAY field (STRING);",
- ),
- ];
-
- types_and_strings.into_iter().for_each(|(field, expected)| {
- let mut s = String::new();
- {
- let mut p = Printer::new(&mut s);
- p.print(&field);
- }
- assert_eq!(&s, expected)
- });
- }
-
- #[inline]
- fn decimal_length_from_precision(precision: usize) -> i32 {
- (10.0_f64.powi(precision as i32).log2() / 8.0).ceil() as i32
- }
-
- #[test]
- fn test_print_flba_logical_types() {
- let types_and_strings = vec![
- (
- Type::primitive_type_builder("field", PhysicalType::FIXED_LEN_BYTE_ARRAY)
- .with_logical_type(None)
- .with_converted_type(ConvertedType::INTERVAL)
- .with_length(12)
- .with_repetition(Repetition::REQUIRED)
- .build()
- .unwrap(),
- "REQUIRED FIXED_LEN_BYTE_ARRAY (12) field (INTERVAL);",
- ),
- (
- Type::primitive_type_builder("field", PhysicalType::FIXED_LEN_BYTE_ARRAY)
- .with_logical_type(Some(LogicalType::UUID(Default::default())))
- .with_length(16)
- .with_repetition(Repetition::REQUIRED)
- .build()
- .unwrap(),
- "REQUIRED FIXED_LEN_BYTE_ARRAY (16) field (UUID);",
- ),
- (
- Type::primitive_type_builder(
- "decimal",
- PhysicalType::FIXED_LEN_BYTE_ARRAY,
- )
- .with_logical_type(Some(LogicalType::DECIMAL(DecimalType {
- precision: 32,
- scale: 20,
- })))
- .with_precision(32)
- .with_scale(20)
- .with_length(decimal_length_from_precision(32))
- .with_repetition(Repetition::REPEATED)
- .build()
- .unwrap(),
- "REPEATED FIXED_LEN_BYTE_ARRAY (14) decimal (DECIMAL(32,20));",
- ),
- ];
-
- types_and_strings.into_iter().for_each(|(field, expected)| {
- let mut s = String::new();
- {
- let mut p = Printer::new(&mut s);
- p.print(&field);
- }
- assert_eq!(&s, expected)
- });
- }
-
- #[test]
- fn test_print_group_type() {
- let mut s = String::new();
- {
- let mut p = Printer::new(&mut s);
- let f1 = Type::primitive_type_builder("f1", PhysicalType::INT32)
- .with_repetition(Repetition::REQUIRED)
- .with_converted_type(ConvertedType::INT_32)
- .with_id(0)
- .build();
- let f2 = Type::primitive_type_builder("f2", PhysicalType::BYTE_ARRAY)
- .with_converted_type(ConvertedType::UTF8)
- .with_id(1)
- .build();
- let f3 = Type::primitive_type_builder("f3", PhysicalType::BYTE_ARRAY)
- .with_logical_type(Some(LogicalType::STRING(Default::default())))
- .with_id(1)
- .build();
- let f4 =
- Type::primitive_type_builder("f4", PhysicalType::FIXED_LEN_BYTE_ARRAY)
- .with_repetition(Repetition::REPEATED)
- .with_converted_type(ConvertedType::INTERVAL)
- .with_length(12)
- .with_id(2)
- .build();
- let mut struct_fields = Vec::new();
- struct_fields.push(Arc::new(f1.unwrap()));
- struct_fields.push(Arc::new(f2.unwrap()));
- struct_fields.push(Arc::new(f3.unwrap()));
- let field = Type::group_type_builder("field")
- .with_repetition(Repetition::OPTIONAL)
- .with_fields(&mut struct_fields)
- .with_id(1)
- .build()
- .unwrap();
- let mut fields = Vec::new();
- fields.push(Arc::new(field));
- fields.push(Arc::new(f4.unwrap()));
- let message = Type::group_type_builder("schema")
- .with_fields(&mut fields)
- .with_id(2)
- .build()
- .unwrap();
- p.print(&message);
- }
- let expected = "message schema {
- OPTIONAL group field {
- REQUIRED INT32 f1 (INT_32);
- OPTIONAL BYTE_ARRAY f2 (UTF8);
- OPTIONAL BYTE_ARRAY f3 (STRING);
- }
- REPEATED FIXED_LEN_BYTE_ARRAY (12) f4 (INTERVAL);
-}";
- assert_eq!(&mut s, expected);
- }
-
- #[test]
- fn test_print_and_parse_primitive() {
- let a2 = Type::primitive_type_builder("a2", PhysicalType::BYTE_ARRAY)
- .with_repetition(Repetition::REPEATED)
- .with_converted_type(ConvertedType::UTF8)
- .build()
- .unwrap();
-
- let a1 = Type::group_type_builder("a1")
- .with_repetition(Repetition::OPTIONAL)
- .with_logical_type(Some(LogicalType::LIST(Default::default())))
- .with_converted_type(ConvertedType::LIST)
- .with_fields(&mut vec![Arc::new(a2)])
- .build()
- .unwrap();
-
- let b3 = Type::primitive_type_builder("b3", PhysicalType::INT32)
- .with_repetition(Repetition::OPTIONAL)
- .build()
- .unwrap();
-
- let b4 = Type::primitive_type_builder("b4", PhysicalType::DOUBLE)
- .with_repetition(Repetition::OPTIONAL)
- .build()
- .unwrap();
-
- let b2 = Type::group_type_builder("b2")
- .with_repetition(Repetition::REPEATED)
- .with_converted_type(ConvertedType::NONE)
- .with_fields(&mut vec![Arc::new(b3), Arc::new(b4)])
- .build()
- .unwrap();
-
- let b1 = Type::group_type_builder("b1")
- .with_repetition(Repetition::OPTIONAL)
- .with_logical_type(Some(LogicalType::LIST(Default::default())))
- .with_converted_type(ConvertedType::LIST)
- .with_fields(&mut vec![Arc::new(b2)])
- .build()
- .unwrap();
-
- let a0 = Type::group_type_builder("a0")
- .with_repetition(Repetition::REQUIRED)
- .with_fields(&mut vec![Arc::new(a1), Arc::new(b1)])
- .build()
- .unwrap();
-
- let message = Type::group_type_builder("root")
- .with_fields(&mut vec![Arc::new(a0)])
- .build()
- .unwrap();
-
- assert_print_parse_message(message);
- }
-
- #[test]
- fn test_print_and_parse_nested() {
- let f1 = Type::primitive_type_builder("f1", PhysicalType::INT32)
- .with_repetition(Repetition::REQUIRED)
- .with_converted_type(ConvertedType::INT_32)
- .build()
- .unwrap();
-
- let f2 = Type::primitive_type_builder("f2", PhysicalType::BYTE_ARRAY)
- .with_repetition(Repetition::OPTIONAL)
- .with_converted_type(ConvertedType::UTF8)
- .build()
- .unwrap();
-
- let field = Type::group_type_builder("field")
- .with_repetition(Repetition::OPTIONAL)
- .with_fields(&mut vec![Arc::new(f1), Arc::new(f2)])
- .build()
- .unwrap();
-
- let f3 = Type::primitive_type_builder("f3", PhysicalType::FIXED_LEN_BYTE_ARRAY)
- .with_repetition(Repetition::REPEATED)
- .with_converted_type(ConvertedType::INTERVAL)
- .with_length(12)
- .build()
- .unwrap();
-
- let message = Type::group_type_builder("schema")
- .with_fields(&mut vec![Arc::new(field), Arc::new(f3)])
- .build()
- .unwrap();
-
- assert_print_parse_message(message);
- }
-
- #[test]
- fn test_print_and_parse_decimal() {
- let f1 = Type::primitive_type_builder("f1", PhysicalType::INT32)
- .with_repetition(Repetition::OPTIONAL)
- .with_logical_type(Some(LogicalType::DECIMAL(DecimalType {
- precision: 9,
- scale: 2,
- })))
- .with_converted_type(ConvertedType::DECIMAL)
- .with_precision(9)
- .with_scale(2)
- .build()
- .unwrap();
-
- let f2 = Type::primitive_type_builder("f2", PhysicalType::INT32)
- .with_repetition(Repetition::OPTIONAL)
- .with_logical_type(Some(LogicalType::DECIMAL(DecimalType {
- precision: 9,
- scale: 0,
- })))
- .with_converted_type(ConvertedType::DECIMAL)
- .with_precision(9)
- .with_scale(0)
- .build()
- .unwrap();
-
- let message = Type::group_type_builder("schema")
- .with_fields(&mut vec![Arc::new(f1), Arc::new(f2)])
- .build()
- .unwrap();
-
- assert_print_parse_message(message);
- }
-}
diff --git a/parquet/src/schema/types.rs b/parquet/src/schema/types.rs
deleted file mode 100644
index 1aa8c26..0000000
--- a/parquet/src/schema/types.rs
+++ /dev/null
@@ -1,2080 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-//! Contains structs and methods to build Parquet schema and schema descriptors.
-
-use std::{collections::HashMap, convert::From, fmt, sync::Arc};
-
-use parquet_format::SchemaElement;
-
-use crate::basic::{
- ConvertedType, LogicalType, Repetition, TimeType, TimeUnit, Type as PhysicalType,
-};
-use crate::errors::{ParquetError, Result};
-
-// ----------------------------------------------------------------------
-// Parquet Type definitions
-
-/// Type alias for `Arc<Type>`.
-pub type TypePtr = Arc<Type>;
-/// Type alias for `Arc<SchemaDescriptor>`.
-pub type SchemaDescPtr = Arc<SchemaDescriptor>;
-/// Type alias for `Arc<ColumnDescriptor>`.
-pub type ColumnDescPtr = Arc<ColumnDescriptor>;
-
-/// Representation of a Parquet type.
-/// Used to describe primitive leaf fields and structs, including top-level schema.
-/// Note that the top-level schema type is represented using `GroupType` whose
-/// repetition is `None`.
-#[derive(Clone, Debug, PartialEq)]
-pub enum Type {
- PrimitiveType {
- basic_info: BasicTypeInfo,
- physical_type: PhysicalType,
- type_length: i32,
- scale: i32,
- precision: i32,
- },
- GroupType {
- basic_info: BasicTypeInfo,
- fields: Vec<TypePtr>,
- },
-}
-
-impl Type {
- /// Creates primitive type builder with provided field name and physical type.
- pub fn primitive_type_builder(
- name: &str,
- physical_type: PhysicalType,
- ) -> PrimitiveTypeBuilder {
- PrimitiveTypeBuilder::new(name, physical_type)
- }
-
- /// Creates group type builder with provided column name.
- pub fn group_type_builder(name: &str) -> GroupTypeBuilder {
- GroupTypeBuilder::new(name)
- }
-
- /// Returns [`BasicTypeInfo`] information about the type.
- pub fn get_basic_info(&self) -> &BasicTypeInfo {
- match *self {
- Type::PrimitiveType { ref basic_info, .. } => &basic_info,
- Type::GroupType { ref basic_info, .. } => &basic_info,
- }
- }
-
- /// Returns this type's field name.
- pub fn name(&self) -> &str {
- self.get_basic_info().name()
- }
-
- /// Gets the fields from this group type.
- /// Note that this will panic if called on a non-group type.
- // TODO: should we return `&[&Type]` here?
- pub fn get_fields(&self) -> &[TypePtr] {
- match *self {
- Type::GroupType { ref fields, .. } => &fields[..],
- _ => panic!("Cannot call get_fields() on a non-group type"),
- }
- }
-
- /// Gets physical type of this primitive type.
- /// Note that this will panic if called on a non-primitive type.
- pub fn get_physical_type(&self) -> PhysicalType {
- match *self {
- Type::PrimitiveType {
- basic_info: _,
- physical_type,
- ..
- } => physical_type,
- _ => panic!("Cannot call get_physical_type() on a non-primitive type"),
- }
- }
-
- /// Gets precision of this primitive type.
- /// Note that this will panic if called on a non-primitive type.
- pub fn get_precision(&self) -> i32 {
- match *self {
- Type::PrimitiveType { precision, .. } => precision,
- _ => panic!("Cannot call get_precision() on non-primitive type"),
- }
- }
-
- /// Gets scale of this primitive type.
- /// Note that this will panic if called on a non-primitive type.
- pub fn get_scale(&self) -> i32 {
- match *self {
- Type::PrimitiveType { scale, .. } => scale,
- _ => panic!("Cannot call get_scale() on non-primitive type"),
- }
- }
-
- /// Checks if `sub_type` schema is part of current schema.
- /// This method can be used to check if projected columns are part of the root schema.
- pub fn check_contains(&self, sub_type: &Type) -> bool {
- // Names match, and repetitions match or not set for both
- let basic_match = self.get_basic_info().name()
- == sub_type.get_basic_info().name()
- && (self.is_schema() && sub_type.is_schema()
- || !self.is_schema()
- && !sub_type.is_schema()
- && self.get_basic_info().repetition()
- == sub_type.get_basic_info().repetition());
-
- match *self {
- Type::PrimitiveType { .. } if basic_match && sub_type.is_primitive() => {
- self.get_physical_type() == sub_type.get_physical_type()
- }
- Type::GroupType { .. } if basic_match && sub_type.is_group() => {
- // build hashmap of name -> TypePtr
- let mut field_map = HashMap::new();
- for field in self.get_fields() {
- field_map.insert(field.name(), field);
- }
-
- for field in sub_type.get_fields() {
- if !field_map
- .get(field.name())
- .map(|tpe| tpe.check_contains(field))
- .unwrap_or(false)
- {
- return false;
- }
- }
- true
- }
- _ => false,
- }
- }
-
- /// Returns `true` if this type is a primitive type, `false` otherwise.
- pub fn is_primitive(&self) -> bool {
- matches!(*self, Type::PrimitiveType { .. })
- }
-
- /// Returns `true` if this type is a group type, `false` otherwise.
- pub fn is_group(&self) -> bool {
- matches!(*self, Type::GroupType { .. })
- }
-
- /// Returns `true` if this type is the top-level schema type (message type).
- pub fn is_schema(&self) -> bool {
- match *self {
- Type::GroupType { ref basic_info, .. } => !basic_info.has_repetition(),
- _ => false,
- }
- }
-
- /// Returns `true` if this type is repeated or optional.
- /// If this type doesn't have repetition defined, we still treat it as optional.
- pub fn is_optional(&self) -> bool {
- self.get_basic_info().has_repetition()
- && self.get_basic_info().repetition() != Repetition::REQUIRED
- }
-}
-
-/// A builder for primitive types. All attributes are optional
-/// except the name and physical type.
-/// Note that if not specified explicitly, `Repetition::OPTIONAL` is used.
-pub struct PrimitiveTypeBuilder<'a> {
- name: &'a str,
- repetition: Repetition,
- physical_type: PhysicalType,
- converted_type: ConvertedType,
- logical_type: Option<LogicalType>,
- length: i32,
- precision: i32,
- scale: i32,
- id: Option<i32>,
-}
-
-impl<'a> PrimitiveTypeBuilder<'a> {
- /// Creates new primitive type builder with provided field name and physical type.
- pub fn new(name: &'a str, physical_type: PhysicalType) -> Self {
- Self {
- name,
- repetition: Repetition::OPTIONAL,
- physical_type,
- converted_type: ConvertedType::NONE,
- logical_type: None,
- length: -1,
- precision: -1,
- scale: -1,
- id: None,
- }
- }
-
- /// Sets [`Repetition`](crate::basic::Repetition) for this field and returns itself.
- pub fn with_repetition(mut self, repetition: Repetition) -> Self {
- self.repetition = repetition;
- self
- }
-
- /// Sets [`ConvertedType`](crate::basic::ConvertedType) for this field and returns itself.
- pub fn with_converted_type(mut self, converted_type: ConvertedType) -> Self {
- self.converted_type = converted_type;
- self
- }
-
- /// Sets [`LogicalType`](crate::basic::LogicalType) for this field and returns itself.
- /// If only the logical type is populated for a primitive type, the converted type
- /// will be automatically populated, and can thus be omitted.
- pub fn with_logical_type(mut self, logical_type: Option<LogicalType>) -> Self {
- self.logical_type = logical_type;
- self
- }
-
- /// Sets type length and returns itself.
- /// This is only applied to FIXED_LEN_BYTE_ARRAY and INT96 (INTERVAL) types, because
- /// they maintain fixed size underlying byte array.
- /// By default, value is `0`.
- pub fn with_length(mut self, length: i32) -> Self {
- self.length = length;
- self
- }
-
- /// Sets precision for Parquet DECIMAL physical type and returns itself.
- /// By default, it equals to `0` and used only for decimal context.
- pub fn with_precision(mut self, precision: i32) -> Self {
- self.precision = precision;
- self
- }
-
- /// Sets scale for Parquet DECIMAL physical type and returns itself.
- /// By default, it equals to `0` and used only for decimal context.
- pub fn with_scale(mut self, scale: i32) -> Self {
- self.scale = scale;
- self
- }
-
- /// Sets optional field id and returns itself.
- pub fn with_id(mut self, id: i32) -> Self {
- self.id = Some(id);
- self
- }
-
- /// Creates a new `PrimitiveType` instance from the collected attributes.
- /// Returns `Err` in case of any building conditions are not met.
- pub fn build(self) -> Result<Type> {
- let mut basic_info = BasicTypeInfo {
- name: String::from(self.name),
- repetition: Some(self.repetition),
- converted_type: self.converted_type,
- logical_type: self.logical_type.clone(),
- id: self.id,
- };
-
- // Check length before logical type, since it is used for logical type validation.
- if self.physical_type == PhysicalType::FIXED_LEN_BYTE_ARRAY && self.length < 0 {
- return Err(general_err!(
- "Invalid FIXED_LEN_BYTE_ARRAY length: {}",
- self.length
- ));
- }
-
- match &self.logical_type {
- Some(logical_type) => {
- // If a converted type is populated, check that it is consistent with
- // its logical type
- if self.converted_type != ConvertedType::NONE {
- if ConvertedType::from(self.logical_type.clone())
- != self.converted_type
- {
- return Err(general_err!(
- "Logical type {:?} is imcompatible with converted type {}",
- logical_type,
- self.converted_type
- ));
- }
- } else {
- // Populate the converted type for backwards compatibility
- basic_info.converted_type = self.logical_type.clone().into();
- }
- // Check that logical type and physical type are compatible
- match (logical_type, self.physical_type) {
- (LogicalType::MAP(_), _) | (LogicalType::LIST(_), _) => {
- return Err(general_err!(
- "{:?} cannot be applied to a primitive type",
- logical_type
- ));
- }
- (LogicalType::ENUM(_), PhysicalType::BYTE_ARRAY) => {}
- (LogicalType::DECIMAL(t), _) => {
- // Check that scale and precision are consistent with legacy values
- if t.scale != self.scale {
- return Err(general_err!(
- "DECIMAL logical type scale {} must match self.scale {}",
- t.scale,
- self.scale
- ));
- }
- if t.precision != self.precision {
- return Err(general_err!(
- "DECIMAL logical type precision {} must match self.precision {}",
- t.precision,
- self.precision
- ));
- }
- self.check_decimal_precision_scale()?;
- }
- (LogicalType::DATE(_), PhysicalType::INT32) => {}
- (
- LogicalType::TIME(TimeType {
- unit: TimeUnit::MILLIS(_),
- ..
- }),
- PhysicalType::INT32,
- ) => {}
- (LogicalType::TIME(t), PhysicalType::INT64) => {
- if t.unit == TimeUnit::MILLIS(Default::default()) {
- return Err(general_err!(
- "Cannot use millisecond unit on INT64 type"
- ));
- }
- }
- (LogicalType::TIMESTAMP(_), PhysicalType::INT64) => {}
- (LogicalType::INTEGER(t), PhysicalType::INT32)
- if t.bit_width <= 32 => {}
- (LogicalType::INTEGER(t), PhysicalType::INT64)
- if t.bit_width == 64 => {}
- // Null type
- (LogicalType::UNKNOWN(_), PhysicalType::INT32) => {}
- (LogicalType::STRING(_), PhysicalType::BYTE_ARRAY) => {}
- (LogicalType::JSON(_), PhysicalType::BYTE_ARRAY) => {}
- (LogicalType::BSON(_), PhysicalType::BYTE_ARRAY) => {}
- (LogicalType::UUID(_), PhysicalType::FIXED_LEN_BYTE_ARRAY) => {}
- (a, b) => {
- return Err(general_err!(
- "Cannot annotate {:?} from {} fields",
- a,
- b
- ))
- }
- }
- }
- None => {}
- }
-
- match self.converted_type {
- ConvertedType::NONE => {}
- ConvertedType::UTF8 | ConvertedType::BSON | ConvertedType::JSON => {
- if self.physical_type != PhysicalType::BYTE_ARRAY {
- return Err(general_err!(
- "{} can only annotate BYTE_ARRAY fields",
- self.converted_type
- ));
- }
- }
- ConvertedType::DECIMAL => {
- self.check_decimal_precision_scale()?;
- }
- ConvertedType::DATE
- | ConvertedType::TIME_MILLIS
- | ConvertedType::UINT_8
- | ConvertedType::UINT_16
- | ConvertedType::UINT_32
- | ConvertedType::INT_8
- | ConvertedType::INT_16
- | ConvertedType::INT_32 => {
- if self.physical_type != PhysicalType::INT32 {
- return Err(general_err!(
- "{} can only annotate INT32",
- self.converted_type
- ));
- }
- }
- ConvertedType::TIME_MICROS
- | ConvertedType::TIMESTAMP_MILLIS
- | ConvertedType::TIMESTAMP_MICROS
- | ConvertedType::UINT_64
- | ConvertedType::INT_64 => {
- if self.physical_type != PhysicalType::INT64 {
- return Err(general_err!(
- "{} can only annotate INT64",
- self.converted_type
- ));
- }
- }
- ConvertedType::INTERVAL => {
- if self.physical_type != PhysicalType::FIXED_LEN_BYTE_ARRAY
- || self.length != 12
- {
- return Err(general_err!(
- "INTERVAL can only annotate FIXED_LEN_BYTE_ARRAY(12)"
- ));
- }
- }
- ConvertedType::ENUM => {
- if self.physical_type != PhysicalType::BYTE_ARRAY {
- return Err(general_err!("ENUM can only annotate BYTE_ARRAY fields"));
- }
- }
- _ => {
- return Err(general_err!(
- "{} cannot be applied to a primitive type",
- self.converted_type
- ));
- }
- }
-
- Ok(Type::PrimitiveType {
- basic_info,
- physical_type: self.physical_type,
- type_length: self.length,
- scale: self.scale,
- precision: self.precision,
- })
- }
-
- #[inline]
- fn check_decimal_precision_scale(&self) -> Result<()> {
- match self.physical_type {
- PhysicalType::INT32
- | PhysicalType::INT64
- | PhysicalType::BYTE_ARRAY
- | PhysicalType::FIXED_LEN_BYTE_ARRAY => (),
- _ => {
- return Err(general_err!(
- "DECIMAL can only annotate INT32, INT64, BYTE_ARRAY and FIXED_LEN_BYTE_ARRAY"
- ));
- }
- }
-
- // Precision is required and must be a non-zero positive integer.
- if self.precision < 1 {
- return Err(general_err!(
- "Invalid DECIMAL precision: {}",
- self.precision
- ));
- }
-
- // Scale must be zero or a positive integer less than the precision.
- if self.scale < 0 {
- return Err(general_err!("Invalid DECIMAL scale: {}", self.scale));
- }
-
- if self.scale >= self.precision {
- return Err(general_err!(
- "Invalid DECIMAL: scale ({}) cannot be greater than or equal to precision \
- ({})",
- self.scale,
- self.precision
- ));
- }
-
- // Check precision and scale based on physical type limitations.
- match self.physical_type {
- PhysicalType::INT32 => {
- if self.precision > 9 {
- return Err(general_err!(
- "Cannot represent INT32 as DECIMAL with precision {}",
- self.precision
- ));
- }
- }
- PhysicalType::INT64 => {
- if self.precision > 18 {
- return Err(general_err!(
- "Cannot represent INT64 as DECIMAL with precision {}",
- self.precision
- ));
- }
- }
- PhysicalType::FIXED_LEN_BYTE_ARRAY => {
- let max_precision =
- (2f64.powi(8 * self.length - 1) - 1f64).log10().floor() as i32;
-
- if self.precision > max_precision {
- return Err(general_err!(
- "Cannot represent FIXED_LEN_BYTE_ARRAY as DECIMAL with length {} and \
- precision {}. The max precision can only be {}",
- self.length,
- self.precision,
- max_precision
- ));
- }
- }
- _ => (), // For BYTE_ARRAY precision is not limited
- }
-
- Ok(())
- }
-}
-
-/// A builder for group types. All attributes are optional except the name.
-/// Note that if not specified explicitly, `None` is used as the repetition of the group,
-/// which means it is a root (message) type.
-pub struct GroupTypeBuilder<'a> {
- name: &'a str,
- repetition: Option<Repetition>,
- converted_type: ConvertedType,
- logical_type: Option<LogicalType>,
- fields: Vec<TypePtr>,
- id: Option<i32>,
-}
-
-impl<'a> GroupTypeBuilder<'a> {
- /// Creates new group type builder with provided field name.
- pub fn new(name: &'a str) -> Self {
- Self {
- name,
- repetition: None,
- converted_type: ConvertedType::NONE,
- logical_type: None,
- fields: Vec::new(),
- id: None,
- }
- }
-
- /// Sets [`Repetition`](crate::basic::Repetition) for this field and returns itself.
- pub fn with_repetition(mut self, repetition: Repetition) -> Self {
- self.repetition = Some(repetition);
- self
- }
-
- /// Sets [`ConvertedType`](crate::basic::ConvertedType) for this field and returns itself.
- pub fn with_converted_type(mut self, converted_type: ConvertedType) -> Self {
- self.converted_type = converted_type;
- self
- }
-
- /// Sets [`LogicalType`](crate::basic::LogicalType) for this field and returns itself.
- pub fn with_logical_type(mut self, logical_type: Option<LogicalType>) -> Self {
- self.logical_type = logical_type;
- self
- }
-
- /// Sets a list of fields that should be child nodes of this field.
- /// Returns updated self.
- pub fn with_fields(mut self, fields: &mut Vec<TypePtr>) -> Self {
- self.fields.append(fields);
- self
- }
-
- /// Sets optional field id and returns itself.
- pub fn with_id(mut self, id: i32) -> Self {
- self.id = Some(id);
- self
- }
-
- /// Creates a new `GroupType` instance from the gathered attributes.
- pub fn build(self) -> Result<Type> {
- let mut basic_info = BasicTypeInfo {
- name: String::from(self.name),
- repetition: self.repetition,
- converted_type: self.converted_type,
- logical_type: self.logical_type.clone(),
- id: self.id,
- };
- // Populate the converted type if only the logical type is populated
- if self.logical_type.is_some() && self.converted_type == ConvertedType::NONE {
- basic_info.converted_type = self.logical_type.into();
- }
- Ok(Type::GroupType {
- basic_info,
- fields: self.fields,
- })
- }
-}
-
-/// Basic type info. This contains information such as the name of the type,
-/// the repetition level, the logical type and the kind of the type (group, primitive).
-#[derive(Clone, Debug, PartialEq)]
-pub struct BasicTypeInfo {
- name: String,
- repetition: Option<Repetition>,
- converted_type: ConvertedType,
- logical_type: Option<LogicalType>,
- id: Option<i32>,
-}
-
-impl BasicTypeInfo {
- /// Returns field name.
- pub fn name(&self) -> &str {
- &self.name
- }
-
- /// Returns `true` if type has repetition field set, `false` otherwise.
- /// This is mostly applied to group type, because primitive type always has
- /// repetition set.
- pub fn has_repetition(&self) -> bool {
- self.repetition.is_some()
- }
-
- /// Returns [`Repetition`](crate::basic::Repetition) value for the type.
- pub fn repetition(&self) -> Repetition {
- assert!(self.repetition.is_some());
- self.repetition.unwrap()
- }
-
- /// Returns [`ConvertedType`](crate::basic::ConvertedType) value for the type.
- pub fn converted_type(&self) -> ConvertedType {
- self.converted_type
- }
-
- /// Returns [`LogicalType`](crate::basic::LogicalType) value for the type.
- pub fn logical_type(&self) -> Option<LogicalType> {
- // Unlike ConvertedType, LogicalType cannot implement Copy, thus we clone it
- self.logical_type.clone()
- }
-
- /// Returns `true` if id is set, `false` otherwise.
- pub fn has_id(&self) -> bool {
- self.id.is_some()
- }
-
- /// Returns id value for the type.
- pub fn id(&self) -> i32 {
- assert!(self.id.is_some());
- self.id.unwrap()
- }
-}
-
-// ----------------------------------------------------------------------
-// Parquet descriptor definitions
-
-/// Represents a path in a nested schema
-#[derive(Clone, PartialEq, Debug, Eq, Hash)]
-pub struct ColumnPath {
- parts: Vec<String>,
-}
-
-impl ColumnPath {
- /// Creates new column path from vector of field names.
- pub fn new(parts: Vec<String>) -> Self {
- ColumnPath { parts }
- }
-
- /// Returns string representation of this column path.
- /// ```rust
- /// use parquet::schema::types::ColumnPath;
- ///
- /// let path = ColumnPath::new(vec!["a".to_string(), "b".to_string(), "c".to_string()]);
- /// assert_eq!(&path.string(), "a.b.c");
- /// ```
- pub fn string(&self) -> String {
- self.parts.join(".")
- }
-
- /// Appends more components to end of column path.
- /// ```rust
- /// use parquet::schema::types::ColumnPath;
- ///
- /// let mut path = ColumnPath::new(vec!["a".to_string(), "b".to_string(), "c"
- /// .to_string()]);
- /// assert_eq!(&path.string(), "a.b.c");
- ///
- /// path.append(vec!["d".to_string(), "e".to_string()]);
- /// assert_eq!(&path.string(), "a.b.c.d.e");
- /// ```
- pub fn append(&mut self, mut tail: Vec<String>) {
- self.parts.append(&mut tail);
- }
-
- pub fn parts(&self) -> &[String] {
- &self.parts
- }
-}
-
-impl fmt::Display for ColumnPath {
- fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
- write!(f, "{:?}", self.string())
- }
-}
-
-impl From<Vec<String>> for ColumnPath {
- fn from(parts: Vec<String>) -> Self {
- ColumnPath { parts }
- }
-}
-
-impl<'a> From<&'a str> for ColumnPath {
- fn from(single_path: &str) -> Self {
- let s = String::from(single_path);
- ColumnPath::from(s)
- }
-}
-
-impl From<String> for ColumnPath {
- fn from(single_path: String) -> Self {
- let v = vec![single_path];
- ColumnPath { parts: v }
- }
-}
-
-impl AsRef<[String]> for ColumnPath {
- fn as_ref(&self) -> &[String] {
- &self.parts
- }
-}
-
-/// A descriptor for leaf-level primitive columns.
-/// This encapsulates information such as definition and repetition levels and is used to
-/// re-assemble nested data.
-#[derive(Debug, PartialEq)]
-pub struct ColumnDescriptor {
- // The "leaf" primitive type of this column
- primitive_type: TypePtr,
-
- // The maximum definition level for this column
- max_def_level: i16,
-
- // The maximum repetition level for this column
- max_rep_level: i16,
-
- // The path of this column. For instance, "a.b.c.d".
- path: ColumnPath,
-}
-
-impl ColumnDescriptor {
- /// Creates new descriptor for leaf-level column.
- pub fn new(
- primitive_type: TypePtr,
- max_def_level: i16,
- max_rep_level: i16,
- path: ColumnPath,
- ) -> Self {
- Self {
- primitive_type,
- max_def_level,
- max_rep_level,
- path,
- }
- }
-
- /// Returns maximum definition level for this column.
- #[inline]
- pub fn max_def_level(&self) -> i16 {
- self.max_def_level
- }
-
- /// Returns maximum repetition level for this column.
- #[inline]
- pub fn max_rep_level(&self) -> i16 {
- self.max_rep_level
- }
-
- /// Returns [`ColumnPath`] for this column.
- pub fn path(&self) -> &ColumnPath {
- &self.path
- }
-
- /// Returns self type [`Type`](crate::schema::types::Type) for this leaf column.
- pub fn self_type(&self) -> &Type {
- self.primitive_type.as_ref()
- }
-
- /// Returns self type [`TypePtr`](crate::schema::types::TypePtr) for this leaf
- /// column.
- pub fn self_type_ptr(&self) -> TypePtr {
- self.primitive_type.clone()
- }
-
- /// Returns column name.
- pub fn name(&self) -> &str {
- self.primitive_type.name()
- }
-
- /// Returns [`ConvertedType`](crate::basic::ConvertedType) for this column.
- pub fn converted_type(&self) -> ConvertedType {
- self.primitive_type.get_basic_info().converted_type()
- }
-
- /// Returns [`LogicalType`](crate::basic::LogicalType) for this column.
- pub fn logical_type(&self) -> Option<LogicalType> {
- self.primitive_type.get_basic_info().logical_type()
- }
-
- /// Returns physical type for this column.
- /// Note that it will panic if called on a non-primitive type.
- pub fn physical_type(&self) -> PhysicalType {
- match self.primitive_type.as_ref() {
- Type::PrimitiveType { physical_type, .. } => *physical_type,
- _ => panic!("Expected primitive type!"),
- }
- }
-
- /// Returns type length for this column.
- /// Note that it will panic if called on a non-primitive type.
- pub fn type_length(&self) -> i32 {
- match self.primitive_type.as_ref() {
- Type::PrimitiveType { type_length, .. } => *type_length,
- _ => panic!("Expected primitive type!"),
- }
- }
-
- /// Returns type precision for this column.
- /// Note that it will panic if called on a non-primitive type.
- pub fn type_precision(&self) -> i32 {
- match self.primitive_type.as_ref() {
- Type::PrimitiveType { precision, .. } => *precision,
- _ => panic!("Expected primitive type!"),
- }
- }
-
- /// Returns type scale for this column.
- /// Note that it will panic if called on a non-primitive type.
- pub fn type_scale(&self) -> i32 {
- match self.primitive_type.as_ref() {
- Type::PrimitiveType { scale, .. } => *scale,
- _ => panic!("Expected primitive type!"),
- }
- }
-}
-
-/// A schema descriptor. This encapsulates the top-level schemas for all the columns,
-/// as well as all descriptors for all the primitive columns.
-pub struct SchemaDescriptor {
- // The top-level schema (the "message" type).
- // This must be a `GroupType` where each field is a root column type in the schema.
- schema: TypePtr,
-
- // All the descriptors for primitive columns in this schema, constructed from
- // `schema` in DFS order.
- leaves: Vec<ColumnDescPtr>,
-
- // Mapping from a leaf column's index to the root column type that it
- // comes from. For instance: the leaf `a.b.c.d` would have a link back to `a`:
- // -- a <-----+
- // -- -- b |
- // -- -- -- c |
- // -- -- -- -- d
- leaf_to_base: Vec<TypePtr>,
-}
-
-impl fmt::Debug for SchemaDescriptor {
- fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
- // Skip leaves and leaf_to_base as they only a cache information already found in `schema`
- f.debug_struct("SchemaDescriptor")
- .field("schema", &self.schema)
- .finish()
- }
-}
-
-impl SchemaDescriptor {
- /// Creates new schema descriptor from Parquet schema.
- pub fn new(tp: TypePtr) -> Self {
- assert!(tp.is_group(), "SchemaDescriptor should take a GroupType");
- let mut leaves = vec![];
- let mut leaf_to_base = Vec::new();
- for f in tp.get_fields() {
- let mut path = vec![];
- build_tree(f, f, 0, 0, &mut leaves, &mut leaf_to_base, &mut path);
- }
-
- Self {
- schema: tp,
- leaves,
- leaf_to_base,
- }
- }
-
- /// Returns [`ColumnDescriptor`] for a field position.
- pub fn column(&self, i: usize) -> ColumnDescPtr {
- assert!(
- i < self.leaves.len(),
- "Index out of bound: {} not in [0, {})",
- i,
- self.leaves.len()
- );
- self.leaves[i].clone()
- }
-
- /// Returns slice of [`ColumnDescriptor`].
- pub fn columns(&self) -> &[ColumnDescPtr] {
- &self.leaves
- }
-
- /// Returns number of leaf-level columns.
- pub fn num_columns(&self) -> usize {
- self.leaves.len()
- }
-
- /// Returns column root [`Type`](crate::schema::types::Type) for a field position.
- pub fn get_column_root(&self, i: usize) -> &Type {
- let result = self.column_root_of(i);
- result.as_ref()
- }
-
- /// Returns column root [`Type`](crate::schema::types::Type) pointer for a field
- /// position.
- pub fn get_column_root_ptr(&self, i: usize) -> TypePtr {
- let result = self.column_root_of(i);
- result.clone()
- }
-
- fn column_root_of(&self, i: usize) -> &Arc<Type> {
- assert!(
- i < self.leaves.len(),
- "Index out of bound: {} not in [0, {})",
- i,
- self.leaves.len()
- );
-
- self.leaf_to_base
- .get(i)
- .unwrap_or_else(|| panic!("Expected a value for index {} but found None", i))
- }
-
- /// Returns schema as [`Type`](crate::schema::types::Type).
- pub fn root_schema(&self) -> &Type {
- self.schema.as_ref()
- }
-
- pub fn root_schema_ptr(&self) -> TypePtr {
- self.schema.clone()
- }
-
- /// Returns schema name.
- pub fn name(&self) -> &str {
- self.schema.name()
- }
-}
-
-fn build_tree<'a>(
- tp: &'a TypePtr,
- base_tp: &TypePtr,
- mut max_rep_level: i16,
- mut max_def_level: i16,
- leaves: &mut Vec<ColumnDescPtr>,
- leaf_to_base: &mut Vec<TypePtr>,
- path_so_far: &mut Vec<&'a str>,
-) {
- assert!(tp.get_basic_info().has_repetition());
-
- path_so_far.push(tp.name());
- match tp.get_basic_info().repetition() {
- Repetition::OPTIONAL => {
- max_def_level += 1;
- }
- Repetition::REPEATED => {
- max_def_level += 1;
- max_rep_level += 1;
- }
- _ => {}
- }
-
- match tp.as_ref() {
- Type::PrimitiveType { .. } => {
- let mut path: Vec<String> = vec![];
- path.extend(path_so_far.iter().copied().map(String::from));
- leaves.push(Arc::new(ColumnDescriptor::new(
- tp.clone(),
- max_def_level,
- max_rep_level,
- ColumnPath::new(path),
- )));
- leaf_to_base.push(base_tp.clone());
- }
- Type::GroupType { ref fields, .. } => {
- for f in fields {
- build_tree(
- f,
- base_tp,
- max_rep_level,
- max_def_level,
- leaves,
- leaf_to_base,
- path_so_far,
- );
- path_so_far.pop();
- }
- }
- }
-}
-
-/// Method to convert from Thrift.
-pub fn from_thrift(elements: &[SchemaElement]) -> Result<TypePtr> {
- let mut index = 0;
- let mut schema_nodes = Vec::new();
- while index < elements.len() {
- let t = from_thrift_helper(elements, index)?;
- index = t.0;
- schema_nodes.push(t.1);
- }
- if schema_nodes.len() != 1 {
- return Err(general_err!(
- "Expected exactly one root node, but found {}",
- schema_nodes.len()
- ));
- }
-
- Ok(schema_nodes.remove(0))
-}
-
-/// Constructs a new Type from the `elements`, starting at index `index`.
-/// The first result is the starting index for the next Type after this one. If it is
-/// equal to `elements.len()`, then this Type is the last one.
-/// The second result is the result Type.
-fn from_thrift_helper(
- elements: &[SchemaElement],
- index: usize,
-) -> Result<(usize, TypePtr)> {
- // Whether or not the current node is root (message type).
- // There is only one message type node in the schema tree.
- let is_root_node = index == 0;
-
- if index > elements.len() {
- return Err(general_err!(
- "Index out of bound, index = {}, len = {}",
- index,
- elements.len()
- ));
- }
- let element = &elements[index];
- let converted_type = ConvertedType::from(element.converted_type);
- // LogicalType is only present in v2 Parquet files. ConvertedType is always
- // populated, regardless of the version of the file (v1 or v2).
- let logical_type = element
- .logical_type
- .as_ref()
- .map(|value| LogicalType::from(value.clone()));
- let field_id = elements[index].field_id;
- match elements[index].num_children {
- // From parquet-format:
- // The children count is used to construct the nested relationship.
- // This field is not set when the element is a primitive type
- // Sometimes parquet-cpp sets num_children field to 0 for primitive types, so we
- // have to handle this case too.
- None | Some(0) => {
- // primitive type
- if elements[index].repetition_type.is_none() {
- return Err(general_err!(
- "Repetition level must be defined for a primitive type"
- ));
- }
- let repetition = Repetition::from(elements[index].repetition_type.unwrap());
- let physical_type = PhysicalType::from(elements[index].type_.unwrap());
- let length = elements[index].type_length.unwrap_or(-1);
- let scale = elements[index].scale.unwrap_or(-1);
- let precision = elements[index].precision.unwrap_or(-1);
- let name = &elements[index].name;
- let mut builder = Type::primitive_type_builder(name, physical_type)
- .with_repetition(repetition)
- .with_converted_type(converted_type)
- .with_logical_type(logical_type)
- .with_length(length)
- .with_precision(precision)
- .with_scale(scale);
- if let Some(id) = field_id {
- builder = builder.with_id(id);
- }
- Ok((index + 1, Arc::new(builder.build()?)))
- }
- Some(n) => {
- let repetition = elements[index].repetition_type.map(Repetition::from);
- let mut fields = vec![];
- let mut next_index = index + 1;
- for _ in 0..n {
- let child_result = from_thrift_helper(elements, next_index as usize)?;
- next_index = child_result.0;
- fields.push(child_result.1);
- }
-
- let mut builder = Type::group_type_builder(&elements[index].name)
- .with_converted_type(converted_type)
- .with_logical_type(logical_type)
- .with_fields(&mut fields);
- if let Some(rep) = repetition {
- // Sometimes parquet-cpp and parquet-mr set repetition level REQUIRED or
- // REPEATED for root node.
- //
- // We only set repetition for group types that are not top-level message
- // type. According to parquet-format:
- // Root of the schema does not have a repetition_type.
- // All other types must have one.
- if !is_root_node {
- builder = builder.with_repetition(rep);
- }
- }
- if let Some(id) = field_id {
- builder = builder.with_id(id);
- }
- Ok((next_index, Arc::new(builder.build().unwrap())))
- }
- }
-}
-
-/// Method to convert to Thrift.
-pub fn to_thrift(schema: &Type) -> Result<Vec<SchemaElement>> {
- if !schema.is_group() {
- return Err(general_err!("Root schema must be Group type"));
- }
- let mut elements: Vec<SchemaElement> = Vec::new();
- to_thrift_helper(schema, &mut elements);
- Ok(elements)
-}
-
-/// Constructs list of `SchemaElement` from the schema using depth-first traversal.
-/// Here we assume that schema is always valid and starts with group type.
-fn to_thrift_helper(schema: &Type, elements: &mut Vec<SchemaElement>) {
- match *schema {
- Type::PrimitiveType {
- ref basic_info,
- physical_type,
- type_length,
- scale,
- precision,
- } => {
- let element = SchemaElement {
- type_: Some(physical_type.into()),
- type_length: if type_length >= 0 {
- Some(type_length)
- } else {
- None
- },
- repetition_type: Some(basic_info.repetition().into()),
- name: basic_info.name().to_owned(),
- num_children: None,
- converted_type: basic_info.converted_type().into(),
- scale: if scale >= 0 { Some(scale) } else { None },
- precision: if precision >= 0 {
- Some(precision)
- } else {
- None
- },
- field_id: if basic_info.has_id() {
- Some(basic_info.id())
- } else {
- None
- },
- logical_type: basic_info.logical_type().map(|value| value.into()),
- };
-
- elements.push(element);
- }
- Type::GroupType {
- ref basic_info,
- ref fields,
- } => {
- let repetition = if basic_info.has_repetition() {
- Some(basic_info.repetition().into())
- } else {
- None
- };
-
- let element = SchemaElement {
- type_: None,
- type_length: None,
- repetition_type: repetition,
- name: basic_info.name().to_owned(),
- num_children: Some(fields.len() as i32),
- converted_type: basic_info.converted_type().into(),
- scale: None,
- precision: None,
- field_id: if basic_info.has_id() {
- Some(basic_info.id())
- } else {
- None
- },
- logical_type: basic_info.logical_type().map(|value| value.into()),
- };
-
- elements.push(element);
-
- // Add child elements for a group
- for field in fields {
- to_thrift_helper(field, elements);
- }
- }
- }
-}
-
-#[cfg(test)]
-mod tests {
- use super::*;
-
- use crate::basic::{DecimalType, IntType};
- use crate::schema::parser::parse_message_type;
-
- // TODO: add tests for v2 types
-
- #[test]
- fn test_primitive_type() {
- let mut result = Type::primitive_type_builder("foo", PhysicalType::INT32)
- .with_logical_type(Some(LogicalType::INTEGER(IntType {
- bit_width: 32,
- is_signed: true,
- })))
- .with_id(0)
- .build();
- assert!(result.is_ok());
-
- if let Ok(tp) = result {
- assert!(tp.is_primitive());
- assert!(!tp.is_group());
- let basic_info = tp.get_basic_info();
- assert_eq!(basic_info.repetition(), Repetition::OPTIONAL);
- assert_eq!(
- basic_info.logical_type(),
- Some(LogicalType::INTEGER(IntType {
- bit_width: 32,
- is_signed: true
- }))
- );
- assert_eq!(basic_info.converted_type(), ConvertedType::INT_32);
- assert_eq!(basic_info.id(), 0);
- match tp {
- Type::PrimitiveType { physical_type, .. } => {
- assert_eq!(physical_type, PhysicalType::INT32);
- }
- _ => panic!(),
- }
- }
-
- // Test illegal inputs with logical type
- result = Type::primitive_type_builder("foo", PhysicalType::INT64)
- .with_repetition(Repetition::REPEATED)
- .with_logical_type(Some(LogicalType::INTEGER(IntType {
- is_signed: true,
- bit_width: 8,
- })))
- .build();
- assert!(result.is_err());
- if let Err(e) = result {
- assert_eq!(
- format!("{}", e),
- "Parquet error: Cannot annotate INTEGER(IntType { bit_width: 8, is_signed: true }) from INT64 fields"
- );
- }
-
- // Test illegal inputs with converted type
- result = Type::primitive_type_builder("foo", PhysicalType::INT64)
- .with_repetition(Repetition::REPEATED)
- .with_converted_type(ConvertedType::BSON)
- .build();
- assert!(result.is_err());
- if let Err(e) = result {
- assert_eq!(
- format!("{}", e),
- "Parquet error: BSON can only annotate BYTE_ARRAY fields"
- );
- }
-
- result = Type::primitive_type_builder("foo", PhysicalType::INT96)
- .with_repetition(Repetition::REQUIRED)
- .with_converted_type(ConvertedType::DECIMAL)
- .with_precision(-1)
- .with_scale(-1)
- .build();
- assert!(result.is_err());
- if let Err(e) = result {
- assert_eq!(
- format!("{}", e),
- "Parquet error: DECIMAL can only annotate INT32, INT64, BYTE_ARRAY and FIXED_LEN_BYTE_ARRAY"
- );
- }
-
- result = Type::primitive_type_builder("foo", PhysicalType::BYTE_ARRAY)
- .with_repetition(Repetition::REQUIRED)
- .with_logical_type(Some(LogicalType::DECIMAL(DecimalType {
- scale: 32,
- precision: 12,
- })))
- .with_precision(-1)
- .with_scale(-1)
- .build();
- assert!(result.is_err());
- if let Err(e) = result {
- assert_eq!(
- format!("{}", e),
- "Parquet error: DECIMAL logical type scale 32 must match self.scale -1"
- );
- }
-
- result = Type::primitive_type_builder("foo", PhysicalType::BYTE_ARRAY)
- .with_repetition(Repetition::REQUIRED)
- .with_converted_type(ConvertedType::DECIMAL)
- .with_precision(-1)
- .with_scale(-1)
- .build();
- assert!(result.is_err());
- if let Err(e) = result {
- assert_eq!(
- format!("{}", e),
- "Parquet error: Invalid DECIMAL precision: -1"
- );
- }
-
- result = Type::primitive_type_builder("foo", PhysicalType::BYTE_ARRAY)
- .with_repetition(Repetition::REQUIRED)
- .with_converted_type(ConvertedType::DECIMAL)
- .with_precision(0)
- .with_scale(-1)
- .build();
- assert!(result.is_err());
- if let Err(e) = result {
- assert_eq!(
- format!("{}", e),
- "Parquet error: Invalid DECIMAL precision: 0"
- );
- }
-
- result = Type::primitive_type_builder("foo", PhysicalType::BYTE_ARRAY)
- .with_repetition(Repetition::REQUIRED)
- .with_converted_type(ConvertedType::DECIMAL)
- .with_precision(1)
- .with_scale(-1)
- .build();
- assert!(result.is_err());
- if let Err(e) = result {
- assert_eq!(format!("{}", e), "Parquet error: Invalid DECIMAL scale: -1");
- }
-
- result = Type::primitive_type_builder("foo", PhysicalType::BYTE_ARRAY)
- .with_repetition(Repetition::REQUIRED)
- .with_converted_type(ConvertedType::DECIMAL)
- .with_precision(1)
- .with_scale(2)
- .build();
- assert!(result.is_err());
- if let Err(e) = result {
- assert_eq!(
- format!("{}", e),
- "Parquet error: Invalid DECIMAL: scale (2) cannot be greater than or equal to precision (1)"
- );
- }
-
- result = Type::primitive_type_builder("foo", PhysicalType::INT32)
- .with_repetition(Repetition::REQUIRED)
- .with_converted_type(ConvertedType::DECIMAL)
- .with_precision(18)
- .with_scale(2)
- .build();
- assert!(result.is_err());
- if let Err(e) = result {
- assert_eq!(
- format!("{}", e),
- "Parquet error: Cannot represent INT32 as DECIMAL with precision 18"
- );
- }
-
- result = Type::primitive_type_builder("foo", PhysicalType::INT64)
- .with_repetition(Repetition::REQUIRED)
- .with_converted_type(ConvertedType::DECIMAL)
- .with_precision(32)
- .with_scale(2)
- .build();
- assert!(result.is_err());
- if let Err(e) = result {
- assert_eq!(
- format!("{}", e),
- "Parquet error: Cannot represent INT64 as DECIMAL with precision 32"
- );
- }
-
- result = Type::primitive_type_builder("foo", PhysicalType::FIXED_LEN_BYTE_ARRAY)
- .with_repetition(Repetition::REQUIRED)
- .with_converted_type(ConvertedType::DECIMAL)
- .with_length(5)
- .with_precision(12)
- .with_scale(2)
- .build();
- assert!(result.is_err());
- if let Err(e) = result {
- assert_eq!(
- format!("{}", e),
- "Parquet error: Cannot represent FIXED_LEN_BYTE_ARRAY as DECIMAL with length 5 and precision 12. The max precision can only be 11"
- );
- }
-
- result = Type::primitive_type_builder("foo", PhysicalType::INT64)
- .with_repetition(Repetition::REQUIRED)
- .with_converted_type(ConvertedType::UINT_8)
- .build();
- assert!(result.is_err());
- if let Err(e) = result {
- assert_eq!(
- format!("{}", e),
- "Parquet error: UINT_8 can only annotate INT32"
- );
- }
-
- result = Type::primitive_type_builder("foo", PhysicalType::INT32)
- .with_repetition(Repetition::REQUIRED)
- .with_converted_type(ConvertedType::TIME_MICROS)
- .build();
- assert!(result.is_err());
- if let Err(e) = result {
- assert_eq!(
- format!("{}", e),
- "Parquet error: TIME_MICROS can only annotate INT64"
- );
- }
-
- result = Type::primitive_type_builder("foo", PhysicalType::BYTE_ARRAY)
- .with_repetition(Repetition::REQUIRED)
- .with_converted_type(ConvertedType::INTERVAL)
- .build();
- assert!(result.is_err());
- if let Err(e) = result {
- assert_eq!(
- format!("{}", e),
- "Parquet error: INTERVAL can only annotate FIXED_LEN_BYTE_ARRAY(12)"
- );
- }
-
- result = Type::primitive_type_builder("foo", PhysicalType::FIXED_LEN_BYTE_ARRAY)
- .with_repetition(Repetition::REQUIRED)
- .with_converted_type(ConvertedType::INTERVAL)
- .with_length(1)
- .build();
- assert!(result.is_err());
- if let Err(e) = result {
- assert_eq!(
- format!("{}", e),
- "Parquet error: INTERVAL can only annotate FIXED_LEN_BYTE_ARRAY(12)"
- );
- }
-
- result = Type::primitive_type_builder("foo", PhysicalType::INT32)
- .with_repetition(Repetition::REQUIRED)
- .with_converted_type(ConvertedType::ENUM)
- .build();
- assert!(result.is_err());
- if let Err(e) = result {
- assert_eq!(
- format!("{}", e),
- "Parquet error: ENUM can only annotate BYTE_ARRAY fields"
- );
- }
-
- result = Type::primitive_type_builder("foo", PhysicalType::INT32)
- .with_repetition(Repetition::REQUIRED)
- .with_converted_type(ConvertedType::MAP)
- .build();
- assert!(result.is_err());
- if let Err(e) = result {
- assert_eq!(
- format!("{}", e),
- "Parquet error: MAP cannot be applied to a primitive type"
- );
- }
-
- result = Type::primitive_type_builder("foo", PhysicalType::FIXED_LEN_BYTE_ARRAY)
- .with_repetition(Repetition::REQUIRED)
- .with_converted_type(ConvertedType::DECIMAL)
- .with_length(-1)
- .build();
- assert!(result.is_err());
- if let Err(e) = result {
- assert_eq!(
- format!("{}", e),
- "Parquet error: Invalid FIXED_LEN_BYTE_ARRAY length: -1"
- );
- }
- }
-
- #[test]
- fn test_group_type() {
- let f1 = Type::primitive_type_builder("f1", PhysicalType::INT32)
- .with_converted_type(ConvertedType::INT_32)
- .with_id(0)
- .build();
- assert!(f1.is_ok());
- let f2 = Type::primitive_type_builder("f2", PhysicalType::BYTE_ARRAY)
- .with_converted_type(ConvertedType::UTF8)
- .with_id(1)
- .build();
- assert!(f2.is_ok());
-
- let mut fields = vec![];
- fields.push(Arc::new(f1.unwrap()));
- fields.push(Arc::new(f2.unwrap()));
-
- let result = Type::group_type_builder("foo")
- .with_repetition(Repetition::REPEATED)
- .with_logical_type(Some(LogicalType::LIST(Default::default())))
- .with_fields(&mut fields)
- .with_id(1)
- .build();
- assert!(result.is_ok());
-
- let tp = result.unwrap();
- let basic_info = tp.get_basic_info();
- assert!(tp.is_group());
- assert!(!tp.is_primitive());
- assert_eq!(basic_info.repetition(), Repetition::REPEATED);
- assert_eq!(
- basic_info.logical_type(),
- Some(LogicalType::LIST(Default::default()))
- );
- assert_eq!(basic_info.converted_type(), ConvertedType::LIST);
- assert_eq!(basic_info.id(), 1);
- assert_eq!(tp.get_fields().len(), 2);
- assert_eq!(tp.get_fields()[0].name(), "f1");
- assert_eq!(tp.get_fields()[1].name(), "f2");
- }
-
- #[test]
- fn test_column_descriptor() {
- let result = test_column_descriptor_helper();
- assert!(
- result.is_ok(),
- "Expected result to be OK but got err:\n {}",
- result.unwrap_err()
- );
- }
-
- fn test_column_descriptor_helper() -> Result<()> {
- let tp = Type::primitive_type_builder("name", PhysicalType::BYTE_ARRAY)
- .with_converted_type(ConvertedType::UTF8)
- .build()?;
-
- let descr = ColumnDescriptor::new(Arc::new(tp), 4, 1, ColumnPath::from("name"));
-
- assert_eq!(descr.path(), &ColumnPath::from("name"));
- assert_eq!(descr.converted_type(), ConvertedType::UTF8);
- assert_eq!(descr.physical_type(), PhysicalType::BYTE_ARRAY);
- assert_eq!(descr.max_def_level(), 4);
- assert_eq!(descr.max_rep_level(), 1);
- assert_eq!(descr.name(), "name");
- assert_eq!(descr.type_length(), -1);
- assert_eq!(descr.type_precision(), -1);
- assert_eq!(descr.type_scale(), -1);
-
- Ok(())
- }
-
- #[test]
- fn test_schema_descriptor() {
- let result = test_schema_descriptor_helper();
- assert!(
- result.is_ok(),
- "Expected result to be OK but got err:\n {}",
- result.unwrap_err()
- );
- }
-
- // A helper fn to avoid handling the results from type creation
- fn test_schema_descriptor_helper() -> Result<()> {
- let mut fields = vec![];
-
- let inta = Type::primitive_type_builder("a", PhysicalType::INT32)
- .with_repetition(Repetition::REQUIRED)
- .with_converted_type(ConvertedType::INT_32)
- .build()?;
- fields.push(Arc::new(inta));
- let intb = Type::primitive_type_builder("b", PhysicalType::INT64)
- .with_converted_type(ConvertedType::INT_64)
- .build()?;
- fields.push(Arc::new(intb));
- let intc = Type::primitive_type_builder("c", PhysicalType::BYTE_ARRAY)
- .with_repetition(Repetition::REPEATED)
- .with_converted_type(ConvertedType::UTF8)
- .build()?;
- fields.push(Arc::new(intc));
-
- // 3-level list encoding
- let item1 = Type::primitive_type_builder("item1", PhysicalType::INT64)
- .with_repetition(Repetition::REQUIRED)
- .with_converted_type(ConvertedType::INT_64)
- .build()?;
- let item2 =
- Type::primitive_type_builder("item2", PhysicalType::BOOLEAN).build()?;
- let item3 = Type::primitive_type_builder("item3", PhysicalType::INT32)
- .with_repetition(Repetition::REPEATED)
- .with_converted_type(ConvertedType::INT_32)
- .build()?;
- let list = Type::group_type_builder("records")
- .with_repetition(Repetition::REPEATED)
- .with_converted_type(ConvertedType::LIST)
- .with_fields(&mut vec![Arc::new(item1), Arc::new(item2), Arc::new(item3)])
- .build()?;
- let bag = Type::group_type_builder("bag")
- .with_repetition(Repetition::OPTIONAL)
- .with_fields(&mut vec![Arc::new(list)])
- .build()?;
- fields.push(Arc::new(bag));
-
- let schema = Type::group_type_builder("schema")
- .with_repetition(Repetition::REPEATED)
- .with_fields(&mut fields)
- .build()?;
- let descr = SchemaDescriptor::new(Arc::new(schema));
-
- let nleaves = 6;
- assert_eq!(descr.num_columns(), nleaves);
-
- // mdef mrep
- // required int32 a 0 0
- // optional int64 b 1 0
- // repeated byte_array c 1 1
- // optional group bag 1 0
- // repeated group records 2 1
- // required int64 item1 2 1
- // optional boolean item2 3 1
- // repeated int32 item3 3 2
- let ex_max_def_levels = vec![0, 1, 1, 2, 3, 3];
- let ex_max_rep_levels = vec![0, 0, 1, 1, 1, 2];
-
- for i in 0..nleaves {
- let col = descr.column(i);
- assert_eq!(col.max_def_level(), ex_max_def_levels[i], "{}", i);
- assert_eq!(col.max_rep_level(), ex_max_rep_levels[i], "{}", i);
- }
-
- assert_eq!(descr.column(0).path().string(), "a");
- assert_eq!(descr.column(1).path().string(), "b");
- assert_eq!(descr.column(2).path().string(), "c");
- assert_eq!(descr.column(3).path().string(), "bag.records.item1");
- assert_eq!(descr.column(4).path().string(), "bag.records.item2");
- assert_eq!(descr.column(5).path().string(), "bag.records.item3");
-
- assert_eq!(descr.get_column_root(0).name(), "a");
- assert_eq!(descr.get_column_root(3).name(), "bag");
- assert_eq!(descr.get_column_root(4).name(), "bag");
- assert_eq!(descr.get_column_root(5).name(), "bag");
-
- Ok(())
- }
-
- #[test]
- fn test_schema_build_tree_def_rep_levels() {
- let message_type = "
- message spark_schema {
- REQUIRED INT32 a;
- OPTIONAL group b {
- OPTIONAL INT32 _1;
- OPTIONAL INT32 _2;
- }
- OPTIONAL group c (LIST) {
- REPEATED group list {
- OPTIONAL INT32 element;
- }
- }
- }
- ";
- let schema = parse_message_type(message_type).expect("should parse schema");
- let descr = SchemaDescriptor::new(Arc::new(schema));
- // required int32 a
- assert_eq!(descr.column(0).max_def_level(), 0);
- assert_eq!(descr.column(0).max_rep_level(), 0);
- // optional int32 b._1
- assert_eq!(descr.column(1).max_def_level(), 2);
- assert_eq!(descr.column(1).max_rep_level(), 0);
- // optional int32 b._2
- assert_eq!(descr.column(2).max_def_level(), 2);
- assert_eq!(descr.column(2).max_rep_level(), 0);
- // repeated optional int32 c.list.element
- assert_eq!(descr.column(3).max_def_level(), 3);
- assert_eq!(descr.column(3).max_rep_level(), 1);
- }
-
- #[test]
- #[should_panic(expected = "Cannot call get_physical_type() on a non-primitive type")]
- fn test_get_physical_type_panic() {
- let list = Type::group_type_builder("records")
- .with_repetition(Repetition::REPEATED)
- .build()
- .unwrap();
- list.get_physical_type();
- }
-
- #[test]
- fn test_get_physical_type_primitive() {
- let f = Type::primitive_type_builder("f", PhysicalType::INT64)
- .build()
- .unwrap();
- assert_eq!(f.get_physical_type(), PhysicalType::INT64);
-
- let f = Type::primitive_type_builder("f", PhysicalType::BYTE_ARRAY)
- .build()
- .unwrap();
- assert_eq!(f.get_physical_type(), PhysicalType::BYTE_ARRAY);
- }
-
- #[test]
- fn test_check_contains_primitive_primitive() {
- // OK
- let f1 = Type::primitive_type_builder("f", PhysicalType::INT32)
- .build()
- .unwrap();
- let f2 = Type::primitive_type_builder("f", PhysicalType::INT32)
- .build()
- .unwrap();
- assert!(f1.check_contains(&f2));
-
- // OK: different logical type does not affect check_contains
- let f1 = Type::primitive_type_builder("f", PhysicalType::INT32)
- .with_converted_type(ConvertedType::UINT_8)
- .build()
- .unwrap();
- let f2 = Type::primitive_type_builder("f", PhysicalType::INT32)
- .with_converted_type(ConvertedType::UINT_16)
- .build()
- .unwrap();
- assert!(f1.check_contains(&f2));
-
- // KO: different name
- let f1 = Type::primitive_type_builder("f1", PhysicalType::INT32)
- .build()
- .unwrap();
- let f2 = Type::primitive_type_builder("f2", PhysicalType::INT32)
- .build()
- .unwrap();
- assert!(!f1.check_contains(&f2));
-
- // KO: different type
- let f1 = Type::primitive_type_builder("f", PhysicalType::INT32)
- .build()
- .unwrap();
- let f2 = Type::primitive_type_builder("f", PhysicalType::INT64)
- .build()
- .unwrap();
- assert!(!f1.check_contains(&f2));
-
- // KO: different repetition
- let f1 = Type::primitive_type_builder("f", PhysicalType::INT32)
- .with_repetition(Repetition::REQUIRED)
- .build()
- .unwrap();
- let f2 = Type::primitive_type_builder("f", PhysicalType::INT32)
- .with_repetition(Repetition::OPTIONAL)
- .build()
- .unwrap();
- assert!(!f1.check_contains(&f2));
- }
-
- // function to create a new group type for testing
- fn test_new_group_type(name: &str, repetition: Repetition, types: Vec<Type>) -> Type {
- let mut fields = Vec::new();
- for tpe in types {
- fields.push(Arc::new(tpe))
- }
- Type::group_type_builder(name)
- .with_repetition(repetition)
- .with_fields(&mut fields)
- .build()
- .unwrap()
- }
-
- #[test]
- fn test_check_contains_group_group() {
- // OK: should match okay with empty fields
- let f1 = Type::group_type_builder("f").build().unwrap();
- let f2 = Type::group_type_builder("f").build().unwrap();
- assert!(f1.check_contains(&f2));
-
- // OK: fields match
- let f1 = test_new_group_type(
- "f",
- Repetition::REPEATED,
- vec![
- Type::primitive_type_builder("f1", PhysicalType::INT32)
- .build()
- .unwrap(),
- Type::primitive_type_builder("f2", PhysicalType::INT64)
- .build()
- .unwrap(),
- ],
- );
- let f2 = test_new_group_type(
- "f",
- Repetition::REPEATED,
- vec![
- Type::primitive_type_builder("f1", PhysicalType::INT32)
- .build()
- .unwrap(),
- Type::primitive_type_builder("f2", PhysicalType::INT64)
- .build()
- .unwrap(),
- ],
- );
- assert!(f1.check_contains(&f2));
-
- // OK: subset of fields
- let f1 = test_new_group_type(
- "f",
- Repetition::REPEATED,
- vec![
- Type::primitive_type_builder("f1", PhysicalType::INT32)
- .build()
- .unwrap(),
- Type::primitive_type_builder("f2", PhysicalType::INT64)
- .build()
- .unwrap(),
- ],
- );
- let f2 = test_new_group_type(
- "f",
- Repetition::REPEATED,
- vec![Type::primitive_type_builder("f2", PhysicalType::INT64)
- .build()
- .unwrap()],
- );
- assert!(f1.check_contains(&f2));
-
- // KO: different name
- let f1 = Type::group_type_builder("f1").build().unwrap();
- let f2 = Type::group_type_builder("f2").build().unwrap();
- assert!(!f1.check_contains(&f2));
-
- // KO: different repetition
- let f1 = Type::group_type_builder("f")
- .with_repetition(Repetition::OPTIONAL)
- .build()
- .unwrap();
- let f2 = Type::group_type_builder("f")
- .with_repetition(Repetition::REPEATED)
- .build()
- .unwrap();
- assert!(!f1.check_contains(&f2));
-
- // KO: different fields
- let f1 = test_new_group_type(
- "f",
- Repetition::REPEATED,
- vec![
- Type::primitive_type_builder("f1", PhysicalType::INT32)
- .build()
- .unwrap(),
- Type::primitive_type_builder("f2", PhysicalType::INT64)
- .build()
- .unwrap(),
- ],
- );
- let f2 = test_new_group_type(
- "f",
- Repetition::REPEATED,
- vec![
- Type::primitive_type_builder("f1", PhysicalType::INT32)
- .build()
- .unwrap(),
- Type::primitive_type_builder("f2", PhysicalType::BOOLEAN)
- .build()
- .unwrap(),
- ],
- );
- assert!(!f1.check_contains(&f2));
-
- // KO: different fields
- let f1 = test_new_group_type(
- "f",
- Repetition::REPEATED,
- vec![
- Type::primitive_type_builder("f1", PhysicalType::INT32)
- .build()
- .unwrap(),
- Type::primitive_type_builder("f2", PhysicalType::INT64)
- .build()
- .unwrap(),
- ],
- );
- let f2 = test_new_group_type(
- "f",
- Repetition::REPEATED,
- vec![Type::primitive_type_builder("f3", PhysicalType::INT32)
- .build()
- .unwrap()],
- );
- assert!(!f1.check_contains(&f2));
- }
-
- #[test]
- fn test_check_contains_group_primitive() {
- // KO: should not match
- let f1 = Type::group_type_builder("f").build().unwrap();
- let f2 = Type::primitive_type_builder("f", PhysicalType::INT64)
- .build()
- .unwrap();
- assert!(!f1.check_contains(&f2));
- assert!(!f2.check_contains(&f1));
-
- // KO: should not match when primitive field is part of group type
- let f1 = test_new_group_type(
- "f",
- Repetition::REPEATED,
- vec![Type::primitive_type_builder("f1", PhysicalType::INT32)
- .build()
- .unwrap()],
- );
- let f2 = Type::primitive_type_builder("f1", PhysicalType::INT32)
- .build()
- .unwrap();
- assert!(!f1.check_contains(&f2));
- assert!(!f2.check_contains(&f1));
-
- // OK: match nested types
- let f1 = test_new_group_type(
- "a",
- Repetition::REPEATED,
- vec![
- test_new_group_type(
- "b",
- Repetition::REPEATED,
- vec![Type::primitive_type_builder("c", PhysicalType::INT32)
- .build()
- .unwrap()],
- ),
- Type::primitive_type_builder("d", PhysicalType::INT64)
- .build()
- .unwrap(),
- Type::primitive_type_builder("e", PhysicalType::BOOLEAN)
- .build()
- .unwrap(),
- ],
- );
- let f2 = test_new_group_type(
- "a",
- Repetition::REPEATED,
- vec![test_new_group_type(
- "b",
- Repetition::REPEATED,
- vec![Type::primitive_type_builder("c", PhysicalType::INT32)
- .build()
- .unwrap()],
- )],
- );
- assert!(f1.check_contains(&f2)); // should match
- assert!(!f2.check_contains(&f1)); // should fail
- }
-
- #[test]
- fn test_schema_type_thrift_conversion_err() {
- let schema = Type::primitive_type_builder("col", PhysicalType::INT32)
- .build()
- .unwrap();
- let thrift_schema = to_thrift(&schema);
- assert!(thrift_schema.is_err());
- if let Err(e) = thrift_schema {
- assert_eq!(
- format!("{}", e),
- "Parquet error: Root schema must be Group type"
- );
- }
- }
-
- #[test]
- fn test_schema_type_thrift_conversion() {
- let message_type = "
- message conversions {
- REQUIRED INT64 id;
- OPTIONAL group int_array_Array (LIST) {
- REPEATED group list {
- OPTIONAL group element (LIST) {
- REPEATED group list {
- OPTIONAL INT32 element;
- }
- }
- }
- }
- OPTIONAL group int_map (MAP) {
- REPEATED group map (MAP_KEY_VALUE) {
- REQUIRED BYTE_ARRAY key (UTF8);
- OPTIONAL INT32 value;
- }
- }
- OPTIONAL group int_Map_Array (LIST) {
- REPEATED group list {
- OPTIONAL group g (MAP) {
- REPEATED group map (MAP_KEY_VALUE) {
- REQUIRED BYTE_ARRAY key (UTF8);
- OPTIONAL group value {
- OPTIONAL group H {
- OPTIONAL group i (LIST) {
- REPEATED group list {
- OPTIONAL DOUBLE element;
- }
- }
- }
- }
- }
- }
- }
- }
- OPTIONAL group nested_struct {
- OPTIONAL INT32 A;
- OPTIONAL group b (LIST) {
- REPEATED group list {
- REQUIRED FIXED_LEN_BYTE_ARRAY (16) element;
- }
- }
- }
- }
- ";
- let expected_schema = parse_message_type(message_type).unwrap();
- let thrift_schema = to_thrift(&expected_schema).unwrap();
- let result_schema = from_thrift(&thrift_schema).unwrap();
- assert_eq!(result_schema, Arc::new(expected_schema));
- }
-
- #[test]
- fn test_schema_type_thrift_conversion_decimal() {
- let message_type = "
- message decimals {
- OPTIONAL INT32 field0;
- OPTIONAL INT64 field1 (DECIMAL (18, 2));
- OPTIONAL FIXED_LEN_BYTE_ARRAY (16) field2 (DECIMAL (38, 18));
- OPTIONAL BYTE_ARRAY field3 (DECIMAL (9));
- }
- ";
- let expected_schema = parse_message_type(message_type).unwrap();
- let thrift_schema = to_thrift(&expected_schema).unwrap();
- let result_schema = from_thrift(&thrift_schema).unwrap();
- assert_eq!(result_schema, Arc::new(expected_schema));
- }
-
- // Tests schema conversion from thrift, when num_children is set to Some(0) for a
- // primitive type.
- #[test]
- fn test_schema_from_thrift_with_num_children_set() {
- // schema definition written by parquet-cpp version 1.3.2-SNAPSHOT
- let message_type = "
- message schema {
- OPTIONAL BYTE_ARRAY id (UTF8);
- OPTIONAL BYTE_ARRAY name (UTF8);
- OPTIONAL BYTE_ARRAY message (UTF8);
- OPTIONAL INT32 type (UINT_8);
- OPTIONAL INT64 author_time (TIMESTAMP_MILLIS);
- OPTIONAL INT64 __index_level_0__;
- }
- ";
-
- let expected_schema = parse_message_type(message_type).unwrap();
- let mut thrift_schema = to_thrift(&expected_schema).unwrap();
- // Change all of None to Some(0)
- for mut elem in &mut thrift_schema[..] {
- if elem.num_children == None {
- elem.num_children = Some(0);
- }
- }
-
- let result_schema = from_thrift(&thrift_schema).unwrap();
- assert_eq!(result_schema, Arc::new(expected_schema));
- }
-
- // Sometimes parquet-cpp sets repetition level for the root node, which is against
- // the format definition, but we need to handle it by setting it back to None.
- #[test]
- fn test_schema_from_thrift_root_has_repetition() {
- // schema definition written by parquet-cpp version 1.3.2-SNAPSHOT
- let message_type = "
- message schema {
- OPTIONAL BYTE_ARRAY a (UTF8);
- OPTIONAL INT32 b (UINT_8);
- }
- ";
-
- let expected_schema = parse_message_type(message_type).unwrap();
- let mut thrift_schema = to_thrift(&expected_schema).unwrap();
- thrift_schema[0].repetition_type = Some(Repetition::REQUIRED.into());
-
- let result_schema = from_thrift(&thrift_schema).unwrap();
- assert_eq!(result_schema, Arc::new(expected_schema));
- }
-}
diff --git a/parquet/src/schema/visitor.rs b/parquet/src/schema/visitor.rs
deleted file mode 100644
index 61bc3be..0000000
--- a/parquet/src/schema/visitor.rs
+++ /dev/null
@@ -1,240 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-use crate::basic::{ConvertedType, Repetition};
-use crate::errors::ParquetError::General;
-use crate::errors::Result;
-use crate::schema::types::{Type, TypePtr};
-
-/// A utility trait to help user to traverse against parquet type.
-pub trait TypeVisitor<R, C> {
- /// Called when a primitive type hit.
- fn visit_primitive(&mut self, primitive_type: TypePtr, context: C) -> Result<R>;
-
- /// Default implementation when visiting a list.
- ///
- /// It checks list type definition and calls `visit_list_with_item` with extracted
- /// item type.
- ///
- /// To fully understand this algorithm, please refer to
- /// [parquet doc](https://github.com/apache/parquet-format/blob/master/LogicalTypes.md).
- fn visit_list(&mut self, list_type: TypePtr, context: C) -> Result<R> {
- match list_type.as_ref() {
- Type::PrimitiveType { .. } => panic!(
- "{:?} is a list type and can't be processed as primitive.",
- list_type
- ),
- Type::GroupType {
- basic_info: _,
- fields,
- } if fields.len() == 1 => {
- let list_item = fields.first().unwrap();
-
- match list_item.as_ref() {
- Type::PrimitiveType { .. } => {
- if list_item.get_basic_info().repetition() == Repetition::REPEATED
- {
- self.visit_list_with_item(
- list_type.clone(),
- list_item.clone(),
- context,
- )
- } else {
- Err(General(
- "Primitive element type of list must be repeated."
- .to_string(),
- ))
- }
- }
- Type::GroupType {
- basic_info: _,
- fields,
- } => {
- if fields.len() == 1
- && list_item.name() != "array"
- && list_item.name() != format!("{}_tuple", list_type.name())
- {
- self.visit_list_with_item(
- list_type.clone(),
- fields.first().unwrap().clone(),
- context,
- )
- } else {
- self.visit_list_with_item(
- list_type.clone(),
- list_item.clone(),
- context,
- )
- }
- }
- }
- }
- _ => Err(General(
- "Group element type of list can only contain one field.".to_string(),
- )),
- }
- }
-
- /// Called when a struct type hit.
- fn visit_struct(&mut self, struct_type: TypePtr, context: C) -> Result<R>;
-
- /// Called when a map type hit.
- fn visit_map(&mut self, map_type: TypePtr, context: C) -> Result<R>;
-
- /// A utility method which detects input type and calls corresponding method.
- fn dispatch(&mut self, cur_type: TypePtr, context: C) -> Result<R> {
- if cur_type.is_primitive() {
- self.visit_primitive(cur_type, context)
- } else {
- match cur_type.get_basic_info().converted_type() {
- ConvertedType::LIST => self.visit_list(cur_type, context),
- ConvertedType::MAP | ConvertedType::MAP_KEY_VALUE => {
- self.visit_map(cur_type, context)
- }
- _ => self.visit_struct(cur_type, context),
- }
- }
- }
-
- /// Called by `visit_list`.
- fn visit_list_with_item(
- &mut self,
- list_type: TypePtr,
- item_type: TypePtr,
- context: C,
- ) -> Result<R>;
-}
-
-#[cfg(test)]
-mod tests {
- use super::TypeVisitor;
- use crate::basic::Type as PhysicalType;
- use crate::errors::Result;
- use crate::schema::parser::parse_message_type;
- use crate::schema::types::TypePtr;
- use std::sync::Arc;
-
- struct TestVisitorContext {}
- struct TestVisitor {
- primitive_visited: bool,
- struct_visited: bool,
- list_visited: bool,
- root_type: TypePtr,
- }
-
- impl TypeVisitor<bool, TestVisitorContext> for TestVisitor {
- fn visit_primitive(
- &mut self,
- primitive_type: TypePtr,
- _context: TestVisitorContext,
- ) -> Result<bool> {
- assert_eq!(
- self.get_field_by_name(primitive_type.name()).as_ref(),
- primitive_type.as_ref()
- );
- self.primitive_visited = true;
- Ok(true)
- }
-
- fn visit_struct(
- &mut self,
- struct_type: TypePtr,
- _context: TestVisitorContext,
- ) -> Result<bool> {
- assert_eq!(
- self.get_field_by_name(struct_type.name()).as_ref(),
- struct_type.as_ref()
- );
- self.struct_visited = true;
- Ok(true)
- }
-
- fn visit_map(
- &mut self,
- _map_type: TypePtr,
- _context: TestVisitorContext,
- ) -> Result<bool> {
- unimplemented!()
- }
-
- fn visit_list_with_item(
- &mut self,
- list_type: TypePtr,
- item_type: TypePtr,
- _context: TestVisitorContext,
- ) -> Result<bool> {
- assert_eq!(
- self.get_field_by_name(list_type.name()).as_ref(),
- list_type.as_ref()
- );
- assert_eq!("element", item_type.name());
- assert_eq!(PhysicalType::INT32, item_type.get_physical_type());
- self.list_visited = true;
- Ok(true)
- }
- }
-
- impl TestVisitor {
- fn new(root: TypePtr) -> Self {
- Self {
- primitive_visited: false,
- struct_visited: false,
- list_visited: false,
- root_type: root,
- }
- }
-
- fn get_field_by_name(&self, name: &str) -> TypePtr {
- self.root_type
- .get_fields()
- .iter()
- .find(|t| t.name() == name)
- .cloned()
- .unwrap()
- }
- }
-
- #[test]
- fn test_visitor() {
- let message_type = "
- message spark_schema {
- REQUIRED INT32 a;
- OPTIONAL group inner_schema {
- REQUIRED INT32 b;
- REQUIRED DOUBLE c;
- }
-
- OPTIONAL group e (LIST) {
- REPEATED group list {
- REQUIRED INT32 element;
- }
- }
- ";
-
- let parquet_type = Arc::new(parse_message_type(&message_type).unwrap());
-
- let mut visitor = TestVisitor::new(parquet_type.clone());
- for f in parquet_type.get_fields() {
- let c = TestVisitorContext {};
- assert!(visitor.dispatch(f.clone(), c).unwrap());
- }
-
- assert!(visitor.struct_visited);
- assert!(visitor.primitive_visited);
- assert!(visitor.list_visited);
- }
-}
diff --git a/parquet/src/util/bit_packing.rs b/parquet/src/util/bit_packing.rs
deleted file mode 100644
index 6b9673f..0000000
--- a/parquet/src/util/bit_packing.rs
+++ /dev/null
@@ -1,3662 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-/// Unpack 32 values with bit width `num_bits` from `in_ptr`, and write to `out_ptr`.
-/// Return the `in_ptr` where the starting offset points to the first byte after all the
-/// bytes that were consumed.
-// TODO: may be better to make these more compact using if-else conditions.
-// However, this may require const generics:
-// https://github.com/rust-lang/rust/issues/44580
-// to eliminate the branching cost.
-// TODO: we should use SIMD instructions to further optimize this. I have explored
-// https://github.com/tantivy-search/bitpacking
-// but the layout it uses for SIMD is different from Parquet.
-// TODO: support packing as well, which is used for encoding.
-pub unsafe fn unpack32(
- mut in_ptr: *const u32,
- out_ptr: *mut u32,
- num_bits: usize,
-) -> *const u32 {
- in_ptr = match num_bits {
- 0 => nullunpacker32(in_ptr, out_ptr),
- 1 => unpack1_32(in_ptr, out_ptr),
- 2 => unpack2_32(in_ptr, out_ptr),
- 3 => unpack3_32(in_ptr, out_ptr),
- 4 => unpack4_32(in_ptr, out_ptr),
- 5 => unpack5_32(in_ptr, out_ptr),
- 6 => unpack6_32(in_ptr, out_ptr),
- 7 => unpack7_32(in_ptr, out_ptr),
- 8 => unpack8_32(in_ptr, out_ptr),
- 9 => unpack9_32(in_ptr, out_ptr),
- 10 => unpack10_32(in_ptr, out_ptr),
- 11 => unpack11_32(in_ptr, out_ptr),
- 12 => unpack12_32(in_ptr, out_ptr),
- 13 => unpack13_32(in_ptr, out_ptr),
- 14 => unpack14_32(in_ptr, out_ptr),
- 15 => unpack15_32(in_ptr, out_ptr),
- 16 => unpack16_32(in_ptr, out_ptr),
- 17 => unpack17_32(in_ptr, out_ptr),
- 18 => unpack18_32(in_ptr, out_ptr),
- 19 => unpack19_32(in_ptr, out_ptr),
- 20 => unpack20_32(in_ptr, out_ptr),
- 21 => unpack21_32(in_ptr, out_ptr),
- 22 => unpack22_32(in_ptr, out_ptr),
- 23 => unpack23_32(in_ptr, out_ptr),
- 24 => unpack24_32(in_ptr, out_ptr),
- 25 => unpack25_32(in_ptr, out_ptr),
- 26 => unpack26_32(in_ptr, out_ptr),
- 27 => unpack27_32(in_ptr, out_ptr),
- 28 => unpack28_32(in_ptr, out_ptr),
- 29 => unpack29_32(in_ptr, out_ptr),
- 30 => unpack30_32(in_ptr, out_ptr),
- 31 => unpack31_32(in_ptr, out_ptr),
- 32 => unpack32_32(in_ptr, out_ptr),
- _ => unimplemented!(),
- };
- in_ptr
-}
-
-unsafe fn nullunpacker32(in_buf: *const u32, mut out: *mut u32) -> *const u32 {
- for _ in 0..32 {
- *out = 0;
- out = out.offset(1);
- }
- in_buf
-}
-
-unsafe fn unpack1_32(in_buf: *const u32, mut out: *mut u32) -> *const u32 {
- *out = (*in_buf) & 1;
- out = out.offset(1);
- *out = ((*in_buf) >> 1) & 1;
- out = out.offset(1);
- *out = ((*in_buf) >> 2) & 1;
- out = out.offset(1);
- *out = ((*in_buf) >> 3) & 1;
- out = out.offset(1);
- *out = ((*in_buf) >> 4) & 1;
- out = out.offset(1);
- *out = ((*in_buf) >> 5) & 1;
- out = out.offset(1);
- *out = ((*in_buf) >> 6) & 1;
- out = out.offset(1);
- *out = ((*in_buf) >> 7) & 1;
- out = out.offset(1);
- *out = ((*in_buf) >> 8) & 1;
- out = out.offset(1);
- *out = ((*in_buf) >> 9) & 1;
- out = out.offset(1);
- *out = ((*in_buf) >> 10) & 1;
- out = out.offset(1);
- *out = ((*in_buf) >> 11) & 1;
- out = out.offset(1);
- *out = ((*in_buf) >> 12) & 1;
- out = out.offset(1);
- *out = ((*in_buf) >> 13) & 1;
- out = out.offset(1);
- *out = ((*in_buf) >> 14) & 1;
- out = out.offset(1);
- *out = ((*in_buf) >> 15) & 1;
- out = out.offset(1);
- *out = ((*in_buf) >> 16) & 1;
- out = out.offset(1);
- *out = ((*in_buf) >> 17) & 1;
- out = out.offset(1);
- *out = ((*in_buf) >> 18) & 1;
- out = out.offset(1);
- *out = ((*in_buf) >> 19) & 1;
- out = out.offset(1);
- *out = ((*in_buf) >> 20) & 1;
- out = out.offset(1);
- *out = ((*in_buf) >> 21) & 1;
- out = out.offset(1);
- *out = ((*in_buf) >> 22) & 1;
- out = out.offset(1);
- *out = ((*in_buf) >> 23) & 1;
- out = out.offset(1);
- *out = ((*in_buf) >> 24) & 1;
- out = out.offset(1);
- *out = ((*in_buf) >> 25) & 1;
- out = out.offset(1);
- *out = ((*in_buf) >> 26) & 1;
- out = out.offset(1);
- *out = ((*in_buf) >> 27) & 1;
- out = out.offset(1);
- *out = ((*in_buf) >> 28) & 1;
- out = out.offset(1);
- *out = ((*in_buf) >> 29) & 1;
- out = out.offset(1);
- *out = ((*in_buf) >> 30) & 1;
- out = out.offset(1);
- *out = (*in_buf) >> 31;
-
- in_buf.offset(1)
-}
-
-unsafe fn unpack2_32(mut in_buf: *const u32, mut out: *mut u32) -> *const u32 {
- *out = (*in_buf) % (1u32 << 2);
- out = out.offset(1);
- *out = ((*in_buf) >> 2) % (1u32 << 2);
- out = out.offset(1);
- *out = ((*in_buf) >> 4) % (1u32 << 2);
- out = out.offset(1);
- *out = ((*in_buf) >> 6) % (1u32 << 2);
- out = out.offset(1);
- *out = ((*in_buf) >> 8) % (1u32 << 2);
- out = out.offset(1);
- *out = ((*in_buf) >> 10) % (1u32 << 2);
- out = out.offset(1);
- *out = ((*in_buf) >> 12) % (1u32 << 2);
- out = out.offset(1);
- *out = ((*in_buf) >> 14) % (1u32 << 2);
- out = out.offset(1);
- *out = ((*in_buf) >> 16) % (1u32 << 2);
- out = out.offset(1);
- *out = ((*in_buf) >> 18) % (1u32 << 2);
- out = out.offset(1);
- *out = ((*in_buf) >> 20) % (1u32 << 2);
- out = out.offset(1);
- *out = ((*in_buf) >> 22) % (1u32 << 2);
- out = out.offset(1);
- *out = ((*in_buf) >> 24) % (1u32 << 2);
- out = out.offset(1);
- *out = ((*in_buf) >> 26) % (1u32 << 2);
- out = out.offset(1);
- *out = ((*in_buf) >> 28) % (1u32 << 2);
- out = out.offset(1);
- *out = (*in_buf) >> 30;
- out = out.offset(1);
- in_buf = in_buf.offset(1);
- *out = (*in_buf) % (1u32 << 2);
- out = out.offset(1);
- *out = ((*in_buf) >> 2) % (1u32 << 2);
- out = out.offset(1);
- *out = ((*in_buf) >> 4) % (1u32 << 2);
- out = out.offset(1);
- *out = ((*in_buf) >> 6) % (1u32 << 2);
- out = out.offset(1);
- *out = ((*in_buf) >> 8) % (1u32 << 2);
- out = out.offset(1);
- *out = ((*in_buf) >> 10) % (1u32 << 2);
- out = out.offset(1);
- *out = ((*in_buf) >> 12) % (1u32 << 2);
- out = out.offset(1);
- *out = ((*in_buf) >> 14) % (1u32 << 2);
- out = out.offset(1);
- *out = ((*in_buf) >> 16) % (1u32 << 2);
- out = out.offset(1);
- *out = ((*in_buf) >> 18) % (1u32 << 2);
- out = out.offset(1);
- *out = ((*in_buf) >> 20) % (1u32 << 2);
- out = out.offset(1);
- *out = ((*in_buf) >> 22) % (1u32 << 2);
- out = out.offset(1);
- *out = ((*in_buf) >> 24) % (1u32 << 2);
- out = out.offset(1);
- *out = ((*in_buf) >> 26) % (1u32 << 2);
- out = out.offset(1);
- *out = ((*in_buf) >> 28) % (1u32 << 2);
- out = out.offset(1);
- *out = (*in_buf) >> 30;
-
- in_buf.offset(1)
-}
-
-unsafe fn unpack3_32(mut in_buf: *const u32, mut out: *mut u32) -> *const u32 {
- *out = (*in_buf) % (1u32 << 3);
- out = out.offset(1);
- *out = ((*in_buf) >> 3) % (1u32 << 3);
- out = out.offset(1);
- *out = ((*in_buf) >> 6) % (1u32 << 3);
- out = out.offset(1);
- *out = ((*in_buf) >> 9) % (1u32 << 3);
- out = out.offset(1);
- *out = ((*in_buf) >> 12) % (1u32 << 3);
- out = out.offset(1);
- *out = ((*in_buf) >> 15) % (1u32 << 3);
- out = out.offset(1);
- *out = ((*in_buf) >> 18) % (1u32 << 3);
- out = out.offset(1);
- *out = ((*in_buf) >> 21) % (1u32 << 3);
- out = out.offset(1);
- *out = ((*in_buf) >> 24) % (1u32 << 3);
- out = out.offset(1);
- *out = ((*in_buf) >> 27) % (1u32 << 3);
- out = out.offset(1);
- *out = (*in_buf) >> 30;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 1)) << (3 - 1);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 1) % (1u32 << 3);
- out = out.offset(1);
- *out = ((*in_buf) >> 4) % (1u32 << 3);
- out = out.offset(1);
- *out = ((*in_buf) >> 7) % (1u32 << 3);
- out = out.offset(1);
- *out = ((*in_buf) >> 10) % (1u32 << 3);
- out = out.offset(1);
- *out = ((*in_buf) >> 13) % (1u32 << 3);
- out = out.offset(1);
- *out = ((*in_buf) >> 16) % (1u32 << 3);
- out = out.offset(1);
- *out = ((*in_buf) >> 19) % (1u32 << 3);
- out = out.offset(1);
- *out = ((*in_buf) >> 22) % (1u32 << 3);
- out = out.offset(1);
- *out = ((*in_buf) >> 25) % (1u32 << 3);
- out = out.offset(1);
- *out = ((*in_buf) >> 28) % (1u32 << 3);
- out = out.offset(1);
- *out = (*in_buf) >> 31;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 2)) << (3 - 2);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 2) % (1u32 << 3);
- out = out.offset(1);
- *out = ((*in_buf) >> 5) % (1u32 << 3);
- out = out.offset(1);
- *out = ((*in_buf) >> 8) % (1u32 << 3);
- out = out.offset(1);
- *out = ((*in_buf) >> 11) % (1u32 << 3);
- out = out.offset(1);
- *out = ((*in_buf) >> 14) % (1u32 << 3);
- out = out.offset(1);
- *out = ((*in_buf) >> 17) % (1u32 << 3);
- out = out.offset(1);
- *out = ((*in_buf) >> 20) % (1u32 << 3);
- out = out.offset(1);
- *out = ((*in_buf) >> 23) % (1u32 << 3);
- out = out.offset(1);
- *out = ((*in_buf) >> 26) % (1u32 << 3);
- out = out.offset(1);
- *out = (*in_buf) >> 29;
-
- in_buf.offset(1)
-}
-
-unsafe fn unpack4_32(mut in_buf: *const u32, mut out: *mut u32) -> *const u32 {
- *out = (*in_buf) % (1u32 << 4);
- out = out.offset(1);
- *out = ((*in_buf) >> 4) % (1u32 << 4);
- out = out.offset(1);
- *out = ((*in_buf) >> 8) % (1u32 << 4);
- out = out.offset(1);
- *out = ((*in_buf) >> 12) % (1u32 << 4);
- out = out.offset(1);
- *out = ((*in_buf) >> 16) % (1u32 << 4);
- out = out.offset(1);
- *out = ((*in_buf) >> 20) % (1u32 << 4);
- out = out.offset(1);
- *out = ((*in_buf) >> 24) % (1u32 << 4);
- out = out.offset(1);
- *out = ((*in_buf) >> 28) % (1u32 << 4);
- out = out.offset(1);
- in_buf = in_buf.offset(1);
-
- *out = (*in_buf) % (1u32 << 4);
- out = out.offset(1);
- *out = ((*in_buf) >> 4) % (1u32 << 4);
- out = out.offset(1);
- *out = ((*in_buf) >> 8) % (1u32 << 4);
- out = out.offset(1);
- *out = ((*in_buf) >> 12) % (1u32 << 4);
- out = out.offset(1);
- *out = ((*in_buf) >> 16) % (1u32 << 4);
- out = out.offset(1);
- *out = ((*in_buf) >> 20) % (1u32 << 4);
- out = out.offset(1);
- *out = ((*in_buf) >> 24) % (1u32 << 4);
- out = out.offset(1);
- *out = ((*in_buf) >> 28) % (1u32 << 4);
- out = out.offset(1);
- in_buf = in_buf.offset(1);
-
- *out = (*in_buf) % (1u32 << 4);
- out = out.offset(1);
- *out = ((*in_buf) >> 4) % (1u32 << 4);
- out = out.offset(1);
- *out = ((*in_buf) >> 8) % (1u32 << 4);
- out = out.offset(1);
- *out = ((*in_buf) >> 12) % (1u32 << 4);
- out = out.offset(1);
- *out = ((*in_buf) >> 16) % (1u32 << 4);
- out = out.offset(1);
- *out = ((*in_buf) >> 20) % (1u32 << 4);
- out = out.offset(1);
- *out = ((*in_buf) >> 24) % (1u32 << 4);
- out = out.offset(1);
- *out = ((*in_buf) >> 28) % (1u32 << 4);
- out = out.offset(1);
- in_buf = in_buf.offset(1);
-
- *out = (*in_buf) % (1u32 << 4);
- out = out.offset(1);
- *out = ((*in_buf) >> 4) % (1u32 << 4);
- out = out.offset(1);
- *out = ((*in_buf) >> 8) % (1u32 << 4);
- out = out.offset(1);
- *out = ((*in_buf) >> 12) % (1u32 << 4);
- out = out.offset(1);
- *out = ((*in_buf) >> 16) % (1u32 << 4);
- out = out.offset(1);
- *out = ((*in_buf) >> 20) % (1u32 << 4);
- out = out.offset(1);
- *out = ((*in_buf) >> 24) % (1u32 << 4);
- out = out.offset(1);
- *out = ((*in_buf) >> 28) % (1u32 << 4);
-
- in_buf.offset(1)
-}
-
-unsafe fn unpack5_32(mut in_buf: *const u32, mut out: *mut u32) -> *const u32 {
- *out = (*in_buf) % (1u32 << 5);
- out = out.offset(1);
- *out = ((*in_buf) >> 5) % (1u32 << 5);
- out = out.offset(1);
- *out = ((*in_buf) >> 10) % (1u32 << 5);
- out = out.offset(1);
- *out = ((*in_buf) >> 15) % (1u32 << 5);
- out = out.offset(1);
- *out = ((*in_buf) >> 20) % (1u32 << 5);
- out = out.offset(1);
- *out = ((*in_buf) >> 25) % (1u32 << 5);
- out = out.offset(1);
- *out = (*in_buf) >> 30;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 3)) << (5 - 3);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 3) % (1u32 << 5);
- out = out.offset(1);
- *out = ((*in_buf) >> 8) % (1u32 << 5);
- out = out.offset(1);
- *out = ((*in_buf) >> 13) % (1u32 << 5);
- out = out.offset(1);
- *out = ((*in_buf) >> 18) % (1u32 << 5);
- out = out.offset(1);
- *out = ((*in_buf) >> 23) % (1u32 << 5);
- out = out.offset(1);
- *out = ((*in_buf) >> 28) % (1u32 << 5);
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 1)) << (5 - 1);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 1) % (1u32 << 5);
- out = out.offset(1);
- *out = ((*in_buf) >> 6) % (1u32 << 5);
- out = out.offset(1);
- *out = ((*in_buf) >> 11) % (1u32 << 5);
- out = out.offset(1);
- *out = ((*in_buf) >> 16) % (1u32 << 5);
- out = out.offset(1);
- *out = ((*in_buf) >> 21) % (1u32 << 5);
- out = out.offset(1);
- *out = ((*in_buf) >> 26) % (1u32 << 5);
- out = out.offset(1);
- *out = (*in_buf) >> 31;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 4)) << (5 - 4);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 4) % (1u32 << 5);
- out = out.offset(1);
- *out = ((*in_buf) >> 9) % (1u32 << 5);
- out = out.offset(1);
- *out = ((*in_buf) >> 14) % (1u32 << 5);
- out = out.offset(1);
- *out = ((*in_buf) >> 19) % (1u32 << 5);
- out = out.offset(1);
- *out = ((*in_buf) >> 24) % (1u32 << 5);
- out = out.offset(1);
- *out = (*in_buf) >> 29;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 2)) << (5 - 2);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 2) % (1u32 << 5);
- out = out.offset(1);
- *out = ((*in_buf) >> 7) % (1u32 << 5);
- out = out.offset(1);
- *out = ((*in_buf) >> 12) % (1u32 << 5);
- out = out.offset(1);
- *out = ((*in_buf) >> 17) % (1u32 << 5);
- out = out.offset(1);
- *out = ((*in_buf) >> 22) % (1u32 << 5);
- out = out.offset(1);
- *out = (*in_buf) >> 27;
-
- in_buf.offset(1)
-}
-
-unsafe fn unpack6_32(mut in_buf: *const u32, mut out: *mut u32) -> *const u32 {
- *out = (*in_buf) % (1u32 << 6);
- out = out.offset(1);
- *out = ((*in_buf) >> 6) % (1u32 << 6);
- out = out.offset(1);
- *out = ((*in_buf) >> 12) % (1u32 << 6);
- out = out.offset(1);
- *out = ((*in_buf) >> 18) % (1u32 << 6);
- out = out.offset(1);
- *out = ((*in_buf) >> 24) % (1u32 << 6);
- out = out.offset(1);
- *out = (*in_buf) >> 30;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 4)) << (6 - 4);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 4) % (1u32 << 6);
- out = out.offset(1);
- *out = ((*in_buf) >> 10) % (1u32 << 6);
- out = out.offset(1);
- *out = ((*in_buf) >> 16) % (1u32 << 6);
- out = out.offset(1);
- *out = ((*in_buf) >> 22) % (1u32 << 6);
- out = out.offset(1);
- *out = (*in_buf) >> 28;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 2)) << (6 - 2);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 2) % (1u32 << 6);
- out = out.offset(1);
- *out = ((*in_buf) >> 8) % (1u32 << 6);
- out = out.offset(1);
- *out = ((*in_buf) >> 14) % (1u32 << 6);
- out = out.offset(1);
- *out = ((*in_buf) >> 20) % (1u32 << 6);
- out = out.offset(1);
- *out = (*in_buf) >> 26;
- in_buf = in_buf.offset(1);
- out = out.offset(1);
-
- *out = (*in_buf) % (1u32 << 6);
- out = out.offset(1);
- *out = ((*in_buf) >> 6) % (1u32 << 6);
- out = out.offset(1);
- *out = ((*in_buf) >> 12) % (1u32 << 6);
- out = out.offset(1);
- *out = ((*in_buf) >> 18) % (1u32 << 6);
- out = out.offset(1);
- *out = ((*in_buf) >> 24) % (1u32 << 6);
- out = out.offset(1);
- *out = (*in_buf) >> 30;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 4)) << (6 - 4);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 4) % (1u32 << 6);
- out = out.offset(1);
- *out = ((*in_buf) >> 10) % (1u32 << 6);
- out = out.offset(1);
- *out = ((*in_buf) >> 16) % (1u32 << 6);
- out = out.offset(1);
- *out = ((*in_buf) >> 22) % (1u32 << 6);
- out = out.offset(1);
- *out = (*in_buf) >> 28;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 2)) << (6 - 2);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 2) % (1u32 << 6);
- out = out.offset(1);
- *out = ((*in_buf) >> 8) % (1u32 << 6);
- out = out.offset(1);
- *out = ((*in_buf) >> 14) % (1u32 << 6);
- out = out.offset(1);
- *out = ((*in_buf) >> 20) % (1u32 << 6);
- out = out.offset(1);
- *out = (*in_buf) >> 26;
-
- in_buf.offset(1)
-}
-
-unsafe fn unpack7_32(mut in_buf: *const u32, mut out: *mut u32) -> *const u32 {
- *out = (*in_buf) % (1u32 << 7);
- out = out.offset(1);
- *out = ((*in_buf) >> 7) % (1u32 << 7);
- out = out.offset(1);
- *out = ((*in_buf) >> 14) % (1u32 << 7);
- out = out.offset(1);
- *out = ((*in_buf) >> 21) % (1u32 << 7);
- out = out.offset(1);
- *out = (*in_buf) >> 28;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 3)) << (7 - 3);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 3) % (1u32 << 7);
- out = out.offset(1);
- *out = ((*in_buf) >> 10) % (1u32 << 7);
- out = out.offset(1);
- *out = ((*in_buf) >> 17) % (1u32 << 7);
- out = out.offset(1);
- *out = ((*in_buf) >> 24) % (1u32 << 7);
- out = out.offset(1);
- *out = (*in_buf) >> 31;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 6)) << (7 - 6);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 6) % (1u32 << 7);
- out = out.offset(1);
- *out = ((*in_buf) >> 13) % (1u32 << 7);
- out = out.offset(1);
- *out = ((*in_buf) >> 20) % (1u32 << 7);
- out = out.offset(1);
- *out = (*in_buf) >> 27;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 2)) << (7 - 2);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 2) % (1u32 << 7);
- out = out.offset(1);
- *out = ((*in_buf) >> 9) % (1u32 << 7);
- out = out.offset(1);
- *out = ((*in_buf) >> 16) % (1u32 << 7);
- out = out.offset(1);
- *out = ((*in_buf) >> 23) % (1u32 << 7);
- out = out.offset(1);
- *out = (*in_buf) >> 30;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 5)) << (7 - 5);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 5) % (1u32 << 7);
- out = out.offset(1);
- *out = ((*in_buf) >> 12) % (1u32 << 7);
- out = out.offset(1);
- *out = ((*in_buf) >> 19) % (1u32 << 7);
- out = out.offset(1);
- *out = (*in_buf) >> 26;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 1)) << (7 - 1);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 1) % (1u32 << 7);
- out = out.offset(1);
- *out = ((*in_buf) >> 8) % (1u32 << 7);
- out = out.offset(1);
- *out = ((*in_buf) >> 15) % (1u32 << 7);
- out = out.offset(1);
- *out = ((*in_buf) >> 22) % (1u32 << 7);
- out = out.offset(1);
- *out = (*in_buf) >> 29;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 4)) << (7 - 4);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 4) % (1u32 << 7);
- out = out.offset(1);
- *out = ((*in_buf) >> 11) % (1u32 << 7);
- out = out.offset(1);
- *out = ((*in_buf) >> 18) % (1u32 << 7);
- out = out.offset(1);
- *out = (*in_buf) >> 25;
-
- in_buf.offset(1)
-}
-
-unsafe fn unpack8_32(mut in_buf: *const u32, mut out: *mut u32) -> *const u32 {
- *out = (*in_buf) % (1u32 << 8);
- out = out.offset(1);
- *out = ((*in_buf) >> 8) % (1u32 << 8);
- out = out.offset(1);
- *out = ((*in_buf) >> 16) % (1u32 << 8);
- out = out.offset(1);
- *out = (*in_buf) >> 24;
- out = out.offset(1);
- in_buf = in_buf.offset(1);
-
- *out = (*in_buf) % (1u32 << 8);
- out = out.offset(1);
- *out = ((*in_buf) >> 8) % (1u32 << 8);
- out = out.offset(1);
- *out = ((*in_buf) >> 16) % (1u32 << 8);
- out = out.offset(1);
- *out = (*in_buf) >> 24;
- out = out.offset(1);
- in_buf = in_buf.offset(1);
-
- *out = (*in_buf) % (1u32 << 8);
- out = out.offset(1);
- *out = ((*in_buf) >> 8) % (1u32 << 8);
- out = out.offset(1);
- *out = ((*in_buf) >> 16) % (1u32 << 8);
- out = out.offset(1);
- *out = (*in_buf) >> 24;
- out = out.offset(1);
- in_buf = in_buf.offset(1);
-
- *out = (*in_buf) % (1u32 << 8);
- out = out.offset(1);
- *out = ((*in_buf) >> 8) % (1u32 << 8);
- out = out.offset(1);
- *out = ((*in_buf) >> 16) % (1u32 << 8);
- out = out.offset(1);
- *out = (*in_buf) >> 24;
- out = out.offset(1);
- in_buf = in_buf.offset(1);
-
- *out = (*in_buf) % (1u32 << 8);
- out = out.offset(1);
- *out = ((*in_buf) >> 8) % (1u32 << 8);
- out = out.offset(1);
- *out = ((*in_buf) >> 16) % (1u32 << 8);
- out = out.offset(1);
- *out = (*in_buf) >> 24;
- out = out.offset(1);
- in_buf = in_buf.offset(1);
-
- *out = (*in_buf) % (1u32 << 8);
- out = out.offset(1);
- *out = ((*in_buf) >> 8) % (1u32 << 8);
- out = out.offset(1);
- *out = ((*in_buf) >> 16) % (1u32 << 8);
- out = out.offset(1);
- *out = (*in_buf) >> 24;
- out = out.offset(1);
- in_buf = in_buf.offset(1);
-
- *out = (*in_buf) % (1u32 << 8);
- out = out.offset(1);
- *out = ((*in_buf) >> 8) % (1u32 << 8);
- out = out.offset(1);
- *out = ((*in_buf) >> 16) % (1u32 << 8);
- out = out.offset(1);
- *out = (*in_buf) >> 24;
- out = out.offset(1);
- in_buf = in_buf.offset(1);
-
- *out = (*in_buf) % (1u32 << 8);
- out = out.offset(1);
- *out = ((*in_buf) >> 8) % (1u32 << 8);
- out = out.offset(1);
- *out = ((*in_buf) >> 16) % (1u32 << 8);
- out = out.offset(1);
- *out = (*in_buf) >> 24;
-
- in_buf.offset(1)
-}
-
-unsafe fn unpack9_32(mut in_buf: *const u32, mut out: *mut u32) -> *const u32 {
- *out = (*in_buf) % (1u32 << 9);
- out = out.offset(1);
- *out = ((*in_buf) >> 9) % (1u32 << 9);
- out = out.offset(1);
- *out = ((*in_buf) >> 18) % (1u32 << 9);
- out = out.offset(1);
- *out = (*in_buf) >> 27;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 4)) << (9 - 4);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 4) % (1u32 << 9);
- out = out.offset(1);
- *out = ((*in_buf) >> 13) % (1u32 << 9);
- out = out.offset(1);
- *out = ((*in_buf) >> 22) % (1u32 << 9);
- out = out.offset(1);
- *out = (*in_buf) >> 31;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 8)) << (9 - 8);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 8) % (1u32 << 9);
- out = out.offset(1);
- *out = ((*in_buf) >> 17) % (1u32 << 9);
- out = out.offset(1);
- *out = (*in_buf) >> 26;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 3)) << (9 - 3);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 3) % (1u32 << 9);
- out = out.offset(1);
- *out = ((*in_buf) >> 12) % (1u32 << 9);
- out = out.offset(1);
- *out = ((*in_buf) >> 21) % (1u32 << 9);
- out = out.offset(1);
- *out = (*in_buf) >> 30;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 7)) << (9 - 7);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 7) % (1u32 << 9);
- out = out.offset(1);
- *out = ((*in_buf) >> 16) % (1u32 << 9);
- out = out.offset(1);
- *out = (*in_buf) >> 25;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 2)) << (9 - 2);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 2) % (1u32 << 9);
- out = out.offset(1);
- *out = ((*in_buf) >> 11) % (1u32 << 9);
- out = out.offset(1);
- *out = ((*in_buf) >> 20) % (1u32 << 9);
- out = out.offset(1);
- *out = (*in_buf) >> 29;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 6)) << (9 - 6);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 6) % (1u32 << 9);
- out = out.offset(1);
- *out = ((*in_buf) >> 15) % (1u32 << 9);
- out = out.offset(1);
- *out = (*in_buf) >> 24;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 1)) << (9 - 1);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 1) % (1u32 << 9);
- out = out.offset(1);
- *out = ((*in_buf) >> 10) % (1u32 << 9);
- out = out.offset(1);
- *out = ((*in_buf) >> 19) % (1u32 << 9);
- out = out.offset(1);
- *out = (*in_buf) >> 28;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 5)) << (9 - 5);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 5) % (1u32 << 9);
- out = out.offset(1);
- *out = ((*in_buf) >> 14) % (1u32 << 9);
- out = out.offset(1);
- *out = (*in_buf) >> 23;
-
- in_buf.offset(1)
-}
-
-unsafe fn unpack10_32(mut in_buf: *const u32, mut out: *mut u32) -> *const u32 {
- *out = (*in_buf) % (1u32 << 10);
- out = out.offset(1);
- *out = ((*in_buf) >> 10) % (1u32 << 10);
- out = out.offset(1);
- *out = ((*in_buf) >> 20) % (1u32 << 10);
- out = out.offset(1);
- *out = (*in_buf) >> 30;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 8)) << (10 - 8);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 8) % (1u32 << 10);
- out = out.offset(1);
- *out = ((*in_buf) >> 18) % (1u32 << 10);
- out = out.offset(1);
- *out = (*in_buf) >> 28;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 6)) << (10 - 6);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 6) % (1u32 << 10);
- out = out.offset(1);
- *out = ((*in_buf) >> 16) % (1u32 << 10);
- out = out.offset(1);
- *out = (*in_buf) >> 26;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 4)) << (10 - 4);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 4) % (1u32 << 10);
- out = out.offset(1);
- *out = ((*in_buf) >> 14) % (1u32 << 10);
- out = out.offset(1);
- *out = (*in_buf) >> 24;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 2)) << (10 - 2);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 2) % (1u32 << 10);
- out = out.offset(1);
- *out = ((*in_buf) >> 12) % (1u32 << 10);
- out = out.offset(1);
- *out = (*in_buf) >> 22;
- in_buf = in_buf.offset(1);
- out = out.offset(1);
-
- *out = (*in_buf) % (1u32 << 10);
- out = out.offset(1);
- *out = ((*in_buf) >> 10) % (1u32 << 10);
- out = out.offset(1);
- *out = ((*in_buf) >> 20) % (1u32 << 10);
- out = out.offset(1);
- *out = (*in_buf) >> 30;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 8)) << (10 - 8);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 8) % (1u32 << 10);
- out = out.offset(1);
- *out = ((*in_buf) >> 18) % (1u32 << 10);
- out = out.offset(1);
- *out = (*in_buf) >> 28;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 6)) << (10 - 6);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 6) % (1u32 << 10);
- out = out.offset(1);
- *out = ((*in_buf) >> 16) % (1u32 << 10);
- out = out.offset(1);
- *out = (*in_buf) >> 26;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 4)) << (10 - 4);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 4) % (1u32 << 10);
- out = out.offset(1);
- *out = ((*in_buf) >> 14) % (1u32 << 10);
- out = out.offset(1);
- *out = (*in_buf) >> 24;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 2)) << (10 - 2);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 2) % (1u32 << 10);
- out = out.offset(1);
- *out = ((*in_buf) >> 12) % (1u32 << 10);
- out = out.offset(1);
- *out = (*in_buf) >> 22;
-
- in_buf.offset(1)
-}
-
-unsafe fn unpack11_32(mut in_buf: *const u32, mut out: *mut u32) -> *const u32 {
- *out = (*in_buf) % (1u32 << 11);
- out = out.offset(1);
- *out = ((*in_buf) >> 11) % (1u32 << 11);
- out = out.offset(1);
- *out = (*in_buf) >> 22;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 1)) << (11 - 1);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 1) % (1u32 << 11);
- out = out.offset(1);
- *out = ((*in_buf) >> 12) % (1u32 << 11);
- out = out.offset(1);
- *out = (*in_buf) >> 23;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 2)) << (11 - 2);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 2) % (1u32 << 11);
- out = out.offset(1);
- *out = ((*in_buf) >> 13) % (1u32 << 11);
- out = out.offset(1);
- *out = (*in_buf) >> 24;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 3)) << (11 - 3);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 3) % (1u32 << 11);
- out = out.offset(1);
- *out = ((*in_buf) >> 14) % (1u32 << 11);
- out = out.offset(1);
- *out = (*in_buf) >> 25;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 4)) << (11 - 4);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 4) % (1u32 << 11);
- out = out.offset(1);
- *out = ((*in_buf) >> 15) % (1u32 << 11);
- out = out.offset(1);
- *out = (*in_buf) >> 26;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 5)) << (11 - 5);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 5) % (1u32 << 11);
- out = out.offset(1);
- *out = ((*in_buf) >> 16) % (1u32 << 11);
- out = out.offset(1);
- *out = (*in_buf) >> 27;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 6)) << (11 - 6);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 6) % (1u32 << 11);
- out = out.offset(1);
- *out = ((*in_buf) >> 17) % (1u32 << 11);
- out = out.offset(1);
- *out = (*in_buf) >> 28;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 7)) << (11 - 7);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 7) % (1u32 << 11);
- out = out.offset(1);
- *out = ((*in_buf) >> 18) % (1u32 << 11);
- out = out.offset(1);
- *out = (*in_buf) >> 29;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 8)) << (11 - 8);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 8) % (1u32 << 11);
- out = out.offset(1);
- *out = ((*in_buf) >> 19) % (1u32 << 11);
- out = out.offset(1);
- *out = (*in_buf) >> 30;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 9)) << (11 - 9);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 9) % (1u32 << 11);
- out = out.offset(1);
- *out = ((*in_buf) >> 20) % (1u32 << 11);
- out = out.offset(1);
- *out = (*in_buf) >> 31;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 10)) << (11 - 10);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 10) % (1u32 << 11);
- out = out.offset(1);
- *out = (*in_buf) >> 21;
-
- in_buf.offset(1)
-}
-
-unsafe fn unpack12_32(mut in_buf: *const u32, mut out: *mut u32) -> *const u32 {
- *out = (*in_buf) % (1u32 << 12);
- out = out.offset(1);
- *out = ((*in_buf) >> 12) % (1u32 << 12);
- out = out.offset(1);
- *out = (*in_buf) >> 24;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 4)) << (12 - 4);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 4) % (1u32 << 12);
- out = out.offset(1);
- *out = ((*in_buf) >> 16) % (1u32 << 12);
- out = out.offset(1);
- *out = (*in_buf) >> 28;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 8)) << (12 - 8);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 8) % (1u32 << 12);
- out = out.offset(1);
- *out = (*in_buf) >> 20;
- in_buf = in_buf.offset(1);
- out = out.offset(1);
-
- *out = (*in_buf) % (1u32 << 12);
- out = out.offset(1);
- *out = ((*in_buf) >> 12) % (1u32 << 12);
- out = out.offset(1);
- *out = (*in_buf) >> 24;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 4)) << (12 - 4);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 4) % (1u32 << 12);
- out = out.offset(1);
- *out = ((*in_buf) >> 16) % (1u32 << 12);
- out = out.offset(1);
- *out = (*in_buf) >> 28;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 8)) << (12 - 8);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 8) % (1u32 << 12);
- out = out.offset(1);
- *out = (*in_buf) >> 20;
- in_buf = in_buf.offset(1);
- out = out.offset(1);
-
- *out = (*in_buf) % (1u32 << 12);
- out = out.offset(1);
- *out = ((*in_buf) >> 12) % (1u32 << 12);
- out = out.offset(1);
- *out = (*in_buf) >> 24;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 4)) << (12 - 4);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 4) % (1u32 << 12);
- out = out.offset(1);
- *out = ((*in_buf) >> 16) % (1u32 << 12);
- out = out.offset(1);
- *out = (*in_buf) >> 28;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 8)) << (12 - 8);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 8) % (1u32 << 12);
- out = out.offset(1);
- *out = (*in_buf) >> 20;
- in_buf = in_buf.offset(1);
- out = out.offset(1);
-
- *out = (*in_buf) % (1u32 << 12);
- out = out.offset(1);
- *out = ((*in_buf) >> 12) % (1u32 << 12);
- out = out.offset(1);
- *out = (*in_buf) >> 24;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 4)) << (12 - 4);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 4) % (1u32 << 12);
- out = out.offset(1);
- *out = ((*in_buf) >> 16) % (1u32 << 12);
- out = out.offset(1);
- *out = (*in_buf) >> 28;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 8)) << (12 - 8);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 8) % (1u32 << 12);
- out = out.offset(1);
- *out = (*in_buf) >> 20;
-
- in_buf.offset(1)
-}
-
-unsafe fn unpack13_32(mut in_buf: *const u32, mut out: *mut u32) -> *const u32 {
- *out = (*in_buf) % (1u32 << 13);
- out = out.offset(1);
- *out = ((*in_buf) >> 13) % (1u32 << 13);
- out = out.offset(1);
- *out = (*in_buf) >> 26;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 7)) << (13 - 7);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 7) % (1u32 << 13);
- out = out.offset(1);
- *out = (*in_buf) >> 20;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 1)) << (13 - 1);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 1) % (1u32 << 13);
- out = out.offset(1);
- *out = ((*in_buf) >> 14) % (1u32 << 13);
- out = out.offset(1);
- *out = (*in_buf) >> 27;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 8)) << (13 - 8);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 8) % (1u32 << 13);
- out = out.offset(1);
- *out = (*in_buf) >> 21;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 2)) << (13 - 2);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 2) % (1u32 << 13);
- out = out.offset(1);
- *out = ((*in_buf) >> 15) % (1u32 << 13);
- out = out.offset(1);
- *out = (*in_buf) >> 28;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 9)) << (13 - 9);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 9) % (1u32 << 13);
- out = out.offset(1);
- *out = (*in_buf) >> 22;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 3)) << (13 - 3);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 3) % (1u32 << 13);
- out = out.offset(1);
- *out = ((*in_buf) >> 16) % (1u32 << 13);
- out = out.offset(1);
- *out = (*in_buf) >> 29;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 10)) << (13 - 10);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 10) % (1u32 << 13);
- out = out.offset(1);
- *out = (*in_buf) >> 23;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 4)) << (13 - 4);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 4) % (1u32 << 13);
- out = out.offset(1);
- *out = ((*in_buf) >> 17) % (1u32 << 13);
- out = out.offset(1);
- *out = (*in_buf) >> 30;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 11)) << (13 - 11);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 11) % (1u32 << 13);
- out = out.offset(1);
- *out = (*in_buf) >> 24;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 5)) << (13 - 5);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 5) % (1u32 << 13);
- out = out.offset(1);
- *out = ((*in_buf) >> 18) % (1u32 << 13);
- out = out.offset(1);
- *out = (*in_buf) >> 31;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 12)) << (13 - 12);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 12) % (1u32 << 13);
- out = out.offset(1);
- *out = (*in_buf) >> 25;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 6)) << (13 - 6);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 6) % (1u32 << 13);
- out = out.offset(1);
- *out = (*in_buf) >> 19;
-
- in_buf.offset(1)
-}
-
-unsafe fn unpack14_32(mut in_buf: *const u32, mut out: *mut u32) -> *const u32 {
- *out = (*in_buf) % (1u32 << 14);
- out = out.offset(1);
- *out = ((*in_buf) >> 14) % (1u32 << 14);
- out = out.offset(1);
- *out = (*in_buf) >> 28;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 10)) << (14 - 10);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 10) % (1u32 << 14);
- out = out.offset(1);
- *out = (*in_buf) >> 24;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 6)) << (14 - 6);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 6) % (1u32 << 14);
- out = out.offset(1);
- *out = (*in_buf) >> 20;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 2)) << (14 - 2);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 2) % (1u32 << 14);
- out = out.offset(1);
- *out = ((*in_buf) >> 16) % (1u32 << 14);
- out = out.offset(1);
- *out = (*in_buf) >> 30;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 12)) << (14 - 12);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 12) % (1u32 << 14);
- out = out.offset(1);
- *out = (*in_buf) >> 26;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 8)) << (14 - 8);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 8) % (1u32 << 14);
- out = out.offset(1);
- *out = (*in_buf) >> 22;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 4)) << (14 - 4);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 4) % (1u32 << 14);
- out = out.offset(1);
- *out = (*in_buf) >> 18;
- in_buf = in_buf.offset(1);
- out = out.offset(1);
-
- *out = (*in_buf) % (1u32 << 14);
- out = out.offset(1);
- *out = ((*in_buf) >> 14) % (1u32 << 14);
- out = out.offset(1);
- *out = (*in_buf) >> 28;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 10)) << (14 - 10);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 10) % (1u32 << 14);
- out = out.offset(1);
- *out = (*in_buf) >> 24;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 6)) << (14 - 6);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 6) % (1u32 << 14);
- out = out.offset(1);
- *out = (*in_buf) >> 20;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 2)) << (14 - 2);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 2) % (1u32 << 14);
- out = out.offset(1);
- *out = ((*in_buf) >> 16) % (1u32 << 14);
- out = out.offset(1);
- *out = (*in_buf) >> 30;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 12)) << (14 - 12);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 12) % (1u32 << 14);
- out = out.offset(1);
- *out = (*in_buf) >> 26;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 8)) << (14 - 8);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 8) % (1u32 << 14);
- out = out.offset(1);
- *out = (*in_buf) >> 22;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 4)) << (14 - 4);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 4) % (1u32 << 14);
- out = out.offset(1);
- *out = (*in_buf) >> 18;
-
- in_buf.offset(1)
-}
-
-unsafe fn unpack15_32(mut in_buf: *const u32, mut out: *mut u32) -> *const u32 {
- *out = (*in_buf) % (1u32 << 15);
- out = out.offset(1);
- *out = ((*in_buf) >> 15) % (1u32 << 15);
- out = out.offset(1);
- *out = (*in_buf) >> 30;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 13)) << (15 - 13);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 13) % (1u32 << 15);
- out = out.offset(1);
- *out = (*in_buf) >> 28;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 11)) << (15 - 11);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 11) % (1u32 << 15);
- out = out.offset(1);
- *out = (*in_buf) >> 26;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 9)) << (15 - 9);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 9) % (1u32 << 15);
- out = out.offset(1);
- *out = (*in_buf) >> 24;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 7)) << (15 - 7);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 7) % (1u32 << 15);
- out = out.offset(1);
- *out = (*in_buf) >> 22;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 5)) << (15 - 5);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 5) % (1u32 << 15);
- out = out.offset(1);
- *out = (*in_buf) >> 20;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 3)) << (15 - 3);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 3) % (1u32 << 15);
- out = out.offset(1);
- *out = (*in_buf) >> 18;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 1)) << (15 - 1);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 1) % (1u32 << 15);
- out = out.offset(1);
- *out = ((*in_buf) >> 16) % (1u32 << 15);
- out = out.offset(1);
- *out = (*in_buf) >> 31;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 14)) << (15 - 14);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 14) % (1u32 << 15);
- out = out.offset(1);
- *out = (*in_buf) >> 29;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 12)) << (15 - 12);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 12) % (1u32 << 15);
- out = out.offset(1);
- *out = (*in_buf) >> 27;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 10)) << (15 - 10);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 10) % (1u32 << 15);
- out = out.offset(1);
- *out = (*in_buf) >> 25;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 8)) << (15 - 8);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 8) % (1u32 << 15);
- out = out.offset(1);
- *out = (*in_buf) >> 23;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 6)) << (15 - 6);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 6) % (1u32 << 15);
- out = out.offset(1);
- *out = (*in_buf) >> 21;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 4)) << (15 - 4);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 4) % (1u32 << 15);
- out = out.offset(1);
- *out = (*in_buf) >> 19;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 2)) << (15 - 2);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 2) % (1u32 << 15);
- out = out.offset(1);
- *out = (*in_buf) >> 17;
-
- in_buf.offset(1)
-}
-
-unsafe fn unpack16_32(mut in_buf: *const u32, mut out: *mut u32) -> *const u32 {
- *out = (*in_buf) % (1u32 << 16);
- out = out.offset(1);
- *out = (*in_buf) >> 16;
- out = out.offset(1);
- in_buf = in_buf.offset(1);
-
- *out = (*in_buf) % (1u32 << 16);
- out = out.offset(1);
- *out = (*in_buf) >> 16;
- out = out.offset(1);
- in_buf = in_buf.offset(1);
-
- *out = (*in_buf) % (1u32 << 16);
- out = out.offset(1);
- *out = (*in_buf) >> 16;
- out = out.offset(1);
- in_buf = in_buf.offset(1);
-
- *out = (*in_buf) % (1u32 << 16);
- out = out.offset(1);
- *out = (*in_buf) >> 16;
- out = out.offset(1);
- in_buf = in_buf.offset(1);
-
- *out = (*in_buf) % (1u32 << 16);
- out = out.offset(1);
- *out = (*in_buf) >> 16;
- out = out.offset(1);
- in_buf = in_buf.offset(1);
-
- *out = (*in_buf) % (1u32 << 16);
- out = out.offset(1);
- *out = (*in_buf) >> 16;
- out = out.offset(1);
- in_buf = in_buf.offset(1);
-
- *out = (*in_buf) % (1u32 << 16);
- out = out.offset(1);
- *out = (*in_buf) >> 16;
- out = out.offset(1);
- in_buf = in_buf.offset(1);
-
- *out = (*in_buf) % (1u32 << 16);
- out = out.offset(1);
- *out = (*in_buf) >> 16;
- out = out.offset(1);
- in_buf = in_buf.offset(1);
-
- *out = (*in_buf) % (1u32 << 16);
- out = out.offset(1);
- *out = (*in_buf) >> 16;
- out = out.offset(1);
- in_buf = in_buf.offset(1);
-
- *out = (*in_buf) % (1u32 << 16);
- out = out.offset(1);
- *out = (*in_buf) >> 16;
- out = out.offset(1);
- in_buf = in_buf.offset(1);
-
- *out = (*in_buf) % (1u32 << 16);
- out = out.offset(1);
- *out = (*in_buf) >> 16;
- out = out.offset(1);
- in_buf = in_buf.offset(1);
-
- *out = (*in_buf) % (1u32 << 16);
- out = out.offset(1);
- *out = (*in_buf) >> 16;
- out = out.offset(1);
- in_buf = in_buf.offset(1);
-
- *out = (*in_buf) % (1u32 << 16);
- out = out.offset(1);
- *out = (*in_buf) >> 16;
- out = out.offset(1);
- in_buf = in_buf.offset(1);
-
- *out = (*in_buf) % (1u32 << 16);
- out = out.offset(1);
- *out = (*in_buf) >> 16;
- out = out.offset(1);
- in_buf = in_buf.offset(1);
-
- *out = (*in_buf) % (1u32 << 16);
- out = out.offset(1);
- *out = (*in_buf) >> 16;
- out = out.offset(1);
- in_buf = in_buf.offset(1);
-
- *out = (*in_buf) % (1u32 << 16);
- out = out.offset(1);
- *out = (*in_buf) >> 16;
-
- in_buf.offset(1)
-}
-
-unsafe fn unpack17_32(mut in_buf: *const u32, mut out: *mut u32) -> *const u32 {
- *out = (*in_buf) % (1u32 << 17);
- out = out.offset(1);
- *out = (*in_buf) >> 17;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 2)) << (17 - 2);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 2) % (1u32 << 17);
- out = out.offset(1);
- *out = (*in_buf) >> 19;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 4)) << (17 - 4);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 4) % (1u32 << 17);
- out = out.offset(1);
- *out = (*in_buf) >> 21;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 6)) << (17 - 6);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 6) % (1u32 << 17);
- out = out.offset(1);
- *out = (*in_buf) >> 23;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 8)) << (17 - 8);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 8) % (1u32 << 17);
- out = out.offset(1);
- *out = (*in_buf) >> 25;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 10)) << (17 - 10);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 10) % (1u32 << 17);
- out = out.offset(1);
- *out = (*in_buf) >> 27;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 12)) << (17 - 12);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 12) % (1u32 << 17);
- out = out.offset(1);
- *out = (*in_buf) >> 29;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 14)) << (17 - 14);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 14) % (1u32 << 17);
- out = out.offset(1);
- *out = (*in_buf) >> 31;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 16)) << (17 - 16);
- out = out.offset(1);
-
- *out = (*in_buf) >> 16;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 1)) << (17 - 1);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 1) % (1u32 << 17);
- out = out.offset(1);
- *out = (*in_buf) >> 18;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 3)) << (17 - 3);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 3) % (1u32 << 17);
- out = out.offset(1);
- *out = (*in_buf) >> 20;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 5)) << (17 - 5);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 5) % (1u32 << 17);
- out = out.offset(1);
- *out = (*in_buf) >> 22;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 7)) << (17 - 7);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 7) % (1u32 << 17);
- out = out.offset(1);
- *out = (*in_buf) >> 24;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 9)) << (17 - 9);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 9) % (1u32 << 17);
- out = out.offset(1);
- *out = (*in_buf) >> 26;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 11)) << (17 - 11);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 11) % (1u32 << 17);
- out = out.offset(1);
- *out = (*in_buf) >> 28;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 13)) << (17 - 13);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 13) % (1u32 << 17);
- out = out.offset(1);
- *out = (*in_buf) >> 30;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 15)) << (17 - 15);
- out = out.offset(1);
-
- *out = (*in_buf) >> 15;
-
- in_buf.offset(1)
-}
-
-unsafe fn unpack18_32(mut in_buf: *const u32, mut out: *mut u32) -> *const u32 {
- *out = (*in_buf) % (1u32 << 18);
- out = out.offset(1);
- *out = (*in_buf) >> 18;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 4)) << (18 - 4);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 4) % (1u32 << 18);
- out = out.offset(1);
- *out = (*in_buf) >> 22;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 8)) << (18 - 8);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 8) % (1u32 << 18);
- out = out.offset(1);
- *out = (*in_buf) >> 26;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 12)) << (18 - 12);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 12) % (1u32 << 18);
- out = out.offset(1);
- *out = (*in_buf) >> 30;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 16)) << (18 - 16);
- out = out.offset(1);
-
- *out = (*in_buf) >> 16;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 2)) << (18 - 2);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 2) % (1u32 << 18);
- out = out.offset(1);
- *out = (*in_buf) >> 20;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 6)) << (18 - 6);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 6) % (1u32 << 18);
- out = out.offset(1);
- *out = (*in_buf) >> 24;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 10)) << (18 - 10);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 10) % (1u32 << 18);
- out = out.offset(1);
- *out = (*in_buf) >> 28;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 14)) << (18 - 14);
- out = out.offset(1);
-
- *out = (*in_buf) >> 14;
- in_buf = in_buf.offset(1);
- out = out.offset(1);
-
- *out = (*in_buf) % (1u32 << 18);
- out = out.offset(1);
- *out = (*in_buf) >> 18;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 4)) << (18 - 4);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 4) % (1u32 << 18);
- out = out.offset(1);
- *out = (*in_buf) >> 22;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 8)) << (18 - 8);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 8) % (1u32 << 18);
- out = out.offset(1);
- *out = (*in_buf) >> 26;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 12)) << (18 - 12);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 12) % (1u32 << 18);
- out = out.offset(1);
- *out = (*in_buf) >> 30;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 16)) << (18 - 16);
- out = out.offset(1);
-
- *out = (*in_buf) >> 16;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 2)) << (18 - 2);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 2) % (1u32 << 18);
- out = out.offset(1);
- *out = (*in_buf) >> 20;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 6)) << (18 - 6);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 6) % (1u32 << 18);
- out = out.offset(1);
- *out = (*in_buf) >> 24;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 10)) << (18 - 10);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 10) % (1u32 << 18);
- out = out.offset(1);
- *out = (*in_buf) >> 28;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 14)) << (18 - 14);
- out = out.offset(1);
-
- *out = (*in_buf) >> 14;
-
- in_buf.offset(1)
-}
-
-unsafe fn unpack19_32(mut in_buf: *const u32, mut out: *mut u32) -> *const u32 {
- *out = (*in_buf) % (1u32 << 19);
- out = out.offset(1);
- *out = (*in_buf) >> 19;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 6)) << (19 - 6);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 6) % (1u32 << 19);
- out = out.offset(1);
- *out = (*in_buf) >> 25;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 12)) << (19 - 12);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 12) % (1u32 << 19);
- out = out.offset(1);
- *out = (*in_buf) >> 31;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 18)) << (19 - 18);
- out = out.offset(1);
-
- *out = (*in_buf) >> 18;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 5)) << (19 - 5);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 5) % (1u32 << 19);
- out = out.offset(1);
- *out = (*in_buf) >> 24;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 11)) << (19 - 11);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 11) % (1u32 << 19);
- out = out.offset(1);
- *out = (*in_buf) >> 30;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 17)) << (19 - 17);
- out = out.offset(1);
-
- *out = (*in_buf) >> 17;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 4)) << (19 - 4);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 4) % (1u32 << 19);
- out = out.offset(1);
- *out = (*in_buf) >> 23;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 10)) << (19 - 10);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 10) % (1u32 << 19);
- out = out.offset(1);
- *out = (*in_buf) >> 29;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 16)) << (19 - 16);
- out = out.offset(1);
-
- *out = (*in_buf) >> 16;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 3)) << (19 - 3);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 3) % (1u32 << 19);
- out = out.offset(1);
- *out = (*in_buf) >> 22;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 9)) << (19 - 9);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 9) % (1u32 << 19);
- out = out.offset(1);
- *out = (*in_buf) >> 28;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 15)) << (19 - 15);
- out = out.offset(1);
-
- *out = (*in_buf) >> 15;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 2)) << (19 - 2);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 2) % (1u32 << 19);
- out = out.offset(1);
- *out = (*in_buf) >> 21;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 8)) << (19 - 8);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 8) % (1u32 << 19);
- out = out.offset(1);
- *out = (*in_buf) >> 27;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 14)) << (19 - 14);
- out = out.offset(1);
-
- *out = (*in_buf) >> 14;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 1)) << (19 - 1);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 1) % (1u32 << 19);
- out = out.offset(1);
- *out = (*in_buf) >> 20;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 7)) << (19 - 7);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 7) % (1u32 << 19);
- out = out.offset(1);
- *out = (*in_buf) >> 26;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 13)) << (19 - 13);
- out = out.offset(1);
-
- *out = (*in_buf) >> 13;
-
- in_buf.offset(1)
-}
-
-unsafe fn unpack20_32(mut in_buf: *const u32, mut out: *mut u32) -> *const u32 {
- *out = (*in_buf) % (1u32 << 20);
- out = out.offset(1);
- *out = (*in_buf) >> 20;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 8)) << (20 - 8);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 8) % (1u32 << 20);
- out = out.offset(1);
- *out = (*in_buf) >> 28;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 16)) << (20 - 16);
- out = out.offset(1);
-
- *out = (*in_buf) >> 16;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 4)) << (20 - 4);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 4) % (1u32 << 20);
- out = out.offset(1);
- *out = (*in_buf) >> 24;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 12)) << (20 - 12);
- out = out.offset(1);
-
- *out = (*in_buf) >> 12;
- in_buf = in_buf.offset(1);
- out = out.offset(1);
-
- *out = (*in_buf) % (1u32 << 20);
- out = out.offset(1);
- *out = (*in_buf) >> 20;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 8)) << (20 - 8);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 8) % (1u32 << 20);
- out = out.offset(1);
- *out = (*in_buf) >> 28;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 16)) << (20 - 16);
- out = out.offset(1);
-
- *out = (*in_buf) >> 16;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 4)) << (20 - 4);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 4) % (1u32 << 20);
- out = out.offset(1);
- *out = (*in_buf) >> 24;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 12)) << (20 - 12);
- out = out.offset(1);
-
- *out = (*in_buf) >> 12;
- in_buf = in_buf.offset(1);
- out = out.offset(1);
-
- *out = (*in_buf) % (1u32 << 20);
- out = out.offset(1);
- *out = (*in_buf) >> 20;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 8)) << (20 - 8);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 8) % (1u32 << 20);
- out = out.offset(1);
- *out = (*in_buf) >> 28;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 16)) << (20 - 16);
- out = out.offset(1);
-
- *out = (*in_buf) >> 16;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 4)) << (20 - 4);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 4) % (1u32 << 20);
- out = out.offset(1);
- *out = (*in_buf) >> 24;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 12)) << (20 - 12);
- out = out.offset(1);
-
- *out = (*in_buf) >> 12;
- in_buf = in_buf.offset(1);
- out = out.offset(1);
-
- *out = (*in_buf) % (1u32 << 20);
- out = out.offset(1);
- *out = (*in_buf) >> 20;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 8)) << (20 - 8);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 8) % (1u32 << 20);
- out = out.offset(1);
- *out = (*in_buf) >> 28;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 16)) << (20 - 16);
- out = out.offset(1);
-
- *out = (*in_buf) >> 16;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 4)) << (20 - 4);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 4) % (1u32 << 20);
- out = out.offset(1);
- *out = (*in_buf) >> 24;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 12)) << (20 - 12);
- out = out.offset(1);
-
- *out = (*in_buf) >> 12;
-
- in_buf.offset(1)
-}
-
-unsafe fn unpack21_32(mut in_buf: *const u32, mut out: *mut u32) -> *const u32 {
- *out = (*in_buf) % (1u32 << 21);
- out = out.offset(1);
- *out = (*in_buf) >> 21;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 10)) << (21 - 10);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 10) % (1u32 << 21);
- out = out.offset(1);
- *out = (*in_buf) >> 31;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 20)) << (21 - 20);
- out = out.offset(1);
-
- *out = (*in_buf) >> 20;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 9)) << (21 - 9);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 9) % (1u32 << 21);
- out = out.offset(1);
- *out = (*in_buf) >> 30;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 19)) << (21 - 19);
- out = out.offset(1);
-
- *out = (*in_buf) >> 19;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 8)) << (21 - 8);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 8) % (1u32 << 21);
- out = out.offset(1);
- *out = (*in_buf) >> 29;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 18)) << (21 - 18);
- out = out.offset(1);
-
- *out = (*in_buf) >> 18;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 7)) << (21 - 7);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 7) % (1u32 << 21);
- out = out.offset(1);
- *out = (*in_buf) >> 28;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 17)) << (21 - 17);
- out = out.offset(1);
-
- *out = (*in_buf) >> 17;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 6)) << (21 - 6);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 6) % (1u32 << 21);
- out = out.offset(1);
- *out = (*in_buf) >> 27;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 16)) << (21 - 16);
- out = out.offset(1);
-
- *out = (*in_buf) >> 16;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 5)) << (21 - 5);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 5) % (1u32 << 21);
- out = out.offset(1);
- *out = (*in_buf) >> 26;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 15)) << (21 - 15);
- out = out.offset(1);
-
- *out = (*in_buf) >> 15;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 4)) << (21 - 4);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 4) % (1u32 << 21);
- out = out.offset(1);
- *out = (*in_buf) >> 25;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 14)) << (21 - 14);
- out = out.offset(1);
-
- *out = (*in_buf) >> 14;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 3)) << (21 - 3);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 3) % (1u32 << 21);
- out = out.offset(1);
- *out = (*in_buf) >> 24;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 13)) << (21 - 13);
- out = out.offset(1);
-
- *out = (*in_buf) >> 13;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 2)) << (21 - 2);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 2) % (1u32 << 21);
- out = out.offset(1);
- *out = (*in_buf) >> 23;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 12)) << (21 - 12);
- out = out.offset(1);
-
- *out = (*in_buf) >> 12;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 1)) << (21 - 1);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 1) % (1u32 << 21);
- out = out.offset(1);
- *out = (*in_buf) >> 22;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 11)) << (21 - 11);
- out = out.offset(1);
-
- *out = (*in_buf) >> 11;
-
- in_buf.offset(1)
-}
-
-unsafe fn unpack22_32(mut in_buf: *const u32, mut out: *mut u32) -> *const u32 {
- *out = (*in_buf) % (1u32 << 22);
- out = out.offset(1);
- *out = (*in_buf) >> 22;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 12)) << (22 - 12);
- out = out.offset(1);
-
- *out = (*in_buf) >> 12;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 2)) << (22 - 2);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 2) % (1u32 << 22);
- out = out.offset(1);
- *out = (*in_buf) >> 24;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 14)) << (22 - 14);
- out = out.offset(1);
-
- *out = (*in_buf) >> 14;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 4)) << (22 - 4);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 4) % (1u32 << 22);
- out = out.offset(1);
- *out = (*in_buf) >> 26;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 16)) << (22 - 16);
- out = out.offset(1);
-
- *out = (*in_buf) >> 16;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 6)) << (22 - 6);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 6) % (1u32 << 22);
- out = out.offset(1);
- *out = (*in_buf) >> 28;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 18)) << (22 - 18);
- out = out.offset(1);
-
- *out = (*in_buf) >> 18;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 8)) << (22 - 8);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 8) % (1u32 << 22);
- out = out.offset(1);
- *out = (*in_buf) >> 30;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 20)) << (22 - 20);
- out = out.offset(1);
-
- *out = (*in_buf) >> 20;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 10)) << (22 - 10);
- out = out.offset(1);
-
- *out = (*in_buf) >> 10;
- in_buf = in_buf.offset(1);
- out = out.offset(1);
-
- *out = (*in_buf) % (1u32 << 22);
- out = out.offset(1);
- *out = (*in_buf) >> 22;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 12)) << (22 - 12);
- out = out.offset(1);
-
- *out = (*in_buf) >> 12;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 2)) << (22 - 2);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 2) % (1u32 << 22);
- out = out.offset(1);
- *out = (*in_buf) >> 24;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 14)) << (22 - 14);
- out = out.offset(1);
-
- *out = (*in_buf) >> 14;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 4)) << (22 - 4);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 4) % (1u32 << 22);
- out = out.offset(1);
- *out = (*in_buf) >> 26;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 16)) << (22 - 16);
- out = out.offset(1);
-
- *out = (*in_buf) >> 16;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 6)) << (22 - 6);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 6) % (1u32 << 22);
- out = out.offset(1);
- *out = (*in_buf) >> 28;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 18)) << (22 - 18);
- out = out.offset(1);
-
- *out = (*in_buf) >> 18;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 8)) << (22 - 8);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 8) % (1u32 << 22);
- out = out.offset(1);
- *out = (*in_buf) >> 30;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 20)) << (22 - 20);
- out = out.offset(1);
-
- *out = (*in_buf) >> 20;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 10)) << (22 - 10);
- out = out.offset(1);
-
- *out = (*in_buf) >> 10;
-
- in_buf.offset(1)
-}
-
-unsafe fn unpack23_32(mut in_buf: *const u32, mut out: *mut u32) -> *const u32 {
- *out = (*in_buf) % (1u32 << 23);
- out = out.offset(1);
- *out = (*in_buf) >> 23;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 14)) << (23 - 14);
- out = out.offset(1);
-
- *out = (*in_buf) >> 14;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 5)) << (23 - 5);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 5) % (1u32 << 23);
- out = out.offset(1);
- *out = (*in_buf) >> 28;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 19)) << (23 - 19);
- out = out.offset(1);
-
- *out = (*in_buf) >> 19;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 10)) << (23 - 10);
- out = out.offset(1);
-
- *out = (*in_buf) >> 10;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 1)) << (23 - 1);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 1) % (1u32 << 23);
- out = out.offset(1);
- *out = (*in_buf) >> 24;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 15)) << (23 - 15);
- out = out.offset(1);
-
- *out = (*in_buf) >> 15;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 6)) << (23 - 6);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 6) % (1u32 << 23);
- out = out.offset(1);
- *out = (*in_buf) >> 29;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 20)) << (23 - 20);
- out = out.offset(1);
-
- *out = (*in_buf) >> 20;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 11)) << (23 - 11);
- out = out.offset(1);
-
- *out = (*in_buf) >> 11;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 2)) << (23 - 2);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 2) % (1u32 << 23);
- out = out.offset(1);
- *out = (*in_buf) >> 25;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 16)) << (23 - 16);
- out = out.offset(1);
-
- *out = (*in_buf) >> 16;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 7)) << (23 - 7);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 7) % (1u32 << 23);
- out = out.offset(1);
- *out = (*in_buf) >> 30;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 21)) << (23 - 21);
- out = out.offset(1);
-
- *out = (*in_buf) >> 21;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 12)) << (23 - 12);
- out = out.offset(1);
-
- *out = (*in_buf) >> 12;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 3)) << (23 - 3);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 3) % (1u32 << 23);
- out = out.offset(1);
- *out = (*in_buf) >> 26;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 17)) << (23 - 17);
- out = out.offset(1);
-
- *out = (*in_buf) >> 17;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 8)) << (23 - 8);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 8) % (1u32 << 23);
- out = out.offset(1);
- *out = (*in_buf) >> 31;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 22)) << (23 - 22);
- out = out.offset(1);
-
- *out = (*in_buf) >> 22;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 13)) << (23 - 13);
- out = out.offset(1);
-
- *out = (*in_buf) >> 13;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 4)) << (23 - 4);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 4) % (1u32 << 23);
- out = out.offset(1);
- *out = (*in_buf) >> 27;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 18)) << (23 - 18);
- out = out.offset(1);
-
- *out = (*in_buf) >> 18;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 9)) << (23 - 9);
- out = out.offset(1);
-
- *out = (*in_buf) >> 9;
-
- in_buf.offset(1)
-}
-
-unsafe fn unpack24_32(mut in_buf: *const u32, mut out: *mut u32) -> *const u32 {
- *out = (*in_buf) % (1u32 << 24);
- out = out.offset(1);
- *out = (*in_buf) >> 24;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 16)) << (24 - 16);
- out = out.offset(1);
-
- *out = (*in_buf) >> 16;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 8)) << (24 - 8);
- out = out.offset(1);
-
- *out = (*in_buf) >> 8;
- in_buf = in_buf.offset(1);
- out = out.offset(1);
-
- *out = (*in_buf) % (1u32 << 24);
- out = out.offset(1);
- *out = (*in_buf) >> 24;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 16)) << (24 - 16);
- out = out.offset(1);
-
- *out = (*in_buf) >> 16;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 8)) << (24 - 8);
- out = out.offset(1);
-
- *out = (*in_buf) >> 8;
- in_buf = in_buf.offset(1);
- out = out.offset(1);
-
- *out = (*in_buf) % (1u32 << 24);
- out = out.offset(1);
- *out = (*in_buf) >> 24;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 16)) << (24 - 16);
- out = out.offset(1);
-
- *out = (*in_buf) >> 16;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 8)) << (24 - 8);
- out = out.offset(1);
-
- *out = (*in_buf) >> 8;
- in_buf = in_buf.offset(1);
- out = out.offset(1);
-
- *out = (*in_buf) % (1u32 << 24);
- out = out.offset(1);
- *out = (*in_buf) >> 24;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 16)) << (24 - 16);
- out = out.offset(1);
-
- *out = (*in_buf) >> 16;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 8)) << (24 - 8);
- out = out.offset(1);
-
- *out = (*in_buf) >> 8;
- in_buf = in_buf.offset(1);
- out = out.offset(1);
-
- *out = (*in_buf) % (1u32 << 24);
- out = out.offset(1);
- *out = (*in_buf) >> 24;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 16)) << (24 - 16);
- out = out.offset(1);
-
- *out = (*in_buf) >> 16;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 8)) << (24 - 8);
- out = out.offset(1);
-
- *out = (*in_buf) >> 8;
- in_buf = in_buf.offset(1);
- out = out.offset(1);
-
- *out = (*in_buf) % (1u32 << 24);
- out = out.offset(1);
- *out = (*in_buf) >> 24;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 16)) << (24 - 16);
- out = out.offset(1);
-
- *out = (*in_buf) >> 16;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 8)) << (24 - 8);
- out = out.offset(1);
-
- *out = (*in_buf) >> 8;
- in_buf = in_buf.offset(1);
- out = out.offset(1);
-
- *out = (*in_buf) % (1u32 << 24);
- out = out.offset(1);
- *out = (*in_buf) >> 24;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 16)) << (24 - 16);
- out = out.offset(1);
-
- *out = (*in_buf) >> 16;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 8)) << (24 - 8);
- out = out.offset(1);
-
- *out = (*in_buf) >> 8;
- in_buf = in_buf.offset(1);
- out = out.offset(1);
-
- *out = (*in_buf) % (1u32 << 24);
- out = out.offset(1);
- *out = (*in_buf) >> 24;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 16)) << (24 - 16);
- out = out.offset(1);
-
- *out = (*in_buf) >> 16;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 8)) << (24 - 8);
- out = out.offset(1);
-
- *out = (*in_buf) >> 8;
-
- in_buf.offset(1)
-}
-
-unsafe fn unpack25_32(mut in_buf: *const u32, mut out: *mut u32) -> *const u32 {
- *out = (*in_buf) % (1u32 << 25);
- out = out.offset(1);
- *out = (*in_buf) >> 25;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 18)) << (25 - 18);
- out = out.offset(1);
-
- *out = (*in_buf) >> 18;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 11)) << (25 - 11);
- out = out.offset(1);
-
- *out = (*in_buf) >> 11;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 4)) << (25 - 4);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 4) % (1u32 << 25);
- out = out.offset(1);
- *out = (*in_buf) >> 29;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 22)) << (25 - 22);
- out = out.offset(1);
-
- *out = (*in_buf) >> 22;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 15)) << (25 - 15);
- out = out.offset(1);
-
- *out = (*in_buf) >> 15;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 8)) << (25 - 8);
- out = out.offset(1);
-
- *out = (*in_buf) >> 8;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 1)) << (25 - 1);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 1) % (1u32 << 25);
- out = out.offset(1);
- *out = (*in_buf) >> 26;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 19)) << (25 - 19);
- out = out.offset(1);
-
- *out = (*in_buf) >> 19;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 12)) << (25 - 12);
- out = out.offset(1);
-
- *out = (*in_buf) >> 12;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 5)) << (25 - 5);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 5) % (1u32 << 25);
- out = out.offset(1);
- *out = (*in_buf) >> 30;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 23)) << (25 - 23);
- out = out.offset(1);
-
- *out = (*in_buf) >> 23;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 16)) << (25 - 16);
- out = out.offset(1);
-
- *out = (*in_buf) >> 16;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 9)) << (25 - 9);
- out = out.offset(1);
-
- *out = (*in_buf) >> 9;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 2)) << (25 - 2);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 2) % (1u32 << 25);
- out = out.offset(1);
- *out = (*in_buf) >> 27;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 20)) << (25 - 20);
- out = out.offset(1);
-
- *out = (*in_buf) >> 20;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 13)) << (25 - 13);
- out = out.offset(1);
-
- *out = (*in_buf) >> 13;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 6)) << (25 - 6);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 6) % (1u32 << 25);
- out = out.offset(1);
- *out = (*in_buf) >> 31;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 24)) << (25 - 24);
- out = out.offset(1);
-
- *out = (*in_buf) >> 24;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 17)) << (25 - 17);
- out = out.offset(1);
-
- *out = (*in_buf) >> 17;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 10)) << (25 - 10);
- out = out.offset(1);
-
- *out = (*in_buf) >> 10;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 3)) << (25 - 3);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 3) % (1u32 << 25);
- out = out.offset(1);
- *out = (*in_buf) >> 28;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 21)) << (25 - 21);
- out = out.offset(1);
-
- *out = (*in_buf) >> 21;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 14)) << (25 - 14);
- out = out.offset(1);
-
- *out = (*in_buf) >> 14;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 7)) << (25 - 7);
- out = out.offset(1);
-
- *out = (*in_buf) >> 7;
-
- in_buf.offset(1)
-}
-
-unsafe fn unpack26_32(mut in_buf: *const u32, mut out: *mut u32) -> *const u32 {
- *out = (*in_buf) % (1u32 << 26);
- out = out.offset(1);
- *out = (*in_buf) >> 26;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 20)) << (26 - 20);
- out = out.offset(1);
-
- *out = (*in_buf) >> 20;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 14)) << (26 - 14);
- out = out.offset(1);
-
- *out = (*in_buf) >> 14;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 8)) << (26 - 8);
- out = out.offset(1);
-
- *out = (*in_buf) >> 8;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 2)) << (26 - 2);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 2) % (1u32 << 26);
- out = out.offset(1);
- *out = (*in_buf) >> 28;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 22)) << (26 - 22);
- out = out.offset(1);
-
- *out = (*in_buf) >> 22;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 16)) << (26 - 16);
- out = out.offset(1);
-
- *out = (*in_buf) >> 16;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 10)) << (26 - 10);
- out = out.offset(1);
-
- *out = (*in_buf) >> 10;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 4)) << (26 - 4);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 4) % (1u32 << 26);
- out = out.offset(1);
- *out = (*in_buf) >> 30;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 24)) << (26 - 24);
- out = out.offset(1);
-
- *out = (*in_buf) >> 24;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 18)) << (26 - 18);
- out = out.offset(1);
-
- *out = (*in_buf) >> 18;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 12)) << (26 - 12);
- out = out.offset(1);
-
- *out = (*in_buf) >> 12;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 6)) << (26 - 6);
- out = out.offset(1);
-
- *out = (*in_buf) >> 6;
- in_buf = in_buf.offset(1);
- out = out.offset(1);
-
- *out = (*in_buf) % (1u32 << 26);
- out = out.offset(1);
- *out = (*in_buf) >> 26;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 20)) << (26 - 20);
- out = out.offset(1);
-
- *out = (*in_buf) >> 20;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 14)) << (26 - 14);
- out = out.offset(1);
-
- *out = (*in_buf) >> 14;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 8)) << (26 - 8);
- out = out.offset(1);
-
- *out = (*in_buf) >> 8;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 2)) << (26 - 2);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 2) % (1u32 << 26);
- out = out.offset(1);
- *out = (*in_buf) >> 28;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 22)) << (26 - 22);
- out = out.offset(1);
-
- *out = (*in_buf) >> 22;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 16)) << (26 - 16);
- out = out.offset(1);
-
- *out = (*in_buf) >> 16;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 10)) << (26 - 10);
- out = out.offset(1);
-
- *out = (*in_buf) >> 10;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 4)) << (26 - 4);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 4) % (1u32 << 26);
- out = out.offset(1);
- *out = (*in_buf) >> 30;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 24)) << (26 - 24);
- out = out.offset(1);
-
- *out = (*in_buf) >> 24;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 18)) << (26 - 18);
- out = out.offset(1);
-
- *out = (*in_buf) >> 18;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 12)) << (26 - 12);
- out = out.offset(1);
-
- *out = (*in_buf) >> 12;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 6)) << (26 - 6);
- out = out.offset(1);
-
- *out = (*in_buf) >> 6;
-
- in_buf.offset(1)
-}
-
-unsafe fn unpack27_32(mut in_buf: *const u32, mut out: *mut u32) -> *const u32 {
- *out = (*in_buf) % (1u32 << 27);
- out = out.offset(1);
- *out = (*in_buf) >> 27;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 22)) << (27 - 22);
- out = out.offset(1);
-
- *out = (*in_buf) >> 22;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 17)) << (27 - 17);
- out = out.offset(1);
-
- *out = (*in_buf) >> 17;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 12)) << (27 - 12);
- out = out.offset(1);
-
- *out = (*in_buf) >> 12;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 7)) << (27 - 7);
- out = out.offset(1);
-
- *out = (*in_buf) >> 7;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 2)) << (27 - 2);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 2) % (1u32 << 27);
- out = out.offset(1);
- *out = (*in_buf) >> 29;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 24)) << (27 - 24);
- out = out.offset(1);
-
- *out = (*in_buf) >> 24;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 19)) << (27 - 19);
- out = out.offset(1);
-
- *out = (*in_buf) >> 19;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 14)) << (27 - 14);
- out = out.offset(1);
-
- *out = (*in_buf) >> 14;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 9)) << (27 - 9);
- out = out.offset(1);
-
- *out = (*in_buf) >> 9;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 4)) << (27 - 4);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 4) % (1u32 << 27);
- out = out.offset(1);
- *out = (*in_buf) >> 31;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 26)) << (27 - 26);
- out = out.offset(1);
-
- *out = (*in_buf) >> 26;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 21)) << (27 - 21);
- out = out.offset(1);
-
- *out = (*in_buf) >> 21;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 16)) << (27 - 16);
- out = out.offset(1);
-
- *out = (*in_buf) >> 16;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 11)) << (27 - 11);
- out = out.offset(1);
-
- *out = (*in_buf) >> 11;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 6)) << (27 - 6);
- out = out.offset(1);
-
- *out = (*in_buf) >> 6;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 1)) << (27 - 1);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 1) % (1u32 << 27);
- out = out.offset(1);
- *out = (*in_buf) >> 28;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 23)) << (27 - 23);
- out = out.offset(1);
-
- *out = (*in_buf) >> 23;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 18)) << (27 - 18);
- out = out.offset(1);
-
- *out = (*in_buf) >> 18;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 13)) << (27 - 13);
- out = out.offset(1);
-
- *out = (*in_buf) >> 13;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 8)) << (27 - 8);
- out = out.offset(1);
-
- *out = (*in_buf) >> 8;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 3)) << (27 - 3);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 3) % (1u32 << 27);
- out = out.offset(1);
- *out = (*in_buf) >> 30;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 25)) << (27 - 25);
- out = out.offset(1);
-
- *out = (*in_buf) >> 25;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 20)) << (27 - 20);
- out = out.offset(1);
-
- *out = (*in_buf) >> 20;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 15)) << (27 - 15);
- out = out.offset(1);
-
- *out = (*in_buf) >> 15;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 10)) << (27 - 10);
- out = out.offset(1);
-
- *out = (*in_buf) >> 10;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 5)) << (27 - 5);
- out = out.offset(1);
-
- *out = (*in_buf) >> 5;
-
- in_buf.offset(1)
-}
-
-unsafe fn unpack28_32(mut in_buf: *const u32, mut out: *mut u32) -> *const u32 {
- *out = (*in_buf) % (1u32 << 28);
- out = out.offset(1);
- *out = (*in_buf) >> 28;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 24)) << (28 - 24);
- out = out.offset(1);
-
- *out = (*in_buf) >> 24;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 20)) << (28 - 20);
- out = out.offset(1);
-
- *out = (*in_buf) >> 20;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 16)) << (28 - 16);
- out = out.offset(1);
-
- *out = (*in_buf) >> 16;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 12)) << (28 - 12);
- out = out.offset(1);
-
- *out = (*in_buf) >> 12;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 8)) << (28 - 8);
- out = out.offset(1);
-
- *out = (*in_buf) >> 8;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 4)) << (28 - 4);
- out = out.offset(1);
-
- *out = (*in_buf) >> 4;
- in_buf = in_buf.offset(1);
- out = out.offset(1);
-
- *out = (*in_buf) % (1u32 << 28);
- out = out.offset(1);
- *out = (*in_buf) >> 28;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 24)) << (28 - 24);
- out = out.offset(1);
-
- *out = (*in_buf) >> 24;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 20)) << (28 - 20);
- out = out.offset(1);
-
- *out = (*in_buf) >> 20;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 16)) << (28 - 16);
- out = out.offset(1);
-
- *out = (*in_buf) >> 16;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 12)) << (28 - 12);
- out = out.offset(1);
-
- *out = (*in_buf) >> 12;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 8)) << (28 - 8);
- out = out.offset(1);
-
- *out = (*in_buf) >> 8;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 4)) << (28 - 4);
- out = out.offset(1);
-
- *out = (*in_buf) >> 4;
- in_buf = in_buf.offset(1);
- out = out.offset(1);
-
- *out = (*in_buf) % (1u32 << 28);
- out = out.offset(1);
- *out = (*in_buf) >> 28;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 24)) << (28 - 24);
- out = out.offset(1);
-
- *out = (*in_buf) >> 24;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 20)) << (28 - 20);
- out = out.offset(1);
-
- *out = (*in_buf) >> 20;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 16)) << (28 - 16);
- out = out.offset(1);
-
- *out = (*in_buf) >> 16;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 12)) << (28 - 12);
- out = out.offset(1);
-
- *out = (*in_buf) >> 12;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 8)) << (28 - 8);
- out = out.offset(1);
-
- *out = (*in_buf) >> 8;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 4)) << (28 - 4);
- out = out.offset(1);
-
- *out = (*in_buf) >> 4;
- in_buf = in_buf.offset(1);
- out = out.offset(1);
-
- *out = (*in_buf) % (1u32 << 28);
- out = out.offset(1);
- *out = (*in_buf) >> 28;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 24)) << (28 - 24);
- out = out.offset(1);
-
- *out = (*in_buf) >> 24;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 20)) << (28 - 20);
- out = out.offset(1);
-
- *out = (*in_buf) >> 20;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 16)) << (28 - 16);
- out = out.offset(1);
-
- *out = (*in_buf) >> 16;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 12)) << (28 - 12);
- out = out.offset(1);
-
- *out = (*in_buf) >> 12;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 8)) << (28 - 8);
- out = out.offset(1);
-
- *out = (*in_buf) >> 8;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 4)) << (28 - 4);
- out = out.offset(1);
-
- *out = (*in_buf) >> 4;
-
- in_buf.offset(1)
-}
-
-unsafe fn unpack29_32(mut in_buf: *const u32, mut out: *mut u32) -> *const u32 {
- *out = (*in_buf) % (1u32 << 29);
- out = out.offset(1);
- *out = (*in_buf) >> 29;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 26)) << (29 - 26);
- out = out.offset(1);
-
- *out = (*in_buf) >> 26;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 23)) << (29 - 23);
- out = out.offset(1);
-
- *out = (*in_buf) >> 23;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 20)) << (29 - 20);
- out = out.offset(1);
-
- *out = (*in_buf) >> 20;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 17)) << (29 - 17);
- out = out.offset(1);
-
- *out = (*in_buf) >> 17;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 14)) << (29 - 14);
- out = out.offset(1);
-
- *out = (*in_buf) >> 14;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 11)) << (29 - 11);
- out = out.offset(1);
-
- *out = (*in_buf) >> 11;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 8)) << (29 - 8);
- out = out.offset(1);
-
- *out = (*in_buf) >> 8;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 5)) << (29 - 5);
- out = out.offset(1);
-
- *out = (*in_buf) >> 5;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 2)) << (29 - 2);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 2) % (1u32 << 29);
- out = out.offset(1);
- *out = (*in_buf) >> 31;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 28)) << (29 - 28);
- out = out.offset(1);
-
- *out = (*in_buf) >> 28;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 25)) << (29 - 25);
- out = out.offset(1);
-
- *out = (*in_buf) >> 25;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 22)) << (29 - 22);
- out = out.offset(1);
-
- *out = (*in_buf) >> 22;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 19)) << (29 - 19);
- out = out.offset(1);
-
- *out = (*in_buf) >> 19;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 16)) << (29 - 16);
- out = out.offset(1);
-
- *out = (*in_buf) >> 16;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 13)) << (29 - 13);
- out = out.offset(1);
-
- *out = (*in_buf) >> 13;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 10)) << (29 - 10);
- out = out.offset(1);
-
- *out = (*in_buf) >> 10;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 7)) << (29 - 7);
- out = out.offset(1);
-
- *out = (*in_buf) >> 7;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 4)) << (29 - 4);
- out = out.offset(1);
-
- *out = (*in_buf) >> 4;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 1)) << (29 - 1);
- out = out.offset(1);
-
- *out = ((*in_buf) >> 1) % (1u32 << 29);
- out = out.offset(1);
- *out = (*in_buf) >> 30;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 27)) << (29 - 27);
- out = out.offset(1);
-
- *out = (*in_buf) >> 27;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 24)) << (29 - 24);
- out = out.offset(1);
-
- *out = (*in_buf) >> 24;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 21)) << (29 - 21);
- out = out.offset(1);
-
- *out = (*in_buf) >> 21;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 18)) << (29 - 18);
- out = out.offset(1);
-
- *out = (*in_buf) >> 18;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 15)) << (29 - 15);
- out = out.offset(1);
-
- *out = (*in_buf) >> 15;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 12)) << (29 - 12);
- out = out.offset(1);
-
- *out = (*in_buf) >> 12;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 9)) << (29 - 9);
- out = out.offset(1);
-
- *out = (*in_buf) >> 9;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 6)) << (29 - 6);
- out = out.offset(1);
-
- *out = (*in_buf) >> 6;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 3)) << (29 - 3);
- out = out.offset(1);
-
- *out = (*in_buf) >> 3;
-
- in_buf.offset(1)
-}
-
-unsafe fn unpack30_32(mut in_buf: *const u32, mut out: *mut u32) -> *const u32 {
- *out = (*in_buf) % (1u32 << 30);
- out = out.offset(1);
- *out = (*in_buf) >> 30;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 28)) << (30 - 28);
- out = out.offset(1);
-
- *out = (*in_buf) >> 28;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 26)) << (30 - 26);
- out = out.offset(1);
-
- *out = (*in_buf) >> 26;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 24)) << (30 - 24);
- out = out.offset(1);
-
- *out = (*in_buf) >> 24;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 22)) << (30 - 22);
- out = out.offset(1);
-
- *out = (*in_buf) >> 22;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 20)) << (30 - 20);
- out = out.offset(1);
-
- *out = (*in_buf) >> 20;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 18)) << (30 - 18);
- out = out.offset(1);
-
- *out = (*in_buf) >> 18;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 16)) << (30 - 16);
- out = out.offset(1);
-
- *out = (*in_buf) >> 16;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 14)) << (30 - 14);
- out = out.offset(1);
-
- *out = (*in_buf) >> 14;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 12)) << (30 - 12);
- out = out.offset(1);
-
- *out = (*in_buf) >> 12;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 10)) << (30 - 10);
- out = out.offset(1);
-
- *out = (*in_buf) >> 10;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 8)) << (30 - 8);
- out = out.offset(1);
-
- *out = (*in_buf) >> 8;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 6)) << (30 - 6);
- out = out.offset(1);
-
- *out = (*in_buf) >> 6;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 4)) << (30 - 4);
- out = out.offset(1);
-
- *out = (*in_buf) >> 4;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 2)) << (30 - 2);
- out = out.offset(1);
-
- *out = (*in_buf) >> 2;
- in_buf = in_buf.offset(1);
- out = out.offset(1);
-
- *out = (*in_buf) % (1u32 << 30);
- out = out.offset(1);
- *out = (*in_buf) >> 30;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 28)) << (30 - 28);
- out = out.offset(1);
-
- *out = (*in_buf) >> 28;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 26)) << (30 - 26);
- out = out.offset(1);
-
- *out = (*in_buf) >> 26;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 24)) << (30 - 24);
- out = out.offset(1);
-
- *out = (*in_buf) >> 24;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 22)) << (30 - 22);
- out = out.offset(1);
-
- *out = (*in_buf) >> 22;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 20)) << (30 - 20);
- out = out.offset(1);
-
- *out = (*in_buf) >> 20;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 18)) << (30 - 18);
- out = out.offset(1);
-
- *out = (*in_buf) >> 18;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 16)) << (30 - 16);
- out = out.offset(1);
-
- *out = (*in_buf) >> 16;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 14)) << (30 - 14);
- out = out.offset(1);
-
- *out = (*in_buf) >> 14;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 12)) << (30 - 12);
- out = out.offset(1);
-
- *out = (*in_buf) >> 12;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 10)) << (30 - 10);
- out = out.offset(1);
-
- *out = (*in_buf) >> 10;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 8)) << (30 - 8);
- out = out.offset(1);
-
- *out = (*in_buf) >> 8;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 6)) << (30 - 6);
- out = out.offset(1);
-
- *out = (*in_buf) >> 6;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 4)) << (30 - 4);
- out = out.offset(1);
-
- *out = (*in_buf) >> 4;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 2)) << (30 - 2);
- out = out.offset(1);
-
- *out = (*in_buf) >> 2;
-
- in_buf.offset(1)
-}
-
-unsafe fn unpack31_32(mut in_buf: *const u32, mut out: *mut u32) -> *const u32 {
- *out = (*in_buf) % (1u32 << 31);
- out = out.offset(1);
- *out = (*in_buf) >> 31;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 30)) << (31 - 30);
- out = out.offset(1);
-
- *out = (*in_buf) >> 30;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 29)) << (31 - 29);
- out = out.offset(1);
-
- *out = (*in_buf) >> 29;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 28)) << (31 - 28);
- out = out.offset(1);
-
- *out = (*in_buf) >> 28;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 27)) << (31 - 27);
- out = out.offset(1);
-
- *out = (*in_buf) >> 27;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 26)) << (31 - 26);
- out = out.offset(1);
-
- *out = (*in_buf) >> 26;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 25)) << (31 - 25);
- out = out.offset(1);
-
- *out = (*in_buf) >> 25;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 24)) << (31 - 24);
- out = out.offset(1);
-
- *out = (*in_buf) >> 24;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 23)) << (31 - 23);
- out = out.offset(1);
-
- *out = (*in_buf) >> 23;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 22)) << (31 - 22);
- out = out.offset(1);
-
- *out = (*in_buf) >> 22;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 21)) << (31 - 21);
- out = out.offset(1);
-
- *out = (*in_buf) >> 21;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 20)) << (31 - 20);
- out = out.offset(1);
-
- *out = (*in_buf) >> 20;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 19)) << (31 - 19);
- out = out.offset(1);
-
- *out = (*in_buf) >> 19;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 18)) << (31 - 18);
- out = out.offset(1);
-
- *out = (*in_buf) >> 18;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 17)) << (31 - 17);
- out = out.offset(1);
-
- *out = (*in_buf) >> 17;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 16)) << (31 - 16);
- out = out.offset(1);
-
- *out = (*in_buf) >> 16;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 15)) << (31 - 15);
- out = out.offset(1);
-
- *out = (*in_buf) >> 15;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 14)) << (31 - 14);
- out = out.offset(1);
-
- *out = (*in_buf) >> 14;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 13)) << (31 - 13);
- out = out.offset(1);
-
- *out = (*in_buf) >> 13;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 12)) << (31 - 12);
- out = out.offset(1);
-
- *out = (*in_buf) >> 12;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 11)) << (31 - 11);
- out = out.offset(1);
-
- *out = (*in_buf) >> 11;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 10)) << (31 - 10);
- out = out.offset(1);
-
- *out = (*in_buf) >> 10;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 9)) << (31 - 9);
- out = out.offset(1);
-
- *out = (*in_buf) >> 9;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 8)) << (31 - 8);
- out = out.offset(1);
-
- *out = (*in_buf) >> 8;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 7)) << (31 - 7);
- out = out.offset(1);
-
- *out = (*in_buf) >> 7;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 6)) << (31 - 6);
- out = out.offset(1);
-
- *out = (*in_buf) >> 6;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 5)) << (31 - 5);
- out = out.offset(1);
-
- *out = (*in_buf) >> 5;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 4)) << (31 - 4);
- out = out.offset(1);
-
- *out = (*in_buf) >> 4;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 3)) << (31 - 3);
- out = out.offset(1);
-
- *out = (*in_buf) >> 3;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 2)) << (31 - 2);
- out = out.offset(1);
-
- *out = (*in_buf) >> 2;
- in_buf = in_buf.offset(1);
- *out |= ((*in_buf) % (1u32 << 1)) << (31 - 1);
- out = out.offset(1);
-
- *out = (*in_buf) >> 1;
-
- in_buf.offset(1)
-}
-
-unsafe fn unpack32_32(mut in_buf: *const u32, mut out: *mut u32) -> *const u32 {
- *out = *in_buf;
- in_buf = in_buf.offset(1);
- out = out.offset(1);
-
- *out = *in_buf;
- in_buf = in_buf.offset(1);
- out = out.offset(1);
-
- *out = *in_buf;
- in_buf = in_buf.offset(1);
- out = out.offset(1);
-
- *out = *in_buf;
- in_buf = in_buf.offset(1);
- out = out.offset(1);
-
- *out = *in_buf;
- in_buf = in_buf.offset(1);
- out = out.offset(1);
-
- *out = *in_buf;
- in_buf = in_buf.offset(1);
- out = out.offset(1);
-
- *out = *in_buf;
- in_buf = in_buf.offset(1);
- out = out.offset(1);
-
- *out = *in_buf;
- in_buf = in_buf.offset(1);
- out = out.offset(1);
-
- *out = *in_buf;
- in_buf = in_buf.offset(1);
- out = out.offset(1);
-
- *out = *in_buf;
- in_buf = in_buf.offset(1);
- out = out.offset(1);
-
- *out = *in_buf;
- in_buf = in_buf.offset(1);
- out = out.offset(1);
-
- *out = *in_buf;
- in_buf = in_buf.offset(1);
- out = out.offset(1);
-
- *out = *in_buf;
- in_buf = in_buf.offset(1);
- out = out.offset(1);
-
- *out = *in_buf;
- in_buf = in_buf.offset(1);
- out = out.offset(1);
-
- *out = *in_buf;
- in_buf = in_buf.offset(1);
- out = out.offset(1);
-
- *out = *in_buf;
- in_buf = in_buf.offset(1);
- out = out.offset(1);
-
- *out = *in_buf;
- in_buf = in_buf.offset(1);
- out = out.offset(1);
-
- *out = *in_buf;
- in_buf = in_buf.offset(1);
- out = out.offset(1);
-
- *out = *in_buf;
- in_buf = in_buf.offset(1);
- out = out.offset(1);
-
- *out = *in_buf;
- in_buf = in_buf.offset(1);
- out = out.offset(1);
-
- *out = *in_buf;
- in_buf = in_buf.offset(1);
- out = out.offset(1);
-
- *out = *in_buf;
- in_buf = in_buf.offset(1);
- out = out.offset(1);
-
- *out = *in_buf;
- in_buf = in_buf.offset(1);
- out = out.offset(1);
-
- *out = *in_buf;
- in_buf = in_buf.offset(1);
- out = out.offset(1);
-
- *out = *in_buf;
- in_buf = in_buf.offset(1);
- out = out.offset(1);
-
- *out = *in_buf;
- in_buf = in_buf.offset(1);
- out = out.offset(1);
-
- *out = *in_buf;
- in_buf = in_buf.offset(1);
- out = out.offset(1);
-
- *out = *in_buf;
- in_buf = in_buf.offset(1);
- out = out.offset(1);
-
- *out = *in_buf;
- in_buf = in_buf.offset(1);
- out = out.offset(1);
-
- *out = *in_buf;
- in_buf = in_buf.offset(1);
- out = out.offset(1);
-
- *out = *in_buf;
- in_buf = in_buf.offset(1);
- out = out.offset(1);
-
- *out = *in_buf;
-
- in_buf.offset(1)
-}
diff --git a/parquet/src/util/bit_util.rs b/parquet/src/util/bit_util.rs
deleted file mode 100644
index 45cfe2b..0000000
--- a/parquet/src/util/bit_util.rs
+++ /dev/null
@@ -1,1150 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-use std::{cmp, mem::size_of};
-
-use crate::data_type::AsBytes;
-use crate::errors::{ParquetError, Result};
-use crate::util::{bit_packing::unpack32, memory::ByteBufferPtr};
-
-#[inline]
-pub fn from_ne_slice<T: FromBytes>(bs: &[u8]) -> T {
- let mut b = T::Buffer::default();
- {
- let b = b.as_mut();
- let bs = &bs[..b.len()];
- b.copy_from_slice(bs);
- }
- T::from_ne_bytes(b)
-}
-
-pub trait FromBytes: Sized {
- type Buffer: AsMut<[u8]> + Default;
- fn from_le_bytes(bs: Self::Buffer) -> Self;
- fn from_be_bytes(bs: Self::Buffer) -> Self;
- fn from_ne_bytes(bs: Self::Buffer) -> Self;
-}
-
-macro_rules! from_le_bytes {
- ($($ty: ty),*) => {
- $(
- impl FromBytes for $ty {
- type Buffer = [u8; size_of::<Self>()];
- fn from_le_bytes(bs: Self::Buffer) -> Self {
- <$ty>::from_le_bytes(bs)
- }
- fn from_be_bytes(bs: Self::Buffer) -> Self {
- <$ty>::from_be_bytes(bs)
- }
- fn from_ne_bytes(bs: Self::Buffer) -> Self {
- <$ty>::from_ne_bytes(bs)
- }
- }
- )*
- };
-}
-
-impl FromBytes for bool {
- type Buffer = [u8; 1];
- fn from_le_bytes(bs: Self::Buffer) -> Self {
- Self::from_ne_bytes(bs)
- }
- fn from_be_bytes(bs: Self::Buffer) -> Self {
- Self::from_ne_bytes(bs)
- }
- fn from_ne_bytes(bs: Self::Buffer) -> Self {
- match bs[0] {
- 0 => false,
- 1 => true,
- _ => panic!("Invalid byte when reading bool"),
- }
- }
-}
-
-from_le_bytes! { u8, u16, u32, u64, i8, i16, i32, i64, f32, f64 }
-
-/// Reads `$size` of bytes from `$src`, and reinterprets them as type `$ty`, in
-/// little-endian order. `$ty` must implement the `Default` trait. Otherwise this won't
-/// compile.
-/// This is copied and modified from byteorder crate.
-macro_rules! read_num_bytes {
- ($ty:ty, $size:expr, $src:expr) => {{
- assert!($size <= $src.len());
- let mut buffer = <$ty as $crate::util::bit_util::FromBytes>::Buffer::default();
- buffer.as_mut()[..$size].copy_from_slice(&$src[..$size]);
- <$ty>::from_ne_bytes(buffer)
- }};
-}
-
-/// Converts value `val` of type `T` to a byte vector, by reading `num_bytes` from `val`.
-/// NOTE: if `val` is less than the size of `T` then it can be truncated.
-#[inline]
-pub fn convert_to_bytes<T>(val: &T, num_bytes: usize) -> Vec<u8>
-where
- T: ?Sized + AsBytes,
-{
- let mut bytes: Vec<u8> = vec![0; num_bytes];
- memcpy_value(val.as_bytes(), num_bytes, &mut bytes);
- bytes
-}
-
-#[inline]
-pub fn memcpy(source: &[u8], target: &mut [u8]) {
- assert!(target.len() >= source.len());
- target[..source.len()].copy_from_slice(source)
-}
-
-#[inline]
-pub fn memcpy_value<T>(source: &T, num_bytes: usize, target: &mut [u8])
-where
- T: ?Sized + AsBytes,
-{
- assert!(
- target.len() >= num_bytes,
- "Not enough space. Only had {} bytes but need to put {} bytes",
- target.len(),
- num_bytes
- );
- memcpy(&source.as_bytes()[..num_bytes], target)
-}
-
-/// Returns the ceil of value/divisor
-#[inline]
-pub fn ceil(value: i64, divisor: i64) -> i64 {
- value / divisor + ((value % divisor != 0) as i64)
-}
-
-/// Returns ceil(log2(x))
-#[inline]
-pub fn log2(mut x: u64) -> i32 {
- if x == 1 {
- return 0;
- }
- x -= 1;
- let mut result = 0;
- while x > 0 {
- x >>= 1;
- result += 1;
- }
- result
-}
-
-/// Returns the `num_bits` least-significant bits of `v`
-#[inline]
-pub fn trailing_bits(v: u64, num_bits: usize) -> u64 {
- if num_bits == 0 {
- return 0;
- }
- if num_bits >= 64 {
- return v;
- }
- let n = 64 - num_bits;
- (v << n) >> n
-}
-
-#[inline]
-pub fn set_array_bit(bits: &mut [u8], i: usize) {
- bits[i / 8] |= 1 << (i % 8);
-}
-
-#[inline]
-pub fn unset_array_bit(bits: &mut [u8], i: usize) {
- bits[i / 8] &= !(1 << (i % 8));
-}
-
-/// Returns the minimum number of bits needed to represent the value 'x'
-#[inline]
-pub fn num_required_bits(x: u64) -> usize {
- for i in (0..64).rev() {
- if x & (1u64 << i) != 0 {
- return i + 1;
- }
- }
- 0
-}
-
-static BIT_MASK: [u8; 8] = [1, 2, 4, 8, 16, 32, 64, 128];
-
-/// Returns whether bit at position `i` in `data` is set or not
-#[inline]
-pub fn get_bit(data: &[u8], i: usize) -> bool {
- (data[i >> 3] & BIT_MASK[i & 7]) != 0
-}
-
-/// Utility class for writing bit/byte streams. This class can write data in either
-/// bit packed or byte aligned fashion.
-pub struct BitWriter {
- buffer: Vec<u8>,
- max_bytes: usize,
- buffered_values: u64,
- byte_offset: usize,
- bit_offset: usize,
- start: usize,
-}
-
-impl BitWriter {
- pub fn new(max_bytes: usize) -> Self {
- Self {
- buffer: vec![0; max_bytes],
- max_bytes,
- buffered_values: 0,
- byte_offset: 0,
- bit_offset: 0,
- start: 0,
- }
- }
-
- /// Initializes the writer from the existing buffer `buffer` and starting
- /// offset `start`.
- pub fn new_from_buf(buffer: Vec<u8>, start: usize) -> Self {
- assert!(start < buffer.len());
- let len = buffer.len();
- Self {
- buffer,
- max_bytes: len,
- buffered_values: 0,
- byte_offset: start,
- bit_offset: 0,
- start,
- }
- }
-
- /// Extend buffer size
- #[inline]
- pub fn extend(&mut self, increment: usize) {
- self.max_bytes += increment;
- let extra = vec![0; increment];
- self.buffer.extend(extra);
- }
-
- /// Report buffer size
- #[inline]
- pub fn capacity(&mut self) -> usize {
- self.max_bytes
- }
-
- /// Consumes and returns the current buffer.
- #[inline]
- pub fn consume(mut self) -> Vec<u8> {
- self.flush();
- self.buffer.truncate(self.byte_offset);
- self.buffer
- }
-
- /// Flushes the internal buffered bits and returns the buffer's content.
- /// This is a borrow equivalent of `consume` method.
- #[inline]
- pub fn flush_buffer(&mut self) -> &[u8] {
- self.flush();
- &self.buffer()[0..self.byte_offset]
- }
-
- /// Clears the internal state so the buffer can be reused.
- #[inline]
- pub fn clear(&mut self) {
- self.buffered_values = 0;
- self.byte_offset = self.start;
- self.bit_offset = 0;
- }
-
- /// Flushes the internal buffered bits and the align the buffer to the next byte.
- #[inline]
- pub fn flush(&mut self) {
- let num_bytes = ceil(self.bit_offset as i64, 8) as usize;
- assert!(self.byte_offset + num_bytes <= self.max_bytes);
- memcpy_value(
- &self.buffered_values,
- num_bytes,
- &mut self.buffer[self.byte_offset..],
- );
- self.buffered_values = 0;
- self.bit_offset = 0;
- self.byte_offset += num_bytes;
- }
-
- /// Advances the current offset by skipping `num_bytes`, flushing the internal bit
- /// buffer first.
- /// This is useful when you want to jump over `num_bytes` bytes and come back later
- /// to fill these bytes.
- ///
- /// Returns error if `num_bytes` is beyond the boundary of the internal buffer.
- /// Otherwise, returns the old offset.
- #[inline]
- pub fn skip(&mut self, num_bytes: usize) -> Result<usize> {
- self.flush();
- assert!(self.byte_offset <= self.max_bytes);
- if self.byte_offset + num_bytes > self.max_bytes {
- return Err(general_err!(
- "Not enough bytes left in BitWriter. Need {} but only have {}",
- self.byte_offset + num_bytes,
- self.max_bytes
- ));
- }
- let result = self.byte_offset;
- self.byte_offset += num_bytes;
- Ok(result)
- }
-
- /// Returns a slice containing the next `num_bytes` bytes starting from the current
- /// offset, and advances the underlying buffer by `num_bytes`.
- /// This is useful when you want to jump over `num_bytes` bytes and come back later
- /// to fill these bytes.
- #[inline]
- pub fn get_next_byte_ptr(&mut self, num_bytes: usize) -> Result<&mut [u8]> {
- let offset = self.skip(num_bytes)?;
- Ok(&mut self.buffer[offset..offset + num_bytes])
- }
-
- #[inline]
- pub fn bytes_written(&self) -> usize {
- self.byte_offset - self.start + ceil(self.bit_offset as i64, 8) as usize
- }
-
- #[inline]
- pub fn buffer(&self) -> &[u8] {
- &self.buffer[self.start..]
- }
-
- #[inline]
- pub fn byte_offset(&self) -> usize {
- self.byte_offset
- }
-
- /// Returns the internal buffer length. This is the maximum number of bytes that this
- /// writer can write. User needs to call `consume` to consume the current buffer
- /// before more data can be written.
- #[inline]
- pub fn buffer_len(&self) -> usize {
- self.max_bytes
- }
-
- pub fn write_at(&mut self, offset: usize, value: u8) {
- self.buffer[offset] = value;
- }
-
- /// Writes the `num_bits` LSB of value `v` to the internal buffer of this writer.
- /// The `num_bits` must not be greater than 64. This is bit packed.
- ///
- /// Returns false if there's not enough room left. True otherwise.
- #[inline]
- pub fn put_value(&mut self, v: u64, num_bits: usize) -> bool {
- assert!(num_bits <= 64);
- assert_eq!(v.checked_shr(num_bits as u32).unwrap_or(0), 0); // covers case v >> 64
-
- if self.byte_offset * 8 + self.bit_offset + num_bits > self.max_bytes as usize * 8
- {
- return false;
- }
-
- self.buffered_values |= v << self.bit_offset;
- self.bit_offset += num_bits;
- if self.bit_offset >= 64 {
- memcpy_value(
- &self.buffered_values,
- 8,
- &mut self.buffer[self.byte_offset..],
- );
- self.byte_offset += 8;
- self.bit_offset -= 64;
- self.buffered_values = 0;
- // Perform checked right shift: v >> offset, where offset < 64, otherwise we
- // shift all bits
- self.buffered_values = v
- .checked_shr((num_bits - self.bit_offset) as u32)
- .unwrap_or(0);
- }
- assert!(self.bit_offset < 64);
- true
- }
-
- /// Writes `val` of `num_bytes` bytes to the next aligned byte. If size of `T` is
- /// larger than `num_bytes`, extra higher ordered bytes will be ignored.
- ///
- /// Returns false if there's not enough room left. True otherwise.
- #[inline]
- pub fn put_aligned<T: AsBytes>(&mut self, val: T, num_bytes: usize) -> bool {
- let result = self.get_next_byte_ptr(num_bytes);
- if result.is_err() {
- // TODO: should we return `Result` for this func?
- return false;
- }
- let mut ptr = result.unwrap();
- memcpy_value(&val, num_bytes, &mut ptr);
- true
- }
-
- /// Writes `val` of `num_bytes` bytes at the designated `offset`. The `offset` is the
- /// offset starting from the beginning of the internal buffer that this writer
- /// maintains. Note that this will overwrite any existing data between `offset` and
- /// `offset + num_bytes`. Also that if size of `T` is larger than `num_bytes`, extra
- /// higher ordered bytes will be ignored.
- ///
- /// Returns false if there's not enough room left, or the `pos` is not valid.
- /// True otherwise.
- #[inline]
- pub fn put_aligned_offset<T: AsBytes>(
- &mut self,
- val: T,
- num_bytes: usize,
- offset: usize,
- ) -> bool {
- if num_bytes + offset > self.max_bytes {
- return false;
- }
- memcpy_value(
- &val,
- num_bytes,
- &mut self.buffer[offset..offset + num_bytes],
- );
- true
- }
-
- /// Writes a VLQ encoded integer `v` to this buffer. The value is byte aligned.
- ///
- /// Returns false if there's not enough room left. True otherwise.
- #[inline]
- pub fn put_vlq_int(&mut self, mut v: u64) -> bool {
- let mut result = true;
- while v & 0xFFFFFFFFFFFFFF80 != 0 {
- result &= self.put_aligned::<u8>(((v & 0x7F) | 0x80) as u8, 1);
- v >>= 7;
- }
- result &= self.put_aligned::<u8>((v & 0x7F) as u8, 1);
- result
- }
-
- /// Writes a zigzag-VLQ encoded (in little endian order) int `v` to this buffer.
- /// Zigzag-VLQ is a variant of VLQ encoding where negative and positive
- /// numbers are encoded in a zigzag fashion.
- /// See: https://developers.google.com/protocol-buffers/docs/encoding
- ///
- /// Returns false if there's not enough room left. True otherwise.
- #[inline]
- pub fn put_zigzag_vlq_int(&mut self, v: i64) -> bool {
- let u: u64 = ((v << 1) ^ (v >> 63)) as u64;
- self.put_vlq_int(u)
- }
-}
-
-/// Maximum byte length for a VLQ encoded integer
-/// MAX_VLQ_BYTE_LEN = 5 for i32, and MAX_VLQ_BYTE_LEN = 10 for i64
-pub const MAX_VLQ_BYTE_LEN: usize = 10;
-
-pub struct BitReader {
- // The byte buffer to read from, passed in by client
- buffer: ByteBufferPtr,
-
- // Bytes are memcpy'd from `buffer` and values are read from this variable.
- // This is faster than reading values byte by byte directly from `buffer`
- buffered_values: u64,
-
- //
- // End Start
- // |............|B|B|B|B|B|B|B|B|..............|
- // ^ ^
- // bit_offset byte_offset
- //
- // Current byte offset in `buffer`
- byte_offset: usize,
-
- // Current bit offset in `buffered_values`
- bit_offset: usize,
-
- // Total number of bytes in `buffer`
- total_bytes: usize,
-}
-
-/// Utility class to read bit/byte stream. This class can read bits or bytes that are
-/// either byte aligned or not.
-impl BitReader {
- pub fn new(buffer: ByteBufferPtr) -> Self {
- let total_bytes = buffer.len();
- let num_bytes = cmp::min(8, total_bytes);
- let buffered_values = read_num_bytes!(u64, num_bytes, buffer.as_ref());
- BitReader {
- buffer,
- buffered_values,
- byte_offset: 0,
- bit_offset: 0,
- total_bytes,
- }
- }
-
- pub fn reset(&mut self, buffer: ByteBufferPtr) {
- self.buffer = buffer;
- self.total_bytes = self.buffer.len();
- let num_bytes = cmp::min(8, self.total_bytes);
- self.buffered_values = read_num_bytes!(u64, num_bytes, self.buffer.as_ref());
- self.byte_offset = 0;
- self.bit_offset = 0;
- }
-
- /// Gets the current byte offset
- #[inline]
- pub fn get_byte_offset(&self) -> usize {
- self.byte_offset + ceil(self.bit_offset as i64, 8) as usize
- }
-
- /// Reads a value of type `T` and of size `num_bits`.
- ///
- /// Returns `None` if there's not enough data available. `Some` otherwise.
- pub fn get_value<T: FromBytes>(&mut self, num_bits: usize) -> Option<T> {
- assert!(num_bits <= 64);
- assert!(num_bits <= size_of::<T>() * 8);
-
- if self.byte_offset * 8 + self.bit_offset + num_bits > self.total_bytes * 8 {
- return None;
- }
-
- let mut v = trailing_bits(self.buffered_values, self.bit_offset + num_bits)
- >> self.bit_offset;
- self.bit_offset += num_bits;
-
- if self.bit_offset >= 64 {
- self.byte_offset += 8;
- self.bit_offset -= 64;
-
- self.reload_buffer_values();
- v |= trailing_bits(self.buffered_values, self.bit_offset)
- .wrapping_shl((num_bits - self.bit_offset) as u32);
- }
-
- // TODO: better to avoid copying here
- Some(from_ne_slice(v.as_bytes()))
- }
-
- pub fn get_batch<T: FromBytes>(&mut self, batch: &mut [T], num_bits: usize) -> usize {
- assert!(num_bits <= 32);
- assert!(num_bits <= size_of::<T>() * 8);
-
- let mut values_to_read = batch.len();
- let needed_bits = num_bits * values_to_read;
- let remaining_bits = (self.total_bytes - self.byte_offset) * 8 - self.bit_offset;
- if remaining_bits < needed_bits {
- values_to_read = remaining_bits / num_bits;
- }
-
- let mut i = 0;
-
- // First align bit offset to byte offset
- if self.bit_offset != 0 {
- while i < values_to_read && self.bit_offset != 0 {
- batch[i] = self
- .get_value(num_bits)
- .expect("expected to have more data");
- i += 1;
- }
- }
-
- unsafe {
- let in_buf = &self.buffer.data()[self.byte_offset..];
- let mut in_ptr = in_buf as *const [u8] as *const u8 as *const u32;
- // FIXME assert!(memory::is_ptr_aligned(in_ptr));
- if size_of::<T>() == 4 {
- while values_to_read - i >= 32 {
- let out_ptr = &mut batch[i..] as *mut [T] as *mut T as *mut u32;
- in_ptr = unpack32(in_ptr, out_ptr, num_bits);
- self.byte_offset += 4 * num_bits;
- i += 32;
- }
- } else {
- let mut out_buf = [0u32; 32];
- let out_ptr = &mut out_buf as &mut [u32] as *mut [u32] as *mut u32;
- while values_to_read - i >= 32 {
- in_ptr = unpack32(in_ptr, out_ptr, num_bits);
- self.byte_offset += 4 * num_bits;
- for n in 0..32 {
- // We need to copy from smaller size to bigger size to avoid
- // overwriting other memory regions.
- if size_of::<T>() > size_of::<u32>() {
- std::ptr::copy_nonoverlapping(
- out_buf[n..].as_ptr() as *const u32,
- &mut batch[i] as *mut T as *mut u32,
- 1,
- );
- } else {
- std::ptr::copy_nonoverlapping(
- out_buf[n..].as_ptr() as *const T,
- &mut batch[i] as *mut T,
- 1,
- );
- }
- i += 1;
- }
- }
- }
- }
-
- assert!(values_to_read - i < 32);
-
- self.reload_buffer_values();
- while i < values_to_read {
- batch[i] = self
- .get_value(num_bits)
- .expect("expected to have more data");
- i += 1;
- }
-
- values_to_read
- }
-
- /// Reads a `num_bytes`-sized value from this buffer and return it.
- /// `T` needs to be a little-endian native type. The value is assumed to be byte
- /// aligned so the bit reader will be advanced to the start of the next byte before
- /// reading the value.
-
- /// Returns `Some` if there's enough bytes left to form a value of `T`.
- /// Otherwise `None`.
- pub fn get_aligned<T: FromBytes>(&mut self, num_bytes: usize) -> Option<T> {
- let bytes_read = ceil(self.bit_offset as i64, 8) as usize;
- if self.byte_offset + bytes_read + num_bytes > self.total_bytes {
- return None;
- }
-
- // Advance byte_offset to next unread byte and read num_bytes
- self.byte_offset += bytes_read;
- let v = read_num_bytes!(T, num_bytes, self.buffer.data()[self.byte_offset..]);
- self.byte_offset += num_bytes;
-
- // Reset buffered_values
- self.bit_offset = 0;
- self.reload_buffer_values();
- Some(v)
- }
-
- /// Reads a VLQ encoded (in little endian order) int from the stream.
- /// The encoded int must start at the beginning of a byte.
- ///
- /// Returns `None` if there's not enough bytes in the stream. `Some` otherwise.
- pub fn get_vlq_int(&mut self) -> Option<i64> {
- let mut shift = 0;
- let mut v: i64 = 0;
- while let Some(byte) = self.get_aligned::<u8>(1) {
- v |= ((byte & 0x7F) as i64) << shift;
- shift += 7;
- assert!(
- shift <= MAX_VLQ_BYTE_LEN * 7,
- "Num of bytes exceed MAX_VLQ_BYTE_LEN ({})",
- MAX_VLQ_BYTE_LEN
- );
- if byte & 0x80 == 0 {
- return Some(v);
- }
- }
- None
- }
-
- /// Reads a zigzag-VLQ encoded (in little endian order) int from the stream
- /// Zigzag-VLQ is a variant of VLQ encoding where negative and positive numbers are
- /// encoded in a zigzag fashion.
- /// See: https://developers.google.com/protocol-buffers/docs/encoding
- ///
- /// Note: the encoded int must start at the beginning of a byte.
- ///
- /// Returns `None` if the number of bytes there's not enough bytes in the stream.
- /// `Some` otherwise.
- #[inline]
- pub fn get_zigzag_vlq_int(&mut self) -> Option<i64> {
- self.get_vlq_int().map(|v| {
- let u = v as u64;
- (u >> 1) as i64 ^ -((u & 1) as i64)
- })
- }
-
- fn reload_buffer_values(&mut self) {
- let bytes_to_read = cmp::min(self.total_bytes - self.byte_offset, 8);
- self.buffered_values =
- read_num_bytes!(u64, bytes_to_read, self.buffer.data()[self.byte_offset..]);
- }
-}
-
-impl From<Vec<u8>> for BitReader {
- #[inline]
- fn from(buffer: Vec<u8>) -> Self {
- BitReader::new(ByteBufferPtr::new(buffer))
- }
-}
-
-#[cfg(test)]
-mod tests {
- use super::super::test_common::*;
- use super::*;
-
- use rand::distributions::{Distribution, Standard};
- use std::fmt::Debug;
-
- #[test]
- fn test_ceil() {
- assert_eq!(ceil(0, 1), 0);
- assert_eq!(ceil(1, 1), 1);
- assert_eq!(ceil(1, 2), 1);
- assert_eq!(ceil(1, 8), 1);
- assert_eq!(ceil(7, 8), 1);
- assert_eq!(ceil(8, 8), 1);
- assert_eq!(ceil(9, 8), 2);
- assert_eq!(ceil(9, 9), 1);
- assert_eq!(ceil(10000000000, 10), 1000000000);
- assert_eq!(ceil(10, 10000000000), 1);
- assert_eq!(ceil(10000000000, 1000000000), 10);
- }
-
- #[test]
- fn test_bit_reader_get_byte_offset() {
- let buffer = vec![255; 10];
- let mut bit_reader = BitReader::from(buffer);
- assert_eq!(bit_reader.get_byte_offset(), 0); // offset (0 bytes, 0 bits)
- bit_reader.get_value::<i32>(6);
- assert_eq!(bit_reader.get_byte_offset(), 1); // offset (0 bytes, 6 bits)
- bit_reader.get_value::<i32>(10);
- assert_eq!(bit_reader.get_byte_offset(), 2); // offset (0 bytes, 16 bits)
- bit_reader.get_value::<i32>(20);
- assert_eq!(bit_reader.get_byte_offset(), 5); // offset (0 bytes, 36 bits)
- bit_reader.get_value::<i32>(30);
- assert_eq!(bit_reader.get_byte_offset(), 9); // offset (8 bytes, 2 bits)
- }
-
- #[test]
- fn test_bit_reader_get_value() {
- let buffer = vec![255, 0];
- let mut bit_reader = BitReader::from(buffer);
- assert_eq!(bit_reader.get_value::<i32>(1), Some(1));
- assert_eq!(bit_reader.get_value::<i32>(2), Some(3));
- assert_eq!(bit_reader.get_value::<i32>(3), Some(7));
- assert_eq!(bit_reader.get_value::<i32>(4), Some(3));
- }
-
- #[test]
- fn test_bit_reader_get_value_boundary() {
- let buffer = vec![10, 0, 0, 0, 20, 0, 30, 0, 0, 0, 40, 0];
- let mut bit_reader = BitReader::from(buffer);
- assert_eq!(bit_reader.get_value::<i64>(32), Some(10));
- assert_eq!(bit_reader.get_value::<i64>(16), Some(20));
- assert_eq!(bit_reader.get_value::<i64>(32), Some(30));
- assert_eq!(bit_reader.get_value::<i64>(16), Some(40));
- }
-
- #[test]
- fn test_bit_reader_get_aligned() {
- // 01110101 11001011
- let buffer = ByteBufferPtr::new(vec![0x75, 0xCB]);
- let mut bit_reader = BitReader::new(buffer.all());
- assert_eq!(bit_reader.get_value::<i32>(3), Some(5));
- assert_eq!(bit_reader.get_aligned::<i32>(1), Some(203));
- assert_eq!(bit_reader.get_value::<i32>(1), None);
- bit_reader.reset(buffer.all());
- assert_eq!(bit_reader.get_aligned::<i32>(3), None);
- }
-
- #[test]
- fn test_bit_reader_get_vlq_int() {
- // 10001001 00000001 11110010 10110101 00000110
- let buffer: Vec<u8> = vec![0x89, 0x01, 0xF2, 0xB5, 0x06];
- let mut bit_reader = BitReader::from(buffer);
- assert_eq!(bit_reader.get_vlq_int(), Some(137));
- assert_eq!(bit_reader.get_vlq_int(), Some(105202));
- }
-
- #[test]
- fn test_bit_reader_get_zigzag_vlq_int() {
- let buffer: Vec<u8> = vec![0, 1, 2, 3];
- let mut bit_reader = BitReader::from(buffer);
- assert_eq!(bit_reader.get_zigzag_vlq_int(), Some(0));
- assert_eq!(bit_reader.get_zigzag_vlq_int(), Some(-1));
- assert_eq!(bit_reader.get_zigzag_vlq_int(), Some(1));
- assert_eq!(bit_reader.get_zigzag_vlq_int(), Some(-2));
- }
-
- #[test]
- fn test_set_array_bit() {
- let mut buffer = vec![0, 0, 0];
- set_array_bit(&mut buffer[..], 1);
- assert_eq!(buffer, vec![2, 0, 0]);
- set_array_bit(&mut buffer[..], 4);
- assert_eq!(buffer, vec![18, 0, 0]);
- unset_array_bit(&mut buffer[..], 1);
- assert_eq!(buffer, vec![16, 0, 0]);
- set_array_bit(&mut buffer[..], 10);
- assert_eq!(buffer, vec![16, 4, 0]);
- set_array_bit(&mut buffer[..], 10);
- assert_eq!(buffer, vec![16, 4, 0]);
- set_array_bit(&mut buffer[..], 11);
- assert_eq!(buffer, vec![16, 12, 0]);
- unset_array_bit(&mut buffer[..], 10);
- assert_eq!(buffer, vec![16, 8, 0]);
- }
-
- #[test]
- fn test_num_required_bits() {
- assert_eq!(num_required_bits(0), 0);
- assert_eq!(num_required_bits(1), 1);
- assert_eq!(num_required_bits(2), 2);
- assert_eq!(num_required_bits(4), 3);
- assert_eq!(num_required_bits(8), 4);
- assert_eq!(num_required_bits(10), 4);
- assert_eq!(num_required_bits(12), 4);
- assert_eq!(num_required_bits(16), 5);
- }
-
- #[test]
- fn test_get_bit() {
- // 00001101
- assert_eq!(true, get_bit(&[0b00001101], 0));
- assert_eq!(false, get_bit(&[0b00001101], 1));
- assert_eq!(true, get_bit(&[0b00001101], 2));
- assert_eq!(true, get_bit(&[0b00001101], 3));
-
- // 01001001 01010010
- assert_eq!(true, get_bit(&[0b01001001, 0b01010010], 0));
- assert_eq!(false, get_bit(&[0b01001001, 0b01010010], 1));
- assert_eq!(false, get_bit(&[0b01001001, 0b01010010], 2));
- assert_eq!(true, get_bit(&[0b01001001, 0b01010010], 3));
- assert_eq!(false, get_bit(&[0b01001001, 0b01010010], 4));
- assert_eq!(false, get_bit(&[0b01001001, 0b01010010], 5));
- assert_eq!(true, get_bit(&[0b01001001, 0b01010010], 6));
- assert_eq!(false, get_bit(&[0b01001001, 0b01010010], 7));
- assert_eq!(false, get_bit(&[0b01001001, 0b01010010], 8));
- assert_eq!(true, get_bit(&[0b01001001, 0b01010010], 9));
- assert_eq!(false, get_bit(&[0b01001001, 0b01010010], 10));
- assert_eq!(false, get_bit(&[0b01001001, 0b01010010], 11));
- assert_eq!(true, get_bit(&[0b01001001, 0b01010010], 12));
- assert_eq!(false, get_bit(&[0b01001001, 0b01010010], 13));
- assert_eq!(true, get_bit(&[0b01001001, 0b01010010], 14));
- assert_eq!(false, get_bit(&[0b01001001, 0b01010010], 15));
- }
-
- #[test]
- fn test_log2() {
- assert_eq!(log2(1), 0);
- assert_eq!(log2(2), 1);
- assert_eq!(log2(3), 2);
- assert_eq!(log2(4), 2);
- assert_eq!(log2(5), 3);
- assert_eq!(log2(5), 3);
- assert_eq!(log2(6), 3);
- assert_eq!(log2(7), 3);
- assert_eq!(log2(8), 3);
- assert_eq!(log2(9), 4);
- }
-
- #[test]
- fn test_skip() {
- let mut writer = BitWriter::new(5);
- let old_offset = writer.skip(1).expect("skip() should return OK");
- writer.put_aligned(42, 4);
- writer.put_aligned_offset(0x10, 1, old_offset);
- let result = writer.consume();
- assert_eq!(result.as_ref(), [0x10, 42, 0, 0, 0]);
-
- writer = BitWriter::new(4);
- let result = writer.skip(5);
- assert!(result.is_err());
- }
-
- #[test]
- fn test_get_next_byte_ptr() {
- let mut writer = BitWriter::new(5);
- {
- let first_byte = writer
- .get_next_byte_ptr(1)
- .expect("get_next_byte_ptr() should return OK");
- first_byte[0] = 0x10;
- }
- writer.put_aligned(42, 4);
- let result = writer.consume();
- assert_eq!(result.as_ref(), [0x10, 42, 0, 0, 0]);
- }
-
- #[test]
- fn test_consume_flush_buffer() {
- let mut writer1 = BitWriter::new(3);
- let mut writer2 = BitWriter::new(3);
- for i in 1..10 {
- writer1.put_value(i, 4);
- writer2.put_value(i, 4);
- }
- let res1 = writer1.flush_buffer();
- let res2 = writer2.consume();
- assert_eq!(res1, &res2[..]);
- }
-
- #[test]
- fn test_put_get_bool() {
- let len = 8;
- let mut writer = BitWriter::new(len);
-
- for i in 0..8 {
- let result = writer.put_value(i % 2, 1);
- assert!(result);
- }
-
- writer.flush();
- {
- let buffer = writer.buffer();
- assert_eq!(buffer[0], 0b10101010);
- }
-
- // Write 00110011
- for i in 0..8 {
- let result = match i {
- 0 | 1 | 4 | 5 => writer.put_value(false as u64, 1),
- _ => writer.put_value(true as u64, 1),
- };
- assert!(result);
- }
- writer.flush();
- {
- let buffer = writer.buffer();
- assert_eq!(buffer[0], 0b10101010);
- assert_eq!(buffer[1], 0b11001100);
- }
-
- let mut reader = BitReader::from(writer.consume());
-
- for i in 0..8 {
- let val = reader
- .get_value::<u8>(1)
- .expect("get_value() should return OK");
- assert_eq!(val, i % 2);
- }
-
- for i in 0..8 {
- let val = reader
- .get_value::<bool>(1)
- .expect("get_value() should return OK");
- match i {
- 0 | 1 | 4 | 5 => assert_eq!(val, false),
- _ => assert_eq!(val, true),
- }
- }
- }
-
- #[test]
- fn test_put_value_roundtrip() {
- test_put_value_rand_numbers(32, 2);
- test_put_value_rand_numbers(32, 3);
- test_put_value_rand_numbers(32, 4);
- test_put_value_rand_numbers(32, 5);
- test_put_value_rand_numbers(32, 6);
- test_put_value_rand_numbers(32, 7);
- test_put_value_rand_numbers(32, 8);
- test_put_value_rand_numbers(64, 16);
- test_put_value_rand_numbers(64, 24);
- test_put_value_rand_numbers(64, 32);
- }
-
- fn test_put_value_rand_numbers(total: usize, num_bits: usize) {
- assert!(num_bits < 64);
- let num_bytes = ceil(num_bits as i64, 8);
- let mut writer = BitWriter::new(num_bytes as usize * total);
- let values: Vec<u64> = random_numbers::<u64>(total)
- .iter()
- .map(|v| v & ((1 << num_bits) - 1))
- .collect();
- (0..total).for_each(|i| {
- assert!(
- writer.put_value(values[i] as u64, num_bits),
- "[{}]: put_value() failed",
- i
- );
- });
-
- let mut reader = BitReader::from(writer.consume());
- (0..total).for_each(|i| {
- let v = reader
- .get_value::<u64>(num_bits)
- .expect("get_value() should return OK");
- assert_eq!(
- v, values[i],
- "[{}]: expected {} but got {}",
- i, values[i], v
- );
- });
- }
-
- #[test]
- fn test_get_batch() {
- const SIZE: &[usize] = &[1, 31, 32, 33, 128, 129];
- for s in SIZE {
- for i in 0..33 {
- match i {
- 0..=8 => test_get_batch_helper::<u8>(*s, i),
- 9..=16 => test_get_batch_helper::<u16>(*s, i),
- _ => test_get_batch_helper::<u32>(*s, i),
- }
- }
- }
- }
-
- fn test_get_batch_helper<T>(total: usize, num_bits: usize)
- where
- T: FromBytes + Default + Clone + Debug + Eq,
- {
- assert!(num_bits <= 32);
- let num_bytes = ceil(num_bits as i64, 8);
- let mut writer = BitWriter::new(num_bytes as usize * total);
-
- let values: Vec<u32> = random_numbers::<u32>(total)
- .iter()
- .map(|v| v & ((1u64 << num_bits) - 1) as u32)
- .collect();
-
- // Generic values used to check against actual values read from `get_batch`.
- let expected_values: Vec<T> =
- values.iter().map(|v| from_ne_slice(v.as_bytes())).collect();
-
- (0..total).for_each(|i| {
- assert!(writer.put_value(values[i] as u64, num_bits));
- });
-
- let buf = writer.consume();
- let mut reader = BitReader::from(buf);
- let mut batch = vec![T::default(); values.len()];
- let values_read = reader.get_batch::<T>(&mut batch, num_bits);
- assert_eq!(values_read, values.len());
- for i in 0..batch.len() {
- assert_eq!(
- batch[i], expected_values[i],
- "num_bits = {}, index = {}",
- num_bits, i
- );
- }
- }
-
- #[test]
- fn test_put_aligned_roundtrip() {
- test_put_aligned_rand_numbers::<u8>(4, 3);
- test_put_aligned_rand_numbers::<u8>(16, 5);
- test_put_aligned_rand_numbers::<i16>(32, 7);
- test_put_aligned_rand_numbers::<i16>(32, 9);
- test_put_aligned_rand_numbers::<i32>(32, 11);
- test_put_aligned_rand_numbers::<i32>(32, 13);
- test_put_aligned_rand_numbers::<i64>(32, 17);
- test_put_aligned_rand_numbers::<i64>(32, 23);
- }
-
- fn test_put_aligned_rand_numbers<T>(total: usize, num_bits: usize)
- where
- T: Copy + FromBytes + AsBytes + Debug + PartialEq,
- Standard: Distribution<T>,
- {
- assert!(num_bits <= 32);
- assert!(total % 2 == 0);
-
- let aligned_value_byte_width = std::mem::size_of::<T>();
- let value_byte_width = ceil(num_bits as i64, 8) as usize;
- let mut writer =
- BitWriter::new((total / 2) * (aligned_value_byte_width + value_byte_width));
- let values: Vec<u32> = random_numbers::<u32>(total / 2)
- .iter()
- .map(|v| v & ((1 << num_bits) - 1))
- .collect();
- let aligned_values = random_numbers::<T>(total / 2);
-
- for i in 0..total {
- let j = i / 2;
- if i % 2 == 0 {
- assert!(
- writer.put_value(values[j] as u64, num_bits),
- "[{}]: put_value() failed",
- i
- );
- } else {
- assert!(
- writer.put_aligned::<T>(aligned_values[j], aligned_value_byte_width),
- "[{}]: put_aligned() failed",
- i
- );
- }
- }
-
- let mut reader = BitReader::from(writer.consume());
- for i in 0..total {
- let j = i / 2;
- if i % 2 == 0 {
- let v = reader
- .get_value::<u64>(num_bits)
- .expect("get_value() should return OK");
- assert_eq!(
- v, values[j] as u64,
- "[{}]: expected {} but got {}",
- i, values[j], v
- );
- } else {
- let v = reader
- .get_aligned::<T>(aligned_value_byte_width)
- .expect("get_aligned() should return OK");
- assert_eq!(
- v, aligned_values[j],
- "[{}]: expected {:?} but got {:?}",
- i, aligned_values[j], v
- );
- }
- }
- }
-
- #[test]
- fn test_put_vlq_int() {
- let total = 64;
- let mut writer = BitWriter::new(total * 32);
- let values = random_numbers::<u32>(total);
- (0..total).for_each(|i| {
- assert!(
- writer.put_vlq_int(values[i] as u64),
- "[{}]; put_vlq_int() failed",
- i
- );
- });
-
- let mut reader = BitReader::from(writer.consume());
- (0..total).for_each(|i| {
- let v = reader
- .get_vlq_int()
- .expect("get_vlq_int() should return OK");
- assert_eq!(
- v as u32, values[i],
- "[{}]: expected {} but got {}",
- i, values[i], v
- );
- });
- }
-
- #[test]
- fn test_put_zigzag_vlq_int() {
- let total = 64;
- let mut writer = BitWriter::new(total * 32);
- let values = random_numbers::<i32>(total);
- (0..total).for_each(|i| {
- assert!(
- writer.put_zigzag_vlq_int(values[i] as i64),
- "[{}]; put_zigzag_vlq_int() failed",
- i
- );
- });
-
- let mut reader = BitReader::from(writer.consume());
- (0..total).for_each(|i| {
- let v = reader
- .get_zigzag_vlq_int()
- .expect("get_zigzag_vlq_int() should return OK");
- assert_eq!(
- v as i32, values[i],
- "[{}]: expected {} but got {}",
- i, values[i], v
- );
- });
- }
-}
diff --git a/parquet/src/util/cursor.rs b/parquet/src/util/cursor.rs
deleted file mode 100644
index eaed6c7..0000000
--- a/parquet/src/util/cursor.rs
+++ /dev/null
@@ -1,261 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-use std::io::{self, Cursor, Error, ErrorKind, Read, Seek, SeekFrom, Write};
-use std::sync::{Arc, Mutex};
-use std::{cmp, fmt};
-
-use crate::file::writer::TryClone;
-
-/// This is object to use if your file is already in memory.
-/// The sliceable cursor is similar to std::io::Cursor, except that it makes it easy to create "cursor slices".
-/// To achieve this, it uses Arc instead of shared references. Indeed reference fields are painful
-/// because the lack of Generic Associated Type implies that you would require complex lifetime propagation when
-/// returning such a cursor.
-#[allow(clippy::rc_buffer)]
-pub struct SliceableCursor {
- inner: Arc<Vec<u8>>,
- start: u64,
- length: usize,
- pos: u64,
-}
-
-impl fmt::Debug for SliceableCursor {
- fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
- f.debug_struct("SliceableCursor")
- .field("start", &self.start)
- .field("length", &self.length)
- .field("pos", &self.pos)
- .field("inner.len", &self.inner.len())
- .finish()
- }
-}
-
-impl SliceableCursor {
- pub fn new(content: impl Into<Arc<Vec<u8>>>) -> Self {
- let inner = content.into();
- let size = inner.len();
- SliceableCursor {
- inner,
- start: 0,
- pos: 0,
- length: size,
- }
- }
-
- /// Create a slice cursor using the same data as a current one.
- pub fn slice(&self, start: u64, length: usize) -> io::Result<Self> {
- let new_start = self.start + start;
- if new_start >= self.inner.len() as u64
- || new_start as usize + length > self.inner.len()
- {
- return Err(Error::new(ErrorKind::InvalidInput, "out of bound"));
- }
- Ok(SliceableCursor {
- inner: Arc::clone(&self.inner),
- start: new_start,
- pos: new_start,
- length,
- })
- }
-
- fn remaining_slice(&self) -> &[u8] {
- let end = self.start as usize + self.length;
- let offset = cmp::min(self.pos, end as u64) as usize;
- &self.inner[offset..end]
- }
-
- /// Get the length of the current cursor slice
- pub fn len(&self) -> u64 {
- self.length as u64
- }
-
- /// return true if the cursor is empty (self.len() == 0)
- pub fn is_empty(&self) -> bool {
- self.len() == 0
- }
-}
-
-/// Implementation inspired by std::io::Cursor
-impl Read for SliceableCursor {
- fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
- let n = Read::read(&mut self.remaining_slice(), buf)?;
- self.pos += n as u64;
- Ok(n)
- }
-}
-
-impl Seek for SliceableCursor {
- fn seek(&mut self, pos: SeekFrom) -> io::Result<u64> {
- let new_pos = match pos {
- SeekFrom::Start(pos) => pos as i64,
- SeekFrom::End(pos) => self.inner.len() as i64 + pos as i64,
- SeekFrom::Current(pos) => self.pos as i64 + pos as i64,
- };
-
- if new_pos < 0 {
- Err(Error::new(
- ErrorKind::InvalidInput,
- format!(
- "Request out of bounds: cur position {} + seek {:?} < 0: {}",
- self.pos, pos, new_pos
- ),
- ))
- } else if new_pos >= self.inner.len() as i64 {
- Err(Error::new(
- ErrorKind::InvalidInput,
- format!(
- "Request out of bounds: cur position {} + seek {:?} >= length {}: {}",
- self.pos,
- pos,
- self.inner.len(),
- new_pos
- ),
- ))
- } else {
- self.pos = new_pos as u64;
- Ok(self.start)
- }
- }
-}
-
-/// Use this type to write Parquet to memory rather than a file.
-#[derive(Debug, Default, Clone)]
-pub struct InMemoryWriteableCursor {
- buffer: Arc<Mutex<Cursor<Vec<u8>>>>,
-}
-
-impl InMemoryWriteableCursor {
- /// Consume this instance and return the underlying buffer as long as there are no other
- /// references to this instance.
- pub fn into_inner(self) -> Option<Vec<u8>> {
- Arc::try_unwrap(self.buffer)
- .ok()
- .and_then(|mutex| mutex.into_inner().ok())
- .map(|cursor| cursor.into_inner())
- }
-
- /// Returns a clone of the underlying buffer
- pub fn data(&self) -> Vec<u8> {
- let inner = self.buffer.lock().unwrap();
- inner.get_ref().to_vec()
- }
-}
-
-impl TryClone for InMemoryWriteableCursor {
- fn try_clone(&self) -> std::io::Result<Self> {
- Ok(Self {
- buffer: self.buffer.clone(),
- })
- }
-}
-
-impl Write for InMemoryWriteableCursor {
- fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
- let mut inner = self.buffer.lock().unwrap();
- inner.write(buf)
- }
-
- fn flush(&mut self) -> std::io::Result<()> {
- let mut inner = self.buffer.lock().unwrap();
- inner.flush()
- }
-}
-
-impl Seek for InMemoryWriteableCursor {
- fn seek(&mut self, pos: SeekFrom) -> std::io::Result<u64> {
- let mut inner = self.buffer.lock().unwrap();
- inner.seek(pos)
- }
-}
-
-#[cfg(test)]
-mod tests {
- use super::*;
-
- /// Create a SliceableCursor of all u8 values in ascending order
- fn get_u8_range() -> SliceableCursor {
- let data: Vec<u8> = (0u8..=255).collect();
- SliceableCursor::new(data)
- }
-
- /// Reads all the bytes in the slice and checks that it matches the u8 range from start to end_included
- fn check_read_all(mut cursor: SliceableCursor, start: u8, end_included: u8) {
- let mut target = vec![];
- let cursor_res = cursor.read_to_end(&mut target);
- println!("{:?}", cursor_res);
- assert!(!cursor_res.is_err(), "reading error");
- assert_eq!((end_included - start) as usize + 1, cursor_res.unwrap());
- assert_eq!((start..=end_included).collect::<Vec<_>>(), target);
- }
-
- #[test]
- fn read_all_whole() {
- let cursor = get_u8_range();
- check_read_all(cursor, 0, 255);
- }
-
- #[test]
- fn read_all_slice() {
- let cursor = get_u8_range().slice(10, 10).expect("error while slicing");
- check_read_all(cursor, 10, 19);
- }
-
- #[test]
- fn seek_cursor_start() {
- let mut cursor = get_u8_range();
-
- cursor.seek(SeekFrom::Start(5)).unwrap();
- check_read_all(cursor, 5, 255);
- }
-
- #[test]
- fn seek_cursor_current() {
- let mut cursor = get_u8_range();
- cursor.seek(SeekFrom::Start(10)).unwrap();
- cursor.seek(SeekFrom::Current(10)).unwrap();
- check_read_all(cursor, 20, 255);
- }
-
- #[test]
- fn seek_cursor_end() {
- let mut cursor = get_u8_range();
-
- cursor.seek(SeekFrom::End(-10)).unwrap();
- check_read_all(cursor, 246, 255);
- }
-
- #[test]
- fn seek_cursor_error_too_long() {
- let mut cursor = get_u8_range();
- let res = cursor.seek(SeekFrom::Start(1000));
- let actual_error = res.expect_err("expected error").to_string();
- let expected_error =
- "Request out of bounds: cur position 0 + seek Start(1000) >= length 256: 1000";
- assert_eq!(actual_error, expected_error);
- }
-
- #[test]
- fn seek_cursor_error_too_short() {
- let mut cursor = get_u8_range();
- let res = cursor.seek(SeekFrom::End(-1000));
- let actual_error = res.expect_err("expected error").to_string();
- let expected_error =
- "Request out of bounds: cur position 0 + seek End(-1000) < 0: -744";
- assert_eq!(actual_error, expected_error);
- }
-}
diff --git a/parquet/src/util/hash_util.rs b/parquet/src/util/hash_util.rs
deleted file mode 100644
index f7849da..0000000
--- a/parquet/src/util/hash_util.rs
+++ /dev/null
@@ -1,172 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-use crate::data_type::AsBytes;
-
-/// Computes hash value for `data`, with a seed value `seed`.
-/// The data type `T` must implement the `AsBytes` trait.
-pub fn hash<T: AsBytes>(data: &T, seed: u32) -> u32 {
- hash_(data.as_bytes(), seed)
-}
-
-fn hash_(data: &[u8], seed: u32) -> u32 {
- #[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
- unsafe {
- if is_x86_feature_detected!("sse4.2") {
- crc32_hash(data, seed)
- } else {
- murmur_hash2_64a(data, seed as u64) as u32
- }
- }
-
- #[cfg(any(target_arch = "aarch64", target_arch = "arm"))]
- unsafe {
- murmur_hash2_64a(data, seed as u64) as u32
- }
-}
-
-const MURMUR_PRIME: u64 = 0xc6a4a7935bd1e995;
-const MURMUR_R: i32 = 47;
-
-/// Rust implementation of MurmurHash2, 64-bit version for 64-bit platforms
-///
-/// SAFTETY Only safe on platforms which support unaligned loads (like x86_64)
-unsafe fn murmur_hash2_64a(data_bytes: &[u8], seed: u64) -> u64 {
- let len = data_bytes.len();
- let len_64 = (len / 8) * 8;
- let data_bytes_64 = std::slice::from_raw_parts(
- &data_bytes[0..len_64] as *const [u8] as *const u64,
- len / 8,
- );
-
- let mut h = seed ^ (MURMUR_PRIME.wrapping_mul(data_bytes.len() as u64));
- for v in data_bytes_64 {
- let mut k = *v;
- k = k.wrapping_mul(MURMUR_PRIME);
- k ^= k >> MURMUR_R;
- k = k.wrapping_mul(MURMUR_PRIME);
- h ^= k;
- h = h.wrapping_mul(MURMUR_PRIME);
- }
-
- let data2 = &data_bytes[len_64..];
-
- let v = len & 7;
- if v == 7 {
- h ^= (data2[6] as u64) << 48;
- }
- if v >= 6 {
- h ^= (data2[5] as u64) << 40;
- }
- if v >= 5 {
- h ^= (data2[4] as u64) << 32;
- }
- if v >= 4 {
- h ^= (data2[3] as u64) << 24;
- }
- if v >= 3 {
- h ^= (data2[2] as u64) << 16;
- }
- if v >= 2 {
- h ^= (data2[1] as u64) << 8;
- }
- if v >= 1 {
- h ^= data2[0] as u64;
- }
- if v > 0 {
- h = h.wrapping_mul(MURMUR_PRIME);
- }
-
- h ^= h >> MURMUR_R;
- h = h.wrapping_mul(MURMUR_PRIME);
- h ^= h >> MURMUR_R;
- h
-}
-
-/// CRC32 hash implementation using SSE4 instructions. Borrowed from Impala.
-#[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
-#[target_feature(enable = "sse4.2")]
-unsafe fn crc32_hash(bytes: &[u8], seed: u32) -> u32 {
- #[cfg(target_arch = "x86")]
- use std::arch::x86::*;
- #[cfg(target_arch = "x86_64")]
- use std::arch::x86_64::*;
-
- let u32_num_bytes = std::mem::size_of::<u32>();
- let mut num_bytes = bytes.len();
- let num_words = num_bytes / u32_num_bytes;
- num_bytes %= u32_num_bytes;
-
- let bytes_u32: &[u32] = std::slice::from_raw_parts(
- &bytes[0..num_words * u32_num_bytes] as *const [u8] as *const u32,
- num_words,
- );
-
- let mut offset = 0;
- let mut hash = seed;
- while offset < num_words {
- hash = _mm_crc32_u32(hash, bytes_u32[offset]);
- offset += 1;
- }
-
- offset = num_words * u32_num_bytes;
- while offset < num_bytes {
- hash = _mm_crc32_u8(hash, bytes[offset]);
- offset += 1;
- }
-
- // The lower half of the CRC hash has poor uniformity, so swap the halves
- // for anyone who only uses the first several bits of the hash.
- hash = (hash << 16) | (hash >> 16);
- hash
-}
-
-#[cfg(test)]
-mod tests {
- use super::*;
-
- #[test]
- fn test_murmur2_64a() {
- unsafe {
- let result = murmur_hash2_64a(b"hello", 123);
- assert_eq!(result, 2597646618390559622);
-
- let result = murmur_hash2_64a(b"helloworld", 123);
- assert_eq!(result, 4934371746140206573);
-
- let result = murmur_hash2_64a(b"helloworldparquet", 123);
- assert_eq!(result, 2392198230801491746);
- }
- }
-
- #[test]
- #[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
- fn test_crc32() {
- if is_x86_feature_detected!("sse4.2") {
- unsafe {
- let result = crc32_hash(b"hello", 123);
- assert_eq!(result, 2927487359);
-
- let result = crc32_hash(b"helloworld", 123);
- assert_eq!(result, 314229527);
-
- let result = crc32_hash(b"helloworldparquet", 123);
- assert_eq!(result, 667078870);
- }
- }
- }
-}
diff --git a/parquet/src/util/io.rs b/parquet/src/util/io.rs
deleted file mode 100644
index 44e99ac..0000000
--- a/parquet/src/util/io.rs
+++ /dev/null
@@ -1,329 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-use std::{cell::RefCell, cmp, fmt, io::*};
-
-use crate::file::{reader::Length, writer::ParquetWriter};
-
-const DEFAULT_BUF_SIZE: usize = 8 * 1024;
-
-// ----------------------------------------------------------------------
-
-/// TryClone tries to clone the type and should maintain the `Seek` position of the given
-/// instance.
-pub trait TryClone: Sized {
- /// Clones the type returning a new instance or an error if it's not possible
- /// to clone it.
- fn try_clone(&self) -> Result<Self>;
-}
-
-/// ParquetReader is the interface which needs to be fulfilled to be able to parse a
-/// parquet source.
-pub trait ParquetReader: Read + Seek + Length + TryClone {}
-impl<T: Read + Seek + Length + TryClone> ParquetReader for T {}
-
-// Read/Write wrappers for `File`.
-
-/// Position trait returns the current position in the stream.
-/// Should be viewed as a lighter version of `Seek` that does not allow seek operations,
-/// and does not require mutable reference for the current position.
-pub trait Position {
- /// Returns position in the stream.
- fn pos(&self) -> u64;
-}
-
-/// Struct that represents a slice of a file data with independent start position and
-/// length. Internally clones provided file handle, wraps with a custom implementation
-/// of BufReader that resets position before any read.
-///
-/// This is workaround and alternative for `file.try_clone()` method. It clones `File`
-/// while preserving independent position, which is not available with `try_clone()`.
-///
-/// Designed after `arrow::io::RandomAccessFile` and `std::io::BufReader`
-pub struct FileSource<R: ParquetReader> {
- reader: RefCell<R>,
- start: u64, // start position in a file
- end: u64, // end position in a file
- buf: Vec<u8>, // buffer where bytes read in advance are stored
- buf_pos: usize, // current position of the reader in the buffer
- buf_cap: usize, // current number of bytes read into the buffer
-}
-
-impl<R: ParquetReader> fmt::Debug for FileSource<R> {
- fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
- f.debug_struct("FileSource")
- .field("reader", &"OPAQUE")
- .field("start", &self.start)
- .field("end", &self.end)
- .field("buf.len", &self.buf.len())
- .field("buf_pos", &self.buf_pos)
- .field("buf_cap", &self.buf_cap)
- .finish()
- }
-}
-
-impl<R: ParquetReader> FileSource<R> {
- /// Creates new file reader with start and length from a file handle
- pub fn new(fd: &R, start: u64, length: usize) -> Self {
- let reader = RefCell::new(fd.try_clone().unwrap());
- Self {
- reader,
- start,
- end: start + length as u64,
- buf: vec![0_u8; DEFAULT_BUF_SIZE],
- buf_pos: 0,
- buf_cap: 0,
- }
- }
-
- fn fill_inner_buf(&mut self) -> Result<&[u8]> {
- if self.buf_pos >= self.buf_cap {
- // If we've reached the end of our internal buffer then we need to fetch
- // some more data from the underlying reader.
- // Branch using `>=` instead of the more correct `==`
- // to tell the compiler that the pos..cap slice is always valid.
- debug_assert!(self.buf_pos == self.buf_cap);
- let mut reader = self.reader.borrow_mut();
- reader.seek(SeekFrom::Start(self.start))?; // always seek to start before reading
- self.buf_cap = reader.read(&mut self.buf)?;
- self.buf_pos = 0;
- }
- Ok(&self.buf[self.buf_pos..self.buf_cap])
- }
-
- fn skip_inner_buf(&mut self, buf: &mut [u8]) -> Result<usize> {
- // discard buffer
- self.buf_pos = 0;
- self.buf_cap = 0;
- // read directly into param buffer
- let mut reader = self.reader.borrow_mut();
- reader.seek(SeekFrom::Start(self.start))?; // always seek to start before reading
- let nread = reader.read(buf)?;
- self.start += nread as u64;
- Ok(nread)
- }
-}
-
-impl<R: ParquetReader> Read for FileSource<R> {
- fn read(&mut self, buf: &mut [u8]) -> Result<usize> {
- let bytes_to_read = cmp::min(buf.len(), (self.end - self.start) as usize);
- let buf = &mut buf[0..bytes_to_read];
-
- // If we don't have any buffered data and we're doing a massive read
- // (larger than our internal buffer), bypass our internal buffer
- // entirely.
- if self.buf_pos == self.buf_cap && buf.len() >= self.buf.len() {
- return self.skip_inner_buf(buf);
- }
- let nread = {
- let mut rem = self.fill_inner_buf()?;
- // copy the data from the inner buffer to the param buffer
- rem.read(buf)?
- };
- // consume from buffer
- self.buf_pos = cmp::min(self.buf_pos + nread, self.buf_cap);
-
- self.start += nread as u64;
- Ok(nread)
- }
-}
-
-impl<R: ParquetReader> Position for FileSource<R> {
- fn pos(&self) -> u64 {
- self.start
- }
-}
-
-impl<R: ParquetReader> Length for FileSource<R> {
- fn len(&self) -> u64 {
- self.end - self.start
- }
-}
-
-/// Struct that represents `File` output stream with position tracking.
-/// Used as a sink in file writer.
-pub struct FileSink<W: ParquetWriter> {
- buf: BufWriter<W>,
- // This is not necessarily position in the underlying file,
- // but rather current position in the sink.
- pos: u64,
-}
-
-impl<W: ParquetWriter> FileSink<W> {
- /// Creates new file sink.
- /// Position is set to whatever position file has.
- pub fn new(buf: &W) -> Self {
- let mut owned_buf = buf.try_clone().unwrap();
- let pos = owned_buf.seek(SeekFrom::Current(0)).unwrap();
- Self {
- buf: BufWriter::new(owned_buf),
- pos,
- }
- }
-}
-
-impl<W: ParquetWriter> Write for FileSink<W> {
- fn write(&mut self, buf: &[u8]) -> Result<usize> {
- let num_bytes = self.buf.write(buf)?;
- self.pos += num_bytes as u64;
- Ok(num_bytes)
- }
-
- fn flush(&mut self) -> Result<()> {
- self.buf.flush()
- }
-}
-
-impl<W: ParquetWriter> Position for FileSink<W> {
- fn pos(&self) -> u64 {
- self.pos
- }
-}
-
-// Position implementation for Cursor to use in various tests.
-impl<'a> Position for Cursor<&'a mut Vec<u8>> {
- fn pos(&self) -> u64 {
- self.position()
- }
-}
-
-#[cfg(test)]
-mod tests {
- use super::*;
-
- use std::iter;
-
- use crate::util::test_common::{get_temp_file, get_test_file};
-
- #[test]
- fn test_io_read_fully() {
- let mut buf = vec![0; 8];
- let mut src = FileSource::new(&get_test_file("alltypes_plain.parquet"), 0, 4);
-
- let bytes_read = src.read(&mut buf[..]).unwrap();
- assert_eq!(bytes_read, 4);
- assert_eq!(buf, vec![b'P', b'A', b'R', b'1', 0, 0, 0, 0]);
- }
-
- #[test]
- fn test_io_read_in_chunks() {
- let mut buf = vec![0; 4];
- let mut src = FileSource::new(&get_test_file("alltypes_plain.parquet"), 0, 4);
-
- let bytes_read = src.read(&mut buf[0..2]).unwrap();
- assert_eq!(bytes_read, 2);
- let bytes_read = src.read(&mut buf[2..]).unwrap();
- assert_eq!(bytes_read, 2);
- assert_eq!(buf, vec![b'P', b'A', b'R', b'1']);
- }
-
- #[test]
- fn test_io_read_pos() {
- let mut src = FileSource::new(&get_test_file("alltypes_plain.parquet"), 0, 4);
-
- let _ = src.read(&mut [0; 1]).unwrap();
- assert_eq!(src.pos(), 1);
-
- let _ = src.read(&mut [0; 4]).unwrap();
- assert_eq!(src.pos(), 4);
- }
-
- #[test]
- fn test_io_read_over_limit() {
- let mut src = FileSource::new(&get_test_file("alltypes_plain.parquet"), 0, 4);
-
- // Read all bytes from source
- let _ = src.read(&mut [0; 128]).unwrap();
- assert_eq!(src.pos(), 4);
-
- // Try reading again, should return 0 bytes.
- let bytes_read = src.read(&mut [0; 128]).unwrap();
- assert_eq!(bytes_read, 0);
- assert_eq!(src.pos(), 4);
- }
-
- #[test]
- fn test_io_seek_switch() {
- let mut buf = vec![0; 4];
- let mut file = get_test_file("alltypes_plain.parquet");
- let mut src = FileSource::new(&file, 0, 4);
-
- file.seek(SeekFrom::Start(5_u64))
- .expect("File seek to a position");
-
- let bytes_read = src.read(&mut buf[..]).unwrap();
- assert_eq!(bytes_read, 4);
- assert_eq!(buf, vec![b'P', b'A', b'R', b'1']);
- }
-
- #[test]
- fn test_io_write_with_pos() {
- let mut file = get_temp_file("file_sink_test", &[b'a', b'b', b'c']);
- file.seek(SeekFrom::Current(3)).unwrap();
-
- // Write into sink
- let mut sink = FileSink::new(&file);
- assert_eq!(sink.pos(), 3);
-
- sink.write_all(&[b'd', b'e', b'f', b'g']).unwrap();
- assert_eq!(sink.pos(), 7);
-
- sink.flush().unwrap();
- assert_eq!(sink.pos(), file.seek(SeekFrom::Current(0)).unwrap());
-
- // Read data using file chunk
- let mut res = vec![0u8; 7];
- let mut chunk =
- FileSource::new(&file, 0, file.metadata().unwrap().len() as usize);
- chunk.read_exact(&mut res[..]).unwrap();
- assert_eq!(res, vec![b'a', b'b', b'c', b'd', b'e', b'f', b'g']);
- }
-
- #[test]
- fn test_io_large_read() {
- // Generate repeated 'abcdef' pattern and write it into a file
- let patterned_data: Vec<u8> = iter::repeat(vec![0, 1, 2, 3, 4, 5])
- .flatten()
- .take(3 * DEFAULT_BUF_SIZE)
- .collect();
- // always use different temp files as test might be run in parallel
- let mut file = get_temp_file("large_file_sink_test", &patterned_data);
-
- // seek the underlying file to the first 'd'
- file.seek(SeekFrom::Start(3)).unwrap();
-
- // create the FileSource reader that starts at pos 1 ('b')
- let mut chunk = FileSource::new(&file, 1, patterned_data.len() - 1);
-
- // read the 'b' at pos 1
- let mut res = vec![0u8; 1];
- chunk.read_exact(&mut res).unwrap();
- assert_eq!(res, &[1]);
-
- // the underlying file is sought to 'e'
- file.seek(SeekFrom::Start(4)).unwrap();
-
- // now read large chunk that starts with 'c' (after 'b')
- let mut res = vec![0u8; 2 * DEFAULT_BUF_SIZE];
- chunk.read_exact(&mut res).unwrap();
- assert_eq!(
- res,
- &patterned_data[2..2 + 2 * DEFAULT_BUF_SIZE],
- "read buf and original data are not equal"
- );
- }
-}
diff --git a/parquet/src/util/memory.rs b/parquet/src/util/memory.rs
deleted file mode 100644
index 1642a4b..0000000
--- a/parquet/src/util/memory.rs
+++ /dev/null
@@ -1,545 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-//! Utility methods and structs for working with memory.
-
-use std::{
- fmt::{Debug, Display, Formatter, Result as FmtResult},
- io::{Result as IoResult, Write},
- mem,
- ops::{Index, IndexMut},
- sync::{
- atomic::{AtomicI64, Ordering},
- Arc, Weak,
- },
-};
-
-// ----------------------------------------------------------------------
-// Memory Tracker classes
-
-/// Reference counted pointer for [`MemTracker`].
-pub type MemTrackerPtr = Arc<MemTracker>;
-/// Non-owning reference for [`MemTracker`].
-pub type WeakMemTrackerPtr = Weak<MemTracker>;
-
-/// Struct to track memory usage information.
-#[derive(Debug)]
-pub struct MemTracker {
- // In the tuple, the first element is the current memory allocated (in bytes),
- // and the second element is the maximum memory allocated so far (in bytes).
- current_memory_usage: AtomicI64,
- max_memory_usage: AtomicI64,
-}
-
-impl MemTracker {
- /// Creates new memory tracker.
- #[inline]
- pub fn new() -> MemTracker {
- MemTracker {
- current_memory_usage: Default::default(),
- max_memory_usage: Default::default(),
- }
- }
-
- /// Returns the current memory consumption, in bytes.
- pub fn memory_usage(&self) -> i64 {
- self.current_memory_usage.load(Ordering::Acquire)
- }
-
- /// Returns the maximum memory consumption so far, in bytes.
- pub fn max_memory_usage(&self) -> i64 {
- self.max_memory_usage.load(Ordering::Acquire)
- }
-
- /// Adds `num_bytes` to the memory consumption tracked by this memory tracker.
- #[inline]
- pub fn alloc(&self, num_bytes: i64) {
- let new_current = self
- .current_memory_usage
- .fetch_add(num_bytes, Ordering::Acquire)
- + num_bytes;
- self.max_memory_usage
- .fetch_max(new_current, Ordering::Acquire);
- }
-}
-
-// ----------------------------------------------------------------------
-// Buffer classes
-
-/// Type alias for [`Buffer`].
-pub type ByteBuffer = Buffer<u8>;
-/// Type alias for [`BufferPtr`].
-pub type ByteBufferPtr = BufferPtr<u8>;
-
-/// A resize-able buffer class with generic member, with optional memory tracker.
-///
-/// Note that a buffer has two attributes:
-/// `capacity` and `size`: the former is the total number of space reserved for
-/// the buffer, while the latter is the actual number of elements.
-/// Invariant: `capacity` >= `size`.
-/// The total allocated bytes for a buffer equals to `capacity * sizeof<T>()`.
-pub struct Buffer<T: Clone> {
- data: Vec<T>,
- mem_tracker: Option<MemTrackerPtr>,
- type_length: usize,
-}
-
-impl<T: Clone> Buffer<T> {
- /// Creates new empty buffer.
- pub fn new() -> Self {
- Buffer {
- data: vec![],
- mem_tracker: None,
- type_length: std::mem::size_of::<T>(),
- }
- }
-
- /// Adds [`MemTracker`] for this buffer.
- #[inline]
- pub fn with_mem_tracker(mut self, mc: MemTrackerPtr) -> Self {
- mc.alloc((self.data.capacity() * self.type_length) as i64);
- self.mem_tracker = Some(mc);
- self
- }
-
- /// Returns slice of data in this buffer.
- #[inline]
- pub fn data(&self) -> &[T] {
- self.data.as_slice()
- }
-
- /// Sets data for this buffer.
- #[inline]
- pub fn set_data(&mut self, new_data: Vec<T>) {
- if let Some(ref mc) = self.mem_tracker {
- let capacity_diff = new_data.capacity() as i64 - self.data.capacity() as i64;
- mc.alloc(capacity_diff * self.type_length as i64);
- }
- self.data = new_data;
- }
-
- /// Resizes underlying data in place to a new length `new_size`.
- ///
- /// If `new_size` is less than current length, data is truncated, otherwise, it is
- /// extended to `new_size` with provided default value `init_value`.
- ///
- /// Memory tracker is also updated, if available.
- #[inline]
- pub fn resize(&mut self, new_size: usize, init_value: T) {
- let old_capacity = self.data.capacity();
- self.data.resize(new_size, init_value);
- if let Some(ref mc) = self.mem_tracker {
- let capacity_diff = self.data.capacity() as i64 - old_capacity as i64;
- mc.alloc(capacity_diff * self.type_length as i64);
- }
- }
-
- /// Clears underlying data.
- #[inline]
- pub fn clear(&mut self) {
- self.data.clear()
- }
-
- /// Reserves capacity `additional_capacity` for underlying data vector.
- ///
- /// Memory tracker is also updated, if available.
- #[inline]
- pub fn reserve(&mut self, additional_capacity: usize) {
- let old_capacity = self.data.capacity();
- self.data.reserve(additional_capacity);
- if self.data.capacity() > old_capacity {
- if let Some(ref mc) = self.mem_tracker {
- let capacity_diff = self.data.capacity() as i64 - old_capacity as i64;
- mc.alloc(capacity_diff * self.type_length as i64);
- }
- }
- }
-
- /// Returns [`BufferPtr`] with buffer data.
- /// Buffer data is reset.
- #[inline]
- pub fn consume(&mut self) -> BufferPtr<T> {
- let old_data = mem::replace(&mut self.data, vec![]);
- let mut result = BufferPtr::new(old_data);
- if let Some(ref mc) = self.mem_tracker {
- result = result.with_mem_tracker(mc.clone());
- }
- result
- }
-
- /// Adds `value` to the buffer.
- #[inline]
- pub fn push(&mut self, value: T) {
- self.data.push(value)
- }
-
- /// Returns current capacity for the buffer.
- #[inline]
- pub fn capacity(&self) -> usize {
- self.data.capacity()
- }
-
- /// Returns current size for the buffer.
- #[inline]
- pub fn size(&self) -> usize {
- self.data.len()
- }
-
- /// Returns `true` if memory tracker is added to buffer, `false` otherwise.
- #[inline]
- pub fn is_mem_tracked(&self) -> bool {
- self.mem_tracker.is_some()
- }
-
- /// Returns memory tracker associated with this buffer.
- /// This may panic, if memory tracker is not set, use method above to check if
- /// memory tracker is available.
- #[inline]
- pub fn mem_tracker(&self) -> &MemTrackerPtr {
- self.mem_tracker.as_ref().unwrap()
- }
-}
-
-impl<T: Sized + Clone> Index<usize> for Buffer<T> {
- type Output = T;
-
- fn index(&self, index: usize) -> &T {
- &self.data[index]
- }
-}
-
-impl<T: Sized + Clone> IndexMut<usize> for Buffer<T> {
- fn index_mut(&mut self, index: usize) -> &mut T {
- &mut self.data[index]
- }
-}
-
-// TODO: implement this for other types
-impl Write for Buffer<u8> {
- #[inline]
- fn write(&mut self, buf: &[u8]) -> IoResult<usize> {
- let old_capacity = self.data.capacity();
- let bytes_written = self.data.write(buf)?;
- if let Some(ref mc) = self.mem_tracker {
- if self.data.capacity() - old_capacity > 0 {
- mc.alloc((self.data.capacity() - old_capacity) as i64)
- }
- }
- Ok(bytes_written)
- }
-
- fn flush(&mut self) -> IoResult<()> {
- // No-op
- self.data.flush()
- }
-}
-
-impl AsRef<[u8]> for Buffer<u8> {
- fn as_ref(&self) -> &[u8] {
- self.data.as_slice()
- }
-}
-
-impl<T: Clone> Drop for Buffer<T> {
- #[inline]
- fn drop(&mut self) {
- if let Some(ref mc) = self.mem_tracker {
- mc.alloc(-((self.data.capacity() * self.type_length) as i64));
- }
- }
-}
-
-// ----------------------------------------------------------------------
-// Immutable Buffer (BufferPtr) classes
-
-/// An representation of a slice on a reference-counting and read-only byte array.
-/// Sub-slices can be further created from this. The byte array will be released
-/// when all slices are dropped.
-#[allow(clippy::rc_buffer)]
-#[derive(Clone, Debug)]
-pub struct BufferPtr<T> {
- data: Arc<Vec<T>>,
- start: usize,
- len: usize,
- // TODO: will this create too many references? rethink about this.
- mem_tracker: Option<MemTrackerPtr>,
-}
-
-impl<T> BufferPtr<T> {
- /// Creates new buffer from a vector.
- pub fn new(v: Vec<T>) -> Self {
- let len = v.len();
- Self {
- data: Arc::new(v),
- start: 0,
- len,
- mem_tracker: None,
- }
- }
-
- /// Returns slice of data in this buffer.
- #[inline]
- pub fn data(&self) -> &[T] {
- &self.data[self.start..self.start + self.len]
- }
-
- /// Updates this buffer with new `start` position and length `len`.
- ///
- /// Range should be within current start position and length.
- #[inline]
- pub fn with_range(mut self, start: usize, len: usize) -> Self {
- self.set_range(start, len);
- self
- }
-
- /// Updates this buffer with new `start` position and length `len`.
- ///
- /// Range should be within current start position and length.
- #[inline]
- pub fn set_range(&mut self, start: usize, len: usize) {
- assert!(self.start <= start && start + len <= self.start + self.len);
- self.start = start;
- self.len = len;
- }
-
- /// Adds memory tracker to this buffer.
- pub fn with_mem_tracker(mut self, mc: MemTrackerPtr) -> Self {
- self.mem_tracker = Some(mc);
- self
- }
-
- /// Returns start position of this buffer.
- #[inline]
- pub fn start(&self) -> usize {
- self.start
- }
-
- /// Returns length of this buffer
- #[inline]
- pub fn len(&self) -> usize {
- self.len
- }
-
- /// Returns whether this buffer is empty
- #[inline]
- pub fn is_empty(&self) -> bool {
- self.len == 0
- }
-
- /// Returns `true` if this buffer has memory tracker, `false` otherwise.
- pub fn is_mem_tracked(&self) -> bool {
- self.mem_tracker.is_some()
- }
-
- /// Returns a shallow copy of the buffer.
- /// Reference counted pointer to the data is copied.
- pub fn all(&self) -> BufferPtr<T> {
- BufferPtr {
- data: self.data.clone(),
- start: self.start,
- len: self.len,
- mem_tracker: self.mem_tracker.as_ref().cloned(),
- }
- }
-
- /// Returns a shallow copy of the buffer that starts with `start` position.
- pub fn start_from(&self, start: usize) -> BufferPtr<T> {
- assert!(start <= self.len);
- BufferPtr {
- data: self.data.clone(),
- start: self.start + start,
- len: self.len - start,
- mem_tracker: self.mem_tracker.as_ref().cloned(),
- }
- }
-
- /// Returns a shallow copy that is a range slice within this buffer.
- pub fn range(&self, start: usize, len: usize) -> BufferPtr<T> {
- assert!(start + len <= self.len);
- BufferPtr {
- data: self.data.clone(),
- start: self.start + start,
- len,
- mem_tracker: self.mem_tracker.as_ref().cloned(),
- }
- }
-}
-
-impl<T: Sized> Index<usize> for BufferPtr<T> {
- type Output = T;
-
- fn index(&self, index: usize) -> &T {
- assert!(index < self.len);
- &self.data[self.start + index]
- }
-}
-
-impl<T: Debug> Display for BufferPtr<T> {
- fn fmt(&self, f: &mut Formatter) -> FmtResult {
- write!(f, "{:?}", self.data)
- }
-}
-
-impl<T> Drop for BufferPtr<T> {
- fn drop(&mut self) {
- if let Some(ref mc) = self.mem_tracker {
- if Arc::strong_count(&self.data) == 1 && Arc::weak_count(&self.data) == 0 {
- mc.alloc(-(self.data.capacity() as i64));
- }
- }
- }
-}
-
-impl AsRef<[u8]> for BufferPtr<u8> {
- #[inline]
- fn as_ref(&self) -> &[u8] {
- &self.data[self.start..self.start + self.len]
- }
-}
-
-#[cfg(test)]
-mod tests {
- use super::*;
-
- #[test]
- fn test_byte_buffer_mem_tracker() {
- let mem_tracker = Arc::new(MemTracker::new());
-
- let mut buffer = ByteBuffer::new().with_mem_tracker(mem_tracker.clone());
- buffer.set_data(vec![0; 10]);
- assert_eq!(mem_tracker.memory_usage(), buffer.capacity() as i64);
- buffer.set_data(vec![0; 20]);
- let capacity = buffer.capacity() as i64;
- assert_eq!(mem_tracker.memory_usage(), capacity);
-
- let max_capacity = {
- let mut buffer2 = ByteBuffer::new().with_mem_tracker(mem_tracker.clone());
- buffer2.reserve(30);
- assert_eq!(
- mem_tracker.memory_usage(),
- buffer2.capacity() as i64 + capacity
- );
- buffer2.set_data(vec![0; 100]);
- assert_eq!(
- mem_tracker.memory_usage(),
- buffer2.capacity() as i64 + capacity
- );
- buffer2.capacity() as i64 + capacity
- };
-
- assert_eq!(mem_tracker.memory_usage(), capacity);
- assert_eq!(mem_tracker.max_memory_usage(), max_capacity);
-
- buffer.reserve(40);
- assert_eq!(mem_tracker.memory_usage(), buffer.capacity() as i64);
-
- buffer.consume();
- assert_eq!(mem_tracker.memory_usage(), buffer.capacity() as i64);
- }
-
- #[test]
- fn test_byte_ptr_mem_tracker() {
- let mem_tracker = Arc::new(MemTracker::new());
-
- let mut buffer = ByteBuffer::new().with_mem_tracker(mem_tracker.clone());
- buffer.set_data(vec![0; 60]);
-
- {
- let buffer_capacity = buffer.capacity() as i64;
- let buf_ptr = buffer.consume();
- assert_eq!(mem_tracker.memory_usage(), buffer_capacity);
- {
- let buf_ptr1 = buf_ptr.all();
- {
- let _ = buf_ptr.start_from(20);
- assert_eq!(mem_tracker.memory_usage(), buffer_capacity);
- }
- assert_eq!(mem_tracker.memory_usage(), buffer_capacity);
- let _ = buf_ptr1.range(30, 20);
- assert_eq!(mem_tracker.memory_usage(), buffer_capacity);
- }
- assert_eq!(mem_tracker.memory_usage(), buffer_capacity);
- }
- assert_eq!(mem_tracker.memory_usage(), buffer.capacity() as i64);
- }
-
- #[test]
- fn test_byte_buffer() {
- let mut buffer = ByteBuffer::new();
- assert_eq!(buffer.size(), 0);
- assert_eq!(buffer.capacity(), 0);
-
- let mut buffer2 = ByteBuffer::new();
- buffer2.reserve(40);
- assert_eq!(buffer2.size(), 0);
- assert_eq!(buffer2.capacity(), 40);
-
- buffer.set_data((0..5).collect());
- assert_eq!(buffer.size(), 5);
- assert_eq!(buffer[4], 4);
-
- buffer.set_data((0..20).collect());
- assert_eq!(buffer.size(), 20);
- assert_eq!(buffer[10], 10);
-
- let expected: Vec<u8> = (0..20).collect();
- {
- let data = buffer.data();
- assert_eq!(data, expected.as_slice());
- }
-
- buffer.reserve(40);
- assert!(buffer.capacity() >= 40);
-
- let byte_ptr = buffer.consume();
- assert_eq!(buffer.size(), 0);
- assert_eq!(byte_ptr.as_ref(), expected.as_slice());
-
- let values: Vec<u8> = (0..30).collect();
- let _ = buffer.write(values.as_slice());
- let _ = buffer.flush();
-
- assert_eq!(buffer.data(), values.as_slice());
- }
-
- #[test]
- fn test_byte_ptr() {
- let values = (0..50).collect();
- let ptr = ByteBufferPtr::new(values);
- assert_eq!(ptr.len(), 50);
- assert_eq!(ptr.start(), 0);
- assert_eq!(ptr[40], 40);
-
- let ptr2 = ptr.all();
- assert_eq!(ptr2.len(), 50);
- assert_eq!(ptr2.start(), 0);
- assert_eq!(ptr2[40], 40);
-
- let ptr3 = ptr.start_from(20);
- assert_eq!(ptr3.len(), 30);
- assert_eq!(ptr3.start(), 20);
- assert_eq!(ptr3[0], 20);
-
- let ptr4 = ptr3.range(10, 10);
- assert_eq!(ptr4.len(), 10);
- assert_eq!(ptr4.start(), 30);
- assert_eq!(ptr4[0], 30);
-
- let expected: Vec<u8> = (30..40).collect();
- assert_eq!(ptr4.as_ref(), expected.as_slice());
- }
-}
diff --git a/parquet/src/util/mod.rs b/parquet/src/util/mod.rs
deleted file mode 100644
index 8f6d85d..0000000
--- a/parquet/src/util/mod.rs
+++ /dev/null
@@ -1,28 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-pub mod io;
-pub mod memory;
-#[macro_use]
-pub mod bit_util;
-mod bit_packing;
-pub mod cursor;
-pub mod hash_util;
-pub(crate) mod test_common;
-pub use self::test_common::page_util::{
- DataPageBuilder, DataPageBuilderImpl, InMemoryPageIterator,
-};
diff --git a/parquet/src/util/test_common/file_util.rs b/parquet/src/util/test_common/file_util.rs
deleted file mode 100644
index 7393b55..0000000
--- a/parquet/src/util/test_common/file_util.rs
+++ /dev/null
@@ -1,73 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-use std::{env, fs, io::Write, path::PathBuf, str::FromStr};
-
-/// Returns path to the test parquet file in 'data' directory
-pub fn get_test_path(file_name: &str) -> PathBuf {
- let mut pathbuf =
- PathBuf::from_str(&arrow::util::test_util::parquet_test_data()).unwrap();
- pathbuf.push(file_name);
- pathbuf
-}
-
-/// Returns file handle for a test parquet file from 'data' directory
-pub fn get_test_file(file_name: &str) -> fs::File {
- let path = get_test_path(file_name);
- fs::File::open(path.as_path()).unwrap_or_else(|err| {
- panic!(
- "Test file {} could not be opened, did you do `git submodule update`?: {}",
- path.display(),
- err
- )
- })
-}
-
-/// Returns file handle for a temp file in 'target' directory with a provided content
-pub fn get_temp_file(file_name: &str, content: &[u8]) -> fs::File {
- // build tmp path to a file in "target/debug/testdata"
- let mut path_buf = env::current_dir().unwrap();
- path_buf.push("target");
- path_buf.push("debug");
- path_buf.push("testdata");
- fs::create_dir_all(&path_buf).unwrap();
- path_buf.push(file_name);
-
- // write file content
- let mut tmp_file = fs::File::create(path_buf.as_path()).unwrap();
- tmp_file.write_all(content).unwrap();
- tmp_file.sync_all().unwrap();
-
- // return file handle for both read and write
- let file = fs::OpenOptions::new()
- .read(true)
- .write(true)
- .open(path_buf.as_path());
- assert!(file.is_ok());
- file.unwrap()
-}
-
-pub fn get_temp_filename() -> PathBuf {
- let mut path_buf = env::current_dir().unwrap();
- path_buf.push("target");
- path_buf.push("debug");
- path_buf.push("testdata");
- fs::create_dir_all(&path_buf).unwrap();
- path_buf.push(rand::random::<i16>().to_string());
-
- path_buf
-}
diff --git a/parquet/src/util/test_common/mod.rs b/parquet/src/util/test_common/mod.rs
deleted file mode 100644
index ed65bbe..0000000
--- a/parquet/src/util/test_common/mod.rs
+++ /dev/null
@@ -1,33 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-pub mod file_util;
-pub mod page_util;
-pub mod rand_gen;
-
-pub use self::rand_gen::random_bools;
-pub use self::rand_gen::random_bytes;
-pub use self::rand_gen::random_numbers;
-pub use self::rand_gen::random_numbers_range;
-pub use self::rand_gen::RandGen;
-
-pub use self::file_util::get_temp_file;
-pub use self::file_util::get_temp_filename;
-pub use self::file_util::get_test_file;
-pub use self::file_util::get_test_path;
-
-pub use self::page_util::make_pages;
diff --git a/parquet/src/util/test_common/page_util.rs b/parquet/src/util/test_common/page_util.rs
deleted file mode 100644
index 581845a..0000000
--- a/parquet/src/util/test_common/page_util.rs
+++ /dev/null
@@ -1,320 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-use crate::basic::Encoding;
-use crate::column::page::PageReader;
-use crate::column::page::{Page, PageIterator};
-use crate::data_type::DataType;
-use crate::encodings::encoding::{get_encoder, DictEncoder, Encoder};
-use crate::encodings::levels::max_buffer_size;
-use crate::encodings::levels::LevelEncoder;
-use crate::errors::Result;
-use crate::schema::types::{ColumnDescPtr, SchemaDescPtr};
-use crate::util::memory::ByteBufferPtr;
-use crate::util::memory::MemTracker;
-use crate::util::memory::MemTrackerPtr;
-use crate::util::test_common::random_numbers_range;
-use rand::distributions::uniform::SampleUniform;
-use std::collections::VecDeque;
-use std::mem;
-use std::sync::Arc;
-
-pub trait DataPageBuilder {
- fn add_rep_levels(&mut self, max_level: i16, rep_levels: &[i16]);
- fn add_def_levels(&mut self, max_level: i16, def_levels: &[i16]);
- fn add_values<T: DataType>(&mut self, encoding: Encoding, values: &[T::T]);
- fn add_indices(&mut self, indices: ByteBufferPtr);
- fn consume(self) -> Page;
-}
-
-/// A utility struct for building data pages (v1 or v2). Callers must call:
-/// - add_rep_levels()
-/// - add_def_levels()
-/// - add_values() for normal data page / add_indices() for dictionary data page
-/// - consume()
-/// in order to populate and obtain a data page.
-pub struct DataPageBuilderImpl {
- desc: ColumnDescPtr,
- encoding: Option<Encoding>,
- mem_tracker: MemTrackerPtr,
- num_values: u32,
- buffer: Vec<u8>,
- rep_levels_byte_len: u32,
- def_levels_byte_len: u32,
- datapage_v2: bool,
-}
-
-impl DataPageBuilderImpl {
- // `num_values` is the number of non-null values to put in the data page.
- // `datapage_v2` flag is used to indicate if the generated data page should use V2
- // format or not.
- pub fn new(desc: ColumnDescPtr, num_values: u32, datapage_v2: bool) -> Self {
- DataPageBuilderImpl {
- desc,
- encoding: None,
- mem_tracker: Arc::new(MemTracker::new()),
- num_values,
- buffer: vec![],
- rep_levels_byte_len: 0,
- def_levels_byte_len: 0,
- datapage_v2,
- }
- }
-
- // Adds levels to the buffer and return number of encoded bytes
- fn add_levels(&mut self, max_level: i16, levels: &[i16]) -> u32 {
- if max_level <= 0 {
- return 0;
- }
- let size = max_buffer_size(Encoding::RLE, max_level, levels.len());
- let mut level_encoder = LevelEncoder::v1(Encoding::RLE, max_level, vec![0; size]);
- level_encoder.put(levels).expect("put() should be OK");
- let encoded_levels = level_encoder.consume().expect("consume() should be OK");
- // Actual encoded bytes (without length offset)
- let encoded_bytes = &encoded_levels[mem::size_of::<i32>()..];
- if self.datapage_v2 {
- // Level encoder always initializes with offset of i32, where it stores
- // length of encoded data; for data page v2 we explicitly
- // store length, therefore we should skip i32 bytes.
- self.buffer.extend_from_slice(encoded_bytes);
- } else {
- self.buffer.extend_from_slice(encoded_levels.as_slice());
- }
- encoded_bytes.len() as u32
- }
-}
-
-impl DataPageBuilder for DataPageBuilderImpl {
- fn add_rep_levels(&mut self, max_levels: i16, rep_levels: &[i16]) {
- self.num_values = rep_levels.len() as u32;
- self.rep_levels_byte_len = self.add_levels(max_levels, rep_levels);
- }
-
- fn add_def_levels(&mut self, max_levels: i16, def_levels: &[i16]) {
- assert!(
- self.num_values == def_levels.len() as u32,
- "Must call `add_rep_levels() first!`"
- );
-
- self.def_levels_byte_len = self.add_levels(max_levels, def_levels);
- }
-
- fn add_values<T: DataType>(&mut self, encoding: Encoding, values: &[T::T]) {
- assert!(
- self.num_values >= values.len() as u32,
- "num_values: {}, values.len(): {}",
- self.num_values,
- values.len()
- );
- self.encoding = Some(encoding);
- let mut encoder: Box<dyn Encoder<T>> =
- get_encoder::<T>(self.desc.clone(), encoding, self.mem_tracker.clone())
- .expect("get_encoder() should be OK");
- encoder.put(values).expect("put() should be OK");
- let encoded_values = encoder
- .flush_buffer()
- .expect("consume_buffer() should be OK");
- self.buffer.extend_from_slice(encoded_values.data());
- }
-
- fn add_indices(&mut self, indices: ByteBufferPtr) {
- self.encoding = Some(Encoding::RLE_DICTIONARY);
- self.buffer.extend_from_slice(indices.data());
- }
-
- fn consume(self) -> Page {
- if self.datapage_v2 {
- Page::DataPageV2 {
- buf: ByteBufferPtr::new(self.buffer),
- num_values: self.num_values,
- encoding: self.encoding.unwrap(),
- num_nulls: 0, /* set to dummy value - don't need this when reading
- * data page */
- num_rows: self.num_values, /* also don't need this when reading
- * data page */
- def_levels_byte_len: self.def_levels_byte_len,
- rep_levels_byte_len: self.rep_levels_byte_len,
- is_compressed: false,
- statistics: None, // set to None, we do not need statistics for tests
- }
- } else {
- Page::DataPage {
- buf: ByteBufferPtr::new(self.buffer),
- num_values: self.num_values,
- encoding: self.encoding.unwrap(),
- def_level_encoding: Encoding::RLE,
- rep_level_encoding: Encoding::RLE,
- statistics: None, // set to None, we do not need statistics for tests
- }
- }
- }
-}
-
-/// A utility page reader which stores pages in memory.
-pub struct InMemoryPageReader<P: Iterator<Item = Page>> {
- page_iter: P,
-}
-
-impl<P: Iterator<Item = Page>> InMemoryPageReader<P> {
- pub fn new(pages: impl IntoIterator<Item = Page, IntoIter = P>) -> Self {
- Self {
- page_iter: pages.into_iter(),
- }
- }
-}
-
-impl<P: Iterator<Item = Page>> PageReader for InMemoryPageReader<P> {
- fn get_next_page(&mut self) -> Result<Option<Page>> {
- Ok(self.page_iter.next())
- }
-}
-
-impl<P: Iterator<Item = Page>> Iterator for InMemoryPageReader<P> {
- type Item = Result<Page>;
-
- fn next(&mut self) -> Option<Self::Item> {
- self.get_next_page().transpose()
- }
-}
-
-/// A utility page iterator which stores page readers in memory, used for tests.
-#[derive(Clone)]
-pub struct InMemoryPageIterator<I: Iterator<Item = Vec<Page>>> {
- schema: SchemaDescPtr,
- column_desc: ColumnDescPtr,
- page_reader_iter: I,
-}
-
-impl<I: Iterator<Item = Vec<Page>>> InMemoryPageIterator<I> {
- pub fn new(
- schema: SchemaDescPtr,
- column_desc: ColumnDescPtr,
- pages: impl IntoIterator<Item = Vec<Page>, IntoIter = I>,
- ) -> Self {
- Self {
- schema,
- column_desc,
- page_reader_iter: pages.into_iter(),
- }
- }
-}
-
-impl<I: Iterator<Item = Vec<Page>>> Iterator for InMemoryPageIterator<I> {
- type Item = Result<Box<dyn PageReader>>;
-
- fn next(&mut self) -> Option<Self::Item> {
- self.page_reader_iter
- .next()
- .map(|x| Ok(Box::new(InMemoryPageReader::new(x)) as Box<dyn PageReader>))
- }
-}
-
-impl<I: Iterator<Item = Vec<Page>>> PageIterator for InMemoryPageIterator<I> {
- fn schema(&mut self) -> Result<SchemaDescPtr> {
- Ok(self.schema.clone())
- }
-
- fn column_schema(&mut self) -> Result<ColumnDescPtr> {
- Ok(self.column_desc.clone())
- }
-}
-
-pub fn make_pages<T: DataType>(
- desc: ColumnDescPtr,
- encoding: Encoding,
- num_pages: usize,
- levels_per_page: usize,
- min: T::T,
- max: T::T,
- def_levels: &mut Vec<i16>,
- rep_levels: &mut Vec<i16>,
- values: &mut Vec<T::T>,
- pages: &mut VecDeque<Page>,
- use_v2: bool,
-) where
- T::T: PartialOrd + SampleUniform + Copy,
-{
- let mut num_values = 0;
- let max_def_level = desc.max_def_level();
- let max_rep_level = desc.max_rep_level();
-
- let mem_tracker = Arc::new(MemTracker::new());
- let mut dict_encoder = DictEncoder::<T>::new(desc.clone(), mem_tracker);
-
- for i in 0..num_pages {
- let mut num_values_cur_page = 0;
- let level_range = i * levels_per_page..(i + 1) * levels_per_page;
-
- if max_def_level > 0 {
- random_numbers_range(levels_per_page, 0, max_def_level + 1, def_levels);
- for dl in &def_levels[level_range.clone()] {
- if *dl == max_def_level {
- num_values_cur_page += 1;
- }
- }
- } else {
- num_values_cur_page = levels_per_page;
- }
- if max_rep_level > 0 {
- random_numbers_range(levels_per_page, 0, max_rep_level + 1, rep_levels);
- }
- random_numbers_range(num_values_cur_page, min, max, values);
-
- // Generate the current page
-
- let mut pb =
- DataPageBuilderImpl::new(desc.clone(), num_values_cur_page as u32, use_v2);
- if max_rep_level > 0 {
- pb.add_rep_levels(max_rep_level, &rep_levels[level_range.clone()]);
- }
- if max_def_level > 0 {
- pb.add_def_levels(max_def_level, &def_levels[level_range]);
- }
-
- let value_range = num_values..num_values + num_values_cur_page;
- match encoding {
- Encoding::PLAIN_DICTIONARY | Encoding::RLE_DICTIONARY => {
- let _ = dict_encoder.put(&values[value_range.clone()]);
- let indices = dict_encoder
- .write_indices()
- .expect("write_indices() should be OK");
- pb.add_indices(indices);
- }
- Encoding::PLAIN => {
- pb.add_values::<T>(encoding, &values[value_range]);
- }
- enc => panic!("Unexpected encoding {}", enc),
- }
-
- let data_page = pb.consume();
- pages.push_back(data_page);
- num_values += num_values_cur_page;
- }
-
- if encoding == Encoding::PLAIN_DICTIONARY || encoding == Encoding::RLE_DICTIONARY {
- let dict = dict_encoder
- .write_dict()
- .expect("write_dict() should be OK");
- let dict_page = Page::DictionaryPage {
- buf: dict,
- num_values: dict_encoder.num_entries() as u32,
- encoding: Encoding::RLE_DICTIONARY,
- is_sorted: false,
- };
- pages.push_front(dict_page);
- }
-}
diff --git a/parquet/src/util/test_common/rand_gen.rs b/parquet/src/util/test_common/rand_gen.rs
deleted file mode 100644
index ea91b28..0000000
--- a/parquet/src/util/test_common/rand_gen.rs
+++ /dev/null
@@ -1,139 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-use rand::{
- distributions::{uniform::SampleUniform, Distribution, Standard},
- thread_rng, Rng,
-};
-
-use crate::data_type::*;
-use crate::util::memory::ByteBufferPtr;
-
-/// Random generator of data type `T` values and sequences.
-pub trait RandGen<T: DataType> {
- fn gen(len: i32) -> T::T;
-
- fn gen_vec(len: i32, total: usize) -> Vec<T::T> {
- let mut result = vec![];
- for _ in 0..total {
- result.push(Self::gen(len))
- }
- result
- }
-}
-
-impl RandGen<BoolType> for BoolType {
- fn gen(_: i32) -> bool {
- thread_rng().gen::<bool>()
- }
-}
-
-impl RandGen<Int32Type> for Int32Type {
- fn gen(_: i32) -> i32 {
- thread_rng().gen::<i32>()
- }
-}
-
-impl RandGen<Int64Type> for Int64Type {
- fn gen(_: i32) -> i64 {
- thread_rng().gen::<i64>()
- }
-}
-
-impl RandGen<Int96Type> for Int96Type {
- fn gen(_: i32) -> Int96 {
- let mut rng = thread_rng();
- let mut result = Int96::new();
- result.set_data(rng.gen::<u32>(), rng.gen::<u32>(), rng.gen::<u32>());
- result
- }
-}
-
-impl RandGen<FloatType> for FloatType {
- fn gen(_: i32) -> f32 {
- thread_rng().gen::<f32>()
- }
-}
-
-impl RandGen<DoubleType> for DoubleType {
- fn gen(_: i32) -> f64 {
- thread_rng().gen::<f64>()
- }
-}
-
-impl RandGen<ByteArrayType> for ByteArrayType {
- fn gen(_: i32) -> ByteArray {
- let mut rng = thread_rng();
- let mut result = ByteArray::new();
- let mut value = vec![];
- let len = rng.gen_range(0..128);
- for _ in 0..len {
- value.push(rng.gen_range(0..255));
- }
- result.set_data(ByteBufferPtr::new(value));
- result
- }
-}
-
-impl RandGen<FixedLenByteArrayType> for FixedLenByteArrayType {
- fn gen(len: i32) -> FixedLenByteArray {
- let mut rng = thread_rng();
- let value_len = if len < 0 {
- rng.gen_range(0..128)
- } else {
- len as usize
- };
- let value = random_bytes(value_len);
- ByteArray::from(value).into()
- }
-}
-
-pub fn random_bytes(n: usize) -> Vec<u8> {
- let mut result = vec![];
- let mut rng = thread_rng();
- for _ in 0..n {
- result.push(rng.gen_range(0..255));
- }
- result
-}
-
-pub fn random_bools(n: usize) -> Vec<bool> {
- let mut result = vec![];
- let mut rng = thread_rng();
- for _ in 0..n {
- result.push(rng.gen::<bool>());
- }
- result
-}
-
-pub fn random_numbers<T>(n: usize) -> Vec<T>
-where
- Standard: Distribution<T>,
-{
- let mut rng = thread_rng();
- Standard.sample_iter(&mut rng).take(n).collect()
-}
-
-pub fn random_numbers_range<T>(n: usize, low: T, high: T, result: &mut Vec<T>)
-where
- T: PartialOrd + SampleUniform + Copy,
-{
- let mut rng = thread_rng();
- for _ in 0..n {
- result.push(rng.gen_range(low..high));
- }
-}
diff --git a/parquet/tests/boolean_writer.rs b/parquet/tests/boolean_writer.rs
deleted file mode 100644
index b9d757e..0000000
--- a/parquet/tests/boolean_writer.rs
+++ /dev/null
@@ -1,100 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-use parquet::column::writer::ColumnWriter;
-use parquet::file::properties::WriterProperties;
-use parquet::file::reader::FileReader;
-use parquet::file::serialized_reader::SerializedFileReader;
-use parquet::file::writer::FileWriter;
-use parquet::file::writer::SerializedFileWriter;
-use parquet::schema::parser::parse_message_type;
-use std::fs;
-use std::path::Path;
-use std::sync::{mpsc, Arc};
-use std::thread;
-use std::time::Duration;
-
-#[test]
-fn it_writes_data_without_hanging() {
- let path = Path::new("it_writes_data_without_hanging.parquet");
-
- let message_type = "
- message BooleanType {
- REQUIRED BOOLEAN DIM0;
- }
-";
- let schema = Arc::new(parse_message_type(message_type).expect("parse schema"));
- let props = Arc::new(WriterProperties::builder().build());
- let file = fs::File::create(&path).expect("create file");
- let mut writer =
- SerializedFileWriter::new(file, schema, props).expect("create parquet writer");
- for _group in 0..1 {
- let mut row_group_writer = writer.next_row_group().expect("get row group writer");
- let values: Vec<i64> = vec![0; 2049];
- let my_bool_values: Vec<bool> = values
- .iter()
- .enumerate()
- .map(|(count, _x)| count % 2 == 0)
- .collect();
- while let Some(mut col_writer) =
- row_group_writer.next_column().expect("next column")
- {
- match col_writer {
- ColumnWriter::BoolColumnWriter(ref mut typed_writer) => {
- typed_writer
- .write_batch(&my_bool_values, None, None)
- .expect("writing bool column");
- }
- _ => {
- panic!("only test boolean values");
- }
- }
- row_group_writer
- .close_column(col_writer)
- .expect("close column");
- }
- let rg_md = row_group_writer.close().expect("close row group");
- println!("total rows written: {}", rg_md.num_rows());
- writer
- .close_row_group(row_group_writer)
- .expect("close row groups");
- }
- writer.close().expect("close writer");
-
- let bytes = fs::read(&path).expect("read file");
- assert_eq!(&bytes[0..4], &[b'P', b'A', b'R', b'1']);
-
- // Now that we have written our data and are happy with it, make
- // sure we can read it back in < 5 seconds...
- let (sender, receiver) = mpsc::channel();
- let _t = thread::spawn(move || {
- let file = fs::File::open(&Path::new("it_writes_data_without_hanging.parquet"))
- .expect("open file");
- let reader = SerializedFileReader::new(file).expect("get serialized reader");
- let iter = reader.get_row_iter(None).expect("get iterator");
- for record in iter {
- println!("reading: {}", record);
- }
- println!("finished reading");
- if let Ok(()) = sender.send(true) {}
- });
- assert_ne!(
- Err(mpsc::RecvTimeoutError::Timeout),
- receiver.recv_timeout(Duration::from_millis(5000))
- );
- fs::remove_file("it_writes_data_without_hanging.parquet").expect("remove file");
-}
diff --git a/parquet/tests/custom_writer.rs b/parquet/tests/custom_writer.rs
deleted file mode 100644
index 0a57e79..0000000
--- a/parquet/tests/custom_writer.rs
+++ /dev/null
@@ -1,100 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-use std::fs::File;
-use std::{
- fs,
- io::{prelude::*, SeekFrom},
- sync::Arc,
-};
-
-use parquet::file::writer::TryClone;
-use parquet::{
- basic::Repetition, basic::Type, file::properties::WriterProperties,
- file::writer::SerializedFileWriter, schema::types,
-};
-use std::env;
-
-// Test creating some sort of custom writer to ensure the
-// appropriate traits are exposed
-struct CustomWriter {
- file: File,
-}
-
-impl Write for CustomWriter {
- fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
- self.file.write(buf)
- }
- fn flush(&mut self) -> std::io::Result<()> {
- self.file.flush()
- }
-}
-
-impl Seek for CustomWriter {
- fn seek(&mut self, pos: SeekFrom) -> std::io::Result<u64> {
- self.file.seek(pos)
- }
-}
-
-impl TryClone for CustomWriter {
- fn try_clone(&self) -> std::io::Result<Self> {
- use std::io::{Error, ErrorKind};
- Err(Error::new(ErrorKind::Other, "Clone not supported"))
- }
-}
-
-#[test]
-fn test_custom_writer() {
- let schema = Arc::new(
- types::Type::group_type_builder("schema")
- .with_fields(&mut vec![Arc::new(
- types::Type::primitive_type_builder("col1", Type::INT32)
- .with_repetition(Repetition::REQUIRED)
- .build()
- .unwrap(),
- )])
- .build()
- .unwrap(),
- );
- let props = Arc::new(WriterProperties::builder().build());
-
- let file = get_temp_file("test_custom_file_writer");
- let test_file = file.try_clone().unwrap();
-
- let writer = CustomWriter { file };
-
- // test is that this file can be created
- let file_writer = SerializedFileWriter::new(writer, schema, props).unwrap();
- std::mem::drop(file_writer);
-
- // ensure the file now exists and has non zero size
- let metadata = test_file.metadata().unwrap();
- assert!(metadata.len() > 0);
-}
-
-/// Returns file handle for a temp file in 'target' directory with a provided content
-fn get_temp_file(file_name: &str) -> fs::File {
- // build tmp path to a file in "target/debug/testdata"
- let mut path_buf = env::current_dir().unwrap();
- path_buf.push("target");
- path_buf.push("debug");
- path_buf.push("testdata");
- fs::create_dir_all(&path_buf).unwrap();
- path_buf.push(file_name);
-
- File::create(path_buf).unwrap()
-}
diff --git a/parquet_derive/Cargo.toml b/parquet_derive/Cargo.toml
deleted file mode 100644
index 1d16137..0000000
--- a/parquet_derive/Cargo.toml
+++ /dev/null
@@ -1,42 +0,0 @@
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-
-[package]
-name = "parquet_derive"
-version = "5.0.0-SNAPSHOT"
-license = "Apache-2.0"
-description = "Derive macros for the Rust implementation of Apache Parquet"
-homepage = "https://github.com/apache/arrow-rs"
-repository = "https://github.com/apache/arrow-rs"
-authors = ["Apache Arrow <dev@arrow.apache.org>"]
-keywords = [ "parquet" ]
-readme = "README.md"
-edition = "2018"
-
-[lib]
-proc-macro = true
-
-[features]
-chrono = []
-bigdecimal = []
-uuid = []
-
-[dependencies]
-proc-macro2 = "1.0"
-quote = "1.0"
-syn = { version = "1.0", features = ["full", "extra-traits"] }
-parquet = { path = "../parquet", version = "5.0.0-SNAPSHOT" }
diff --git a/parquet_derive/README.md b/parquet_derive/README.md
deleted file mode 100644
index 47ac8ae..0000000
--- a/parquet_derive/README.md
+++ /dev/null
@@ -1,98 +0,0 @@
-<!---
- Licensed to the Apache Software Foundation (ASF) under one
- or more contributor license agreements. See the NOTICE file
- distributed with this work for additional information
- regarding copyright ownership. The ASF licenses this file
- to you under the Apache License, Version 2.0 (the
- "License"); you may not use this file except in compliance
- with the License. You may obtain a copy of the License at
-
- http://www.apache.org/licenses/LICENSE-2.0
-
- Unless required by applicable law or agreed to in writing,
- software distributed under the License is distributed on an
- "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
- KIND, either express or implied. See the License for the
- specific language governing permissions and limitations
- under the License.
--->
-
-# Parquet Derive
-
-A crate for deriving `RecordWriter` for arbitrary, _simple_ structs. This does not generate writers for arbitrarily nested
-structures. It only works for primitives and a few generic structures and
-various levels of reference. Please see features checklist for what is currently
-supported.
-
-Derive also has some support for the chrono time library. You must must enable the `chrono` feature to get this support.
-
-## Usage
-Add this to your Cargo.toml:
-```toml
-[dependencies]
-parquet = "4.0.0-SNAPSHOT"
-parquet_derive = "4.0.0-SNAPSHOT"
-```
-
-and this to your crate root:
-```rust
-extern crate parquet;
-#[macro_use] extern crate parquet_derive;
-```
-
-Example usage of deriving a `RecordWriter` for your struct:
-
-```rust
-use parquet;
-use parquet::record::RecordWriter;
-
-#[derive(ParquetRecordWriter)]
-struct ACompleteRecord<'a> {
- pub a_bool: bool,
- pub a_str: &'a str,
- pub a_string: String,
- pub a_borrowed_string: &'a String,
- pub maybe_a_str: Option<&'a str>,
- pub magic_number: i32,
- pub low_quality_pi: f32,
- pub high_quality_pi: f64,
- pub maybe_pi: Option<f32>,
- pub maybe_best_pi: Option<f64>,
-}
-
-// Initialize your parquet file
-let mut writer = SerializedFileWriter::new(file, schema, props).unwrap();
-let mut row_group = writer.next_row_group().unwrap();
-
-// Build up your records
-let chunks = vec![ACompleteRecord{...}];
-
-// The derived `RecordWriter` takes over here
-(&chunks[..]).write_to_row_group(&mut row_group);
-
-writer.close_row_group(row_group).unwrap();
-writer.close().unwrap();
-```
-
-## Features
-- [X] Support writing `String`, `&str`, `bool`, `i32`, `f32`, `f64`, `Vec<u8>`
-- [ ] Support writing dictionaries
-- [X] Support writing logical types like timestamp
-- [X] Derive definition_levels for `Option`
-- [ ] Derive definition levels for nested structures
-- [ ] Derive writing tuple struct
-- [ ] Derive writing `tuple` container types
-
-## Requirements
-- Same as `parquet-rs`
-
-## Test
-Testing a `*_derive` crate requires an intermediate crate. Go to `parquet_derive_test` and run `cargo test` for
-unit tests.
-
-## Docs
-To build documentation, run `cargo doc --no-deps`.
-To compile and view in the browser, run `cargo doc --no-deps --open`.
-
-## License
-Licensed under the Apache License, Version 2.0: http://www.apache.org/licenses/LICENSE-2.0.
diff --git a/parquet_derive/src/lib.rs b/parquet_derive/src/lib.rs
deleted file mode 100644
index 279d0f7..0000000
--- a/parquet_derive/src/lib.rs
+++ /dev/null
@@ -1,126 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-#![recursion_limit = "128"]
-
-extern crate proc_macro;
-extern crate proc_macro2;
-extern crate syn;
-#[macro_use]
-extern crate quote;
-
-extern crate parquet;
-
-use syn::{parse_macro_input, Data, DataStruct, DeriveInput};
-
-mod parquet_field;
-
-/// Derive flat, simple RecordWriter implementations. Works by parsing
-/// a struct tagged with `#[derive(ParquetRecordWriter)]` and emitting
-/// the correct writing code for each field of the struct. Column writers
-/// are generated in the order they are defined.
-///
-/// It is up to the programmer to keep the order of the struct
-/// fields lined up with the schema.
-///
-/// Example:
-///
-/// ```ignore
-/// use parquet;
-/// use parquet::record::RecordWriter;
-/// use parquet::schema::parser::parse_message_type;
-///
-/// use std::sync::Arc;
-//
-/// #[derive(ParquetRecordWriter)]
-/// struct ACompleteRecord<'a> {
-/// pub a_bool: bool,
-/// pub a_str: &'a str,
-/// }
-///
-/// let schema_str = "message schema {
-/// REQUIRED boolean a_bool;
-/// REQUIRED BINARY a_str (UTF8);
-/// }";
-///
-/// pub fn write_some_records() {
-/// let samples = vec![
-/// ACompleteRecord {
-/// a_bool: true,
-/// a_str: "I'm true"
-/// },
-/// ACompleteRecord {
-/// a_bool: false,
-/// a_str: "I'm false"
-/// }
-/// ];
-///
-/// let schema = Arc::new(parse_message_type(schema_str).unwrap());
-///
-/// let props = Arc::new(WriterProperties::builder().build());
-/// let mut writer = SerializedFileWriter::new(file, schema, props).unwrap();
-///
-/// let mut row_group = writer.next_row_group().unwrap();
-/// samples.as_slice().write_to_row_group(&mut row_group).unwrap();
-/// writer.close_row_group(row_group).unwrap();
-/// writer.close().unwrap();
-/// }
-/// ```
-///
-#[proc_macro_derive(ParquetRecordWriter)]
-pub fn parquet_record_writer(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
- let input: DeriveInput = parse_macro_input!(input as DeriveInput);
- let fields = match input.data {
- Data::Struct(DataStruct { fields, .. }) => fields,
- Data::Enum(_) => unimplemented!("Enum currently is not supported"),
- Data::Union(_) => unimplemented!("Union currently is not supported"),
- };
-
- let field_infos: Vec<_> = fields
- .iter()
- .map(|f: &syn::Field| parquet_field::Field::from(f))
- .collect();
-
- let writer_snippets: Vec<proc_macro2::TokenStream> =
- field_infos.iter().map(|x| x.writer_snippet()).collect();
-
- let derived_for = input.ident;
- let generics = input.generics;
-
- (quote! {
- impl#generics RecordWriter<#derived_for#generics> for &[#derived_for#generics] {
- fn write_to_row_group(&self, row_group_writer: &mut Box<parquet::file::writer::RowGroupWriter>) -> Result<(), parquet::errors::ParquetError> {
- let mut row_group_writer = row_group_writer;
- let records = &self; // Used by all the writer snippets to be more clear
-
- #(
- {
- let mut some_column_writer = row_group_writer.next_column().unwrap();
- if let Some(mut column_writer) = some_column_writer {
- #writer_snippets
- row_group_writer.close_column(column_writer)?;
- } else {
- return Err(parquet::errors::ParquetError::General("Failed to get next column".into()))
- }
- }
- );*
-
- Ok(())
- }
- }
- }).into()
-}
diff --git a/parquet_derive/src/parquet_field.rs b/parquet_derive/src/parquet_field.rs
deleted file mode 100644
index 328f4a6..0000000
--- a/parquet_derive/src/parquet_field.rs
+++ /dev/null
@@ -1,920 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-#[derive(Debug, PartialEq)]
-pub struct Field {
- ident: syn::Ident,
- ty: Type,
- is_a_byte_buf: bool,
- third_party_type: Option<ThirdPartyType>,
-}
-
-/// Use third party libraries, detected
-/// at compile time. These libraries will
-/// be written to parquet as their preferred
-/// physical type.
-///
-/// ChronoNaiveDateTime is written as i64
-/// ChronoNaiveDate is written as i32
-#[derive(Debug, PartialEq)]
-enum ThirdPartyType {
- ChronoNaiveDateTime,
- ChronoNaiveDate,
- Uuid,
-}
-
-impl Field {
- pub fn from(f: &syn::Field) -> Self {
- let ty = Type::from(f);
- let is_a_byte_buf = ty.physical_type() == parquet::basic::Type::BYTE_ARRAY;
-
- let third_party_type = match &ty.last_part()[..] {
- "NaiveDateTime" => Some(ThirdPartyType::ChronoNaiveDateTime),
- "NaiveDate" => Some(ThirdPartyType::ChronoNaiveDate),
- "Uuid" => Some(ThirdPartyType::Uuid),
- _ => None,
- };
-
- Field {
- ident: f
- .ident
- .clone()
- .expect("Only structs with named fields are currently supported"),
- ty,
- is_a_byte_buf,
- third_party_type,
- }
- }
-
- /// Takes the parsed field of the struct and emits a valid
- /// column writer snippet. Should match exactly what you
- /// would write by hand.
- ///
- /// Can only generate writers for basic structs, for example:
- ///
- /// struct Record {
- /// a_bool: bool,
- /// maybe_a_bool: Option<bool>
- /// }
- ///
- /// but not
- ///
- /// struct UnsupportedNestedRecord {
- /// a_property: bool,
- /// nested_record: Record
- /// }
- ///
- /// because this parsing logic is not sophisticated enough for definition
- /// levels beyond 2.
- pub fn writer_snippet(&self) -> proc_macro2::TokenStream {
- let ident = &self.ident;
- let column_writer = self.ty.column_writer();
-
- let vals_builder = match &self.ty {
- Type::TypePath(_) => self.copied_direct_vals(),
- Type::Option(ref first_type) => match **first_type {
- Type::TypePath(_) => self.option_into_vals(),
- Type::Reference(_, ref second_type) => match **second_type {
- Type::TypePath(_) => self.option_into_vals(),
- _ => unimplemented!("Unsupported type encountered"),
- },
- ref f => unimplemented!("Unsupported: {:#?}", f),
- },
- Type::Reference(_, ref first_type) => match **first_type {
- Type::TypePath(_) => self.copied_direct_vals(),
- Type::Option(ref second_type) => match **second_type {
- Type::TypePath(_) => self.option_into_vals(),
- Type::Reference(_, ref second_type) => match **second_type {
- Type::TypePath(_) => self.option_into_vals(),
- _ => unimplemented!("Unsupported type encountered"),
- },
- ref f => unimplemented!("Unsupported: {:#?}", f),
- },
- ref f => unimplemented!("Unsupported: {:#?}", f),
- },
- f => unimplemented!("Unsupported: {:#?}", f),
- };
-
- let definition_levels = match &self.ty {
- Type::TypePath(_) => None,
- Type::Option(ref first_type) => match **first_type {
- Type::TypePath(_) => Some(self.optional_definition_levels()),
- Type::Option(_) => unimplemented!("Unsupported nesting encountered"),
- Type::Reference(_, ref second_type)
- | Type::Vec(ref second_type)
- | Type::Array(ref second_type) => match **second_type {
- Type::TypePath(_) => Some(self.optional_definition_levels()),
- _ => unimplemented!("Unsupported nesting encountered"),
- },
- },
- Type::Reference(_, ref first_type)
- | Type::Vec(ref first_type)
- | Type::Array(ref first_type) => match **first_type {
- Type::TypePath(_) => None,
- Type::Reference(_, ref second_type)
- | Type::Vec(ref second_type)
- | Type::Array(ref second_type)
- | Type::Option(ref second_type) => match **second_type {
- Type::TypePath(_) => Some(self.optional_definition_levels()),
- Type::Reference(_, ref third_type) => match **third_type {
- Type::TypePath(_) => Some(self.optional_definition_levels()),
- _ => unimplemented!("Unsupported definition encountered"),
- },
- _ => unimplemented!("Unsupported definition encountered"),
- },
- },
- };
-
- // "vals" is the run of primitive data being written for the column
- // "definition_levels" is a vector of bools which controls whether a value is missing or present
- // this TokenStream is only one part of the code for writing a column and
- // it relies on values calculated in prior code snippets, namely "definition_levels" and "vals_builder".
- // All the context is put together in this functions final quote and
- // this expression just switches between non-nullable and nullable write statements
- let write_batch_expr = if definition_levels.is_some() {
- quote! {
- if let #column_writer(ref mut typed) = column_writer {
- typed.write_batch(&vals[..], Some(&definition_levels[..]), None)?;
- } else {
- panic!("Schema and struct disagree on type for {}", stringify!{#ident})
- }
- }
- } else {
- quote! {
- if let #column_writer(ref mut typed) = column_writer {
- typed.write_batch(&vals[..], None, None)?;
- } else {
- panic!("Schema and struct disagree on type for {}", stringify!{#ident})
- }
- }
- };
-
- quote! {
- {
- #definition_levels
-
- #vals_builder
-
- #write_batch_expr
- }
- }
- }
-
- fn option_into_vals(&self) -> proc_macro2::TokenStream {
- let field_name = &self.ident;
- let is_a_byte_buf = self.is_a_byte_buf;
- let is_a_timestamp =
- self.third_party_type == Some(ThirdPartyType::ChronoNaiveDateTime);
- let is_a_date = self.third_party_type == Some(ThirdPartyType::ChronoNaiveDate);
- let is_a_uuid = self.third_party_type == Some(ThirdPartyType::Uuid);
- let copy_to_vec = !matches!(
- self.ty.physical_type(),
- parquet::basic::Type::BYTE_ARRAY | parquet::basic::Type::FIXED_LEN_BYTE_ARRAY
- );
-
- let binding = if copy_to_vec {
- quote! { let Some(inner) = rec.#field_name }
- } else {
- quote! { let Some(ref inner) = rec.#field_name }
- };
-
- let some = if is_a_timestamp {
- quote! { Some(inner.timestamp_millis()) }
- } else if is_a_date {
- quote! { Some(inner.signed_duration_since(chrono::NaiveDate::from_ymd(1970, 1, 1)).num_days() as i32) }
- } else if is_a_uuid {
- quote! { Some((&inner.to_string()[..]).into()) }
- } else if is_a_byte_buf {
- quote! { Some((&inner[..]).into())}
- } else {
- quote! { Some(inner) }
- };
-
- quote! {
- let vals: Vec<_> = records.iter().filter_map(|rec| {
- if #binding {
- #some
- } else {
- None
- }
- }).collect();
- }
- }
-
- fn copied_direct_vals(&self) -> proc_macro2::TokenStream {
- let field_name = &self.ident;
- let is_a_byte_buf = self.is_a_byte_buf;
- let is_a_timestamp =
- self.third_party_type == Some(ThirdPartyType::ChronoNaiveDateTime);
- let is_a_date = self.third_party_type == Some(ThirdPartyType::ChronoNaiveDate);
- let is_a_uuid = self.third_party_type == Some(ThirdPartyType::Uuid);
-
- let access = if is_a_timestamp {
- quote! { rec.#field_name.timestamp_millis() }
- } else if is_a_date {
- quote! { rec.#field_name.signed_duration_since(chrono::NaiveDate::from_ymd(1970, 1, 1)).num_days() as i32 }
- } else if is_a_uuid {
- quote! { (&rec.#field_name.to_string()[..]).into() }
- } else if is_a_byte_buf {
- quote! { (&rec.#field_name[..]).into() }
- } else {
- quote! { rec.#field_name }
- };
-
- quote! {
- let vals: Vec<_> = records.iter().map(|rec| #access).collect();
- }
- }
-
- fn optional_definition_levels(&self) -> proc_macro2::TokenStream {
- let field_name = &self.ident;
-
- quote! {
- let definition_levels: Vec<i16> = self
- .iter()
- .map(|rec| if rec.#field_name.is_some() { 1 } else { 0 })
- .collect();
- }
- }
-}
-
-#[allow(clippy::enum_variant_names)]
-#[allow(clippy::large_enum_variant)]
-#[derive(Debug, PartialEq)]
-enum Type {
- Array(Box<Type>),
- Option(Box<Type>),
- Vec(Box<Type>),
- TypePath(syn::Type),
- Reference(Option<syn::Lifetime>, Box<Type>),
-}
-
-impl Type {
- /// Takes a rust type and returns the appropriate
- /// parquet-rs column writer
- fn column_writer(&self) -> syn::TypePath {
- use parquet::basic::Type as BasicType;
-
- match self.physical_type() {
- BasicType::BOOLEAN => {
- syn::parse_quote!(parquet::column::writer::ColumnWriter::BoolColumnWriter)
- }
- BasicType::INT32 => syn::parse_quote!(
- parquet::column::writer::ColumnWriter::Int32ColumnWriter
- ),
- BasicType::INT64 => syn::parse_quote!(
- parquet::column::writer::ColumnWriter::Int64ColumnWriter
- ),
- BasicType::INT96 => syn::parse_quote!(
- parquet::column::writer::ColumnWriter::Int96ColumnWriter
- ),
- BasicType::FLOAT => syn::parse_quote!(
- parquet::column::writer::ColumnWriter::FloatColumnWriter
- ),
- BasicType::DOUBLE => syn::parse_quote!(
- parquet::column::writer::ColumnWriter::DoubleColumnWriter
- ),
- BasicType::BYTE_ARRAY => syn::parse_quote!(
- parquet::column::writer::ColumnWriter::ByteArrayColumnWriter
- ),
- BasicType::FIXED_LEN_BYTE_ARRAY => syn::parse_quote!(
- parquet::column::writer::ColumnWriter::FixedLenByteArrayColumnWriter
- ),
- }
- }
-
- /// Helper to simplify a nested field definition to its leaf type
- ///
- /// Ex:
- /// Option<&String> => Type::TypePath(String)
- /// &Option<i32> => Type::TypePath(i32)
- /// Vec<Vec<u8>> => Type::Vec(u8)
- ///
- /// Useful in determining the physical type of a field and the
- /// definition levels.
- fn leaf_type_recursive(&self) -> &Type {
- self.leaf_type_recursive_helper(self, None)
- }
-
- fn leaf_type_recursive_helper<'a>(
- &'a self,
- ty: &'a Type,
- parent_ty: Option<&'a Type>,
- ) -> &Type {
- match ty {
- Type::TypePath(_) => parent_ty.unwrap_or(ty),
- Type::Option(ref first_type)
- | Type::Vec(ref first_type)
- | Type::Array(ref first_type)
- | Type::Reference(_, ref first_type) => {
- self.leaf_type_recursive_helper(first_type, Some(ty))
- }
- }
- }
-
- /// Helper method to further unwrap leaf_type() to get inner-most
- /// type information, useful for determining the physical type
- /// and normalizing the type paths.
- fn inner_type(&self) -> &syn::Type {
- let leaf_type = self.leaf_type_recursive();
-
- match leaf_type {
- Type::TypePath(ref type_) => type_,
- Type::Option(ref first_type)
- | Type::Vec(ref first_type)
- | Type::Array(ref first_type)
- | Type::Reference(_, ref first_type) => match **first_type {
- Type::TypePath(ref type_) => type_,
- _ => unimplemented!("leaf_type() should only return shallow types"),
- },
- }
- }
-
- /// Helper to normalize a type path by extracting the
- /// most identifiable part
- ///
- /// Ex:
- /// std::string::String => String
- /// Vec<u8> => Vec<u8>
- /// chrono::NaiveDateTime => NaiveDateTime
- ///
- /// Does run the risk of mis-identifying a type if import
- /// rename is in play. Please note procedural macros always
- /// run before type resolution so this is a risk the user
- /// takes on when renaming imports.
- fn last_part(&self) -> String {
- let inner_type = self.inner_type();
- let inner_type_str = (quote! { #inner_type }).to_string();
-
- inner_type_str
- .split("::")
- .last()
- .unwrap()
- .trim()
- .to_string()
- }
-
- /// Converts rust types to parquet physical types.
- ///
- /// Ex:
- /// [u8; 10] => FIXED_LEN_BYTE_ARRAY
- /// Vec<u8> => BYTE_ARRAY
- /// String => BYTE_ARRAY
- /// i32 => INT32
- fn physical_type(&self) -> parquet::basic::Type {
- use parquet::basic::Type as BasicType;
-
- let last_part = self.last_part();
- let leaf_type = self.leaf_type_recursive();
-
- match leaf_type {
- Type::Array(ref first_type) => {
- if let Type::TypePath(_) = **first_type {
- if last_part == "u8" {
- return BasicType::FIXED_LEN_BYTE_ARRAY;
- }
- }
- }
- Type::Vec(ref first_type) => {
- if let Type::TypePath(_) = **first_type {
- if last_part == "u8" {
- return BasicType::BYTE_ARRAY;
- }
- }
- }
- _ => (),
- }
-
- match last_part.trim() {
- "bool" => BasicType::BOOLEAN,
- "u8" | "u16" | "u32" => BasicType::INT32,
- "i8" | "i16" | "i32" | "NaiveDate" => BasicType::INT32,
- "u64" | "i64" | "usize" | "NaiveDateTime" => BasicType::INT64,
- "f32" => BasicType::FLOAT,
- "f64" => BasicType::DOUBLE,
- "String" | "str" | "Uuid" => BasicType::BYTE_ARRAY,
- f => unimplemented!("{} currently is not supported", f),
- }
- }
-
- /// Convert a parsed rust field AST in to a more easy to manipulate
- /// parquet_derive::Field
- fn from(f: &syn::Field) -> Self {
- Type::from_type(f, &f.ty)
- }
-
- fn from_type(f: &syn::Field, ty: &syn::Type) -> Self {
- match ty {
- syn::Type::Path(ref p) => Type::from_type_path(f, p),
- syn::Type::Reference(ref tr) => Type::from_type_reference(f, tr),
- syn::Type::Array(ref ta) => Type::from_type_array(f, ta),
- other => unimplemented!(
- "Unable to derive {:?} - it is currently an unsupported type\n{:#?}",
- f.ident.as_ref().unwrap(),
- other
- ),
- }
- }
-
- fn from_type_path(f: &syn::Field, p: &syn::TypePath) -> Self {
- let last_segment = p.path.segments.last().unwrap();
-
- let is_vec =
- last_segment.ident == syn::Ident::new("Vec", proc_macro2::Span::call_site());
- let is_option = last_segment.ident
- == syn::Ident::new("Option", proc_macro2::Span::call_site());
-
- if is_vec || is_option {
- let generic_type = match &last_segment.arguments {
- syn::PathArguments::AngleBracketed(angle_args) => {
- assert_eq!(angle_args.args.len(), 1);
- let first_arg = &angle_args.args[0];
-
- match first_arg {
- syn::GenericArgument::Type(ref typath) => typath.clone(),
- other => unimplemented!("Unsupported: {:#?}", other),
- }
- }
- other => unimplemented!("Unsupported: {:#?}", other),
- };
-
- if is_vec {
- Type::Vec(Box::new(Type::from_type(f, &generic_type)))
- } else {
- Type::Option(Box::new(Type::from_type(f, &generic_type)))
- }
- } else {
- Type::TypePath(syn::Type::Path(p.clone()))
- }
- }
-
- fn from_type_reference(f: &syn::Field, tr: &syn::TypeReference) -> Self {
- let lifetime = tr.lifetime.clone();
- let inner_type = Type::from_type(f, tr.elem.as_ref());
- Type::Reference(lifetime, Box::new(inner_type))
- }
-
- fn from_type_array(f: &syn::Field, ta: &syn::TypeArray) -> Self {
- let inner_type = Type::from_type(f, ta.elem.as_ref());
- Type::Array(Box::new(inner_type))
- }
-}
-
-#[cfg(test)]
-mod test {
- use super::*;
- use syn::{self, Data, DataStruct, DeriveInput};
-
- fn extract_fields(input: proc_macro2::TokenStream) -> Vec<syn::Field> {
- let input: DeriveInput = syn::parse2(input).unwrap();
-
- let fields = match input.data {
- Data::Struct(DataStruct { fields, .. }) => fields,
- _ => panic!("Input must be a struct"),
- };
-
- fields.iter().map(|field| field.to_owned()).collect()
- }
-
- #[test]
- fn test_generating_a_simple_writer_snippet() {
- let snippet: proc_macro2::TokenStream = quote! {
- struct ABoringStruct {
- counter: usize,
- }
- };
-
- let fields = extract_fields(snippet);
- let counter = Field::from(&fields[0]);
-
- let snippet = counter.writer_snippet().to_string();
- assert_eq!(snippet,
- (quote!{
- {
- let vals : Vec < _ > = records . iter ( ) . map ( | rec | rec . counter ) . collect ( );
-
- if let parquet::column::writer::ColumnWriter::Int64ColumnWriter ( ref mut typed ) = column_writer {
- typed . write_batch ( & vals [ .. ] , None , None ) ?;
- } else {
- panic!("Schema and struct disagree on type for {}" , stringify!{ counter } )
- }
- }
- }).to_string()
- )
- }
-
- #[test]
- fn test_optional_to_writer_snippet() {
- let struct_def: proc_macro2::TokenStream = quote! {
- struct StringBorrower<'a> {
- optional_str: Option<&'a str>,
- optional_string: &Option<String>,
- optional_dumb_int: &Option<&i32>,
- }
- };
-
- let fields = extract_fields(struct_def);
-
- let optional = Field::from(&fields[0]);
- let snippet = optional.writer_snippet();
- assert_eq!(snippet.to_string(),
- (quote! {
- {
- let definition_levels : Vec < i16 > = self . iter ( ) . map ( | rec | if rec . optional_str . is_some ( ) { 1 } else { 0 } ) . collect ( ) ;
-
- let vals: Vec <_> = records.iter().filter_map( |rec| {
- if let Some ( ref inner ) = rec . optional_str {
- Some ( (&inner[..]).into() )
- } else {
- None
- }
- }).collect();
-
- if let parquet::column::writer::ColumnWriter::ByteArrayColumnWriter ( ref mut typed ) = column_writer {
- typed . write_batch ( & vals [ .. ] , Some(&definition_levels[..]) , None ) ? ;
- } else {
- panic!("Schema and struct disagree on type for {}" , stringify ! { optional_str } )
- }
- }
- }
- ).to_string());
-
- let optional = Field::from(&fields[1]);
- let snippet = optional.writer_snippet();
- assert_eq!(snippet.to_string(),
- (quote!{
- {
- let definition_levels : Vec < i16 > = self . iter ( ) . map ( | rec | if rec . optional_string . is_some ( ) { 1 } else { 0 } ) . collect ( ) ;
-
- let vals: Vec <_> = records.iter().filter_map( |rec| {
- if let Some ( ref inner ) = rec . optional_string {
- Some ( (&inner[..]).into() )
- } else {
- None
- }
- }).collect();
-
- if let parquet::column::writer::ColumnWriter::ByteArrayColumnWriter ( ref mut typed ) = column_writer {
- typed . write_batch ( & vals [ .. ] , Some(&definition_levels[..]) , None ) ? ;
- } else {
- panic!("Schema and struct disagree on type for {}" , stringify ! { optional_string } )
- }
- }
- }).to_string());
-
- let optional = Field::from(&fields[2]);
- let snippet = optional.writer_snippet();
- assert_eq!(snippet.to_string(),
- (quote!{
- {
- let definition_levels : Vec < i16 > = self . iter ( ) . map ( | rec | if rec . optional_dumb_int . is_some ( ) { 1 } else { 0 } ) . collect ( ) ;
-
- let vals: Vec <_> = records.iter().filter_map( |rec| {
- if let Some ( inner ) = rec . optional_dumb_int {
- Some ( inner )
- } else {
- None
- }
- }).collect();
-
- if let parquet::column::writer::ColumnWriter::Int32ColumnWriter ( ref mut typed ) = column_writer {
- typed . write_batch ( & vals [ .. ] , Some(&definition_levels[..]) , None ) ? ;
- } else {
- panic!("Schema and struct disagree on type for {}" , stringify ! { optional_dumb_int } )
- }
- }
- }).to_string());
- }
-
- #[test]
- fn test_converting_to_column_writer_type() {
- let snippet: proc_macro2::TokenStream = quote! {
- struct ABasicStruct {
- yes_no: bool,
- name: String,
- }
- };
-
- let fields = extract_fields(snippet);
- let processed: Vec<_> = fields.iter().map(|field| Field::from(field)).collect();
-
- let column_writers: Vec<_> = processed
- .iter()
- .map(|field| field.ty.column_writer())
- .collect();
-
- assert_eq!(
- column_writers,
- vec![
- syn::parse_quote!(
- parquet::column::writer::ColumnWriter::BoolColumnWriter
- ),
- syn::parse_quote!(
- parquet::column::writer::ColumnWriter::ByteArrayColumnWriter
- )
- ]
- );
- }
-
- #[test]
- fn convert_basic_struct() {
- let snippet: proc_macro2::TokenStream = quote! {
- struct ABasicStruct {
- yes_no: bool,
- name: String,
- }
- };
-
- let fields = extract_fields(snippet);
- let processed: Vec<_> = fields.iter().map(|field| Field::from(field)).collect();
- assert_eq!(processed.len(), 2);
-
- assert_eq!(
- processed,
- vec![
- Field {
- ident: syn::Ident::new("yes_no", proc_macro2::Span::call_site()),
- ty: Type::TypePath(syn::parse_quote!(bool)),
- is_a_byte_buf: false,
- third_party_type: None,
- },
- Field {
- ident: syn::Ident::new("name", proc_macro2::Span::call_site()),
- ty: Type::TypePath(syn::parse_quote!(String)),
- is_a_byte_buf: true,
- third_party_type: None,
- }
- ]
- )
- }
-
- #[test]
- fn test_get_inner_type() {
- let snippet: proc_macro2::TokenStream = quote! {
- struct LotsOfInnerTypes {
- a_vec: Vec<u8>,
- a_option: std::option::Option<bool>,
- a_silly_string: std::string::String,
- a_complicated_thing: std::option::Option<std::result::Result<(),()>>,
- }
- };
-
- let fields = extract_fields(snippet);
- let converted_fields: Vec<_> =
- fields.iter().map(|field| Type::from(field)).collect();
- let inner_types: Vec<_> = converted_fields
- .iter()
- .map(|field| field.inner_type())
- .collect();
- let inner_types_strs: Vec<_> = inner_types
- .iter()
- .map(|ty| (quote! { #ty }).to_string())
- .collect();
-
- assert_eq!(
- inner_types_strs,
- vec![
- "u8",
- "bool",
- "std :: string :: String",
- "std :: result :: Result < () , () >"
- ]
- )
- }
-
- #[test]
- fn test_physical_type() {
- use parquet::basic::Type as BasicType;
- let snippet: proc_macro2::TokenStream = quote! {
- struct LotsOfInnerTypes {
- a_buf: Vec<u8>,
- a_number: i32,
- a_verbose_option: std::option::Option<bool>,
- a_silly_string: std::string::String,
- a_fix_byte_buf: [u8; 10],
- a_complex_option: Option<&Vec<u8>>,
- a_complex_vec: &Vec<&Option<u8>>,
- }
- };
-
- let fields = extract_fields(snippet);
- let converted_fields: Vec<_> =
- fields.iter().map(|field| Type::from(field)).collect();
- let physical_types: Vec<_> = converted_fields
- .iter()
- .map(|ty| ty.physical_type())
- .collect();
-
- assert_eq!(
- physical_types,
- vec![
- BasicType::BYTE_ARRAY,
- BasicType::INT32,
- BasicType::BOOLEAN,
- BasicType::BYTE_ARRAY,
- BasicType::FIXED_LEN_BYTE_ARRAY,
- BasicType::BYTE_ARRAY,
- BasicType::INT32
- ]
- )
- }
-
- #[test]
- fn test_convert_comprehensive_owned_struct() {
- let snippet: proc_macro2::TokenStream = quote! {
- struct VecHolder {
- a_vec: Vec<u8>,
- a_option: std::option::Option<bool>,
- a_silly_string: std::string::String,
- a_complicated_thing: std::option::Option<std::result::Result<(),()>>,
- }
- };
-
- let fields = extract_fields(snippet);
- let converted_fields: Vec<_> =
- fields.iter().map(|field| Type::from(field)).collect();
-
- assert_eq!(
- converted_fields,
- vec![
- Type::Vec(Box::new(Type::TypePath(syn::parse_quote!(u8)))),
- Type::Option(Box::new(Type::TypePath(syn::parse_quote!(bool)))),
- Type::TypePath(syn::parse_quote!(std::string::String)),
- Type::Option(Box::new(Type::TypePath(
- syn::parse_quote!(std::result::Result<(),()>)
- ))),
- ]
- );
- }
-
- #[test]
- fn test_convert_borrowed_struct() {
- let snippet: proc_macro2::TokenStream = quote! {
- struct Borrower<'a> {
- a_str: &'a str,
- a_borrowed_option: &'a Option<bool>,
- so_many_borrows: &'a Option<&'a str>,
- }
- };
-
- let fields = extract_fields(snippet);
- let types: Vec<_> = fields.iter().map(|field| Type::from(field)).collect();
-
- assert_eq!(
- types,
- vec![
- Type::Reference(
- Some(syn::Lifetime::new("'a", proc_macro2::Span::call_site())),
- Box::new(Type::TypePath(syn::parse_quote!(str)))
- ),
- Type::Reference(
- Some(syn::Lifetime::new("'a", proc_macro2::Span::call_site())),
- Box::new(Type::Option(Box::new(Type::TypePath(syn::parse_quote!(
- bool
- )))))
- ),
- Type::Reference(
- Some(syn::Lifetime::new("'a", proc_macro2::Span::call_site())),
- Box::new(Type::Option(Box::new(Type::Reference(
- Some(syn::Lifetime::new("'a", proc_macro2::Span::call_site())),
- Box::new(Type::TypePath(syn::parse_quote!(str)))
- ))))
- ),
- ]
- );
- }
-
- #[test]
- #[cfg(feature = "chrono")]
- fn test_chrono_timestamp_millis() {
- let snippet: proc_macro2::TokenStream = quote! {
- struct ATimestampStruct {
- henceforth: chrono::NaiveDateTime,
- maybe_happened: Option<&chrono::NaiveDateTime>,
- }
- };
-
- let fields = extract_fields(snippet);
- let when = Field::from(&fields[0]);
- assert_eq!(when.writer_snippet().to_string(),(quote!{
- {
- let vals : Vec<_> = records.iter().map(|rec| rec.henceforth.timestamp_millis() ).collect();
- if let parquet::column::writer::ColumnWriter::Int64ColumnWriter(ref mut typed) = column_writer {
- typed.write_batch(&vals[..], None, None) ?;
- } else {
- panic!("Schema and struct disagree on type for {}" , stringify!{ henceforth })
- }
- }
- }).to_string());
-
- let maybe_happened = Field::from(&fields[1]);
- assert_eq!(maybe_happened.writer_snippet().to_string(),(quote!{
- {
- let definition_levels : Vec<i16> = self.iter().map(|rec| if rec.maybe_happened.is_some() { 1 } else { 0 }).collect();
- let vals : Vec<_> = records.iter().filter_map(|rec| {
- rec.maybe_happened.map(|inner| { inner.timestamp_millis() })
- }).collect();
-
- if let parquet::column::writer::ColumnWriter::Int64ColumnWriter(ref mut typed) = column_writer {
- typed.write_batch(&vals[..], Some(&definition_levels[..]), None) ?;
- } else {
- panic!("Schema and struct disagree on type for {}" , stringify!{ maybe_happened })
- }
- }
- }).to_string());
- }
-
- #[test]
- #[cfg(feature = "chrono")]
- fn test_chrono_date() {
- let snippet: proc_macro2::TokenStream = quote! {
- struct ATimestampStruct {
- henceforth: chrono::NaiveDate,
- maybe_happened: Option<&chrono::NaiveDate>,
- }
- };
-
- let fields = extract_fields(snippet);
- let when = Field::from(&fields[0]);
- assert_eq!(when.writer_snippet().to_string(),(quote!{
- {
- let vals : Vec<_> = records.iter().map(|rec| rec.henceforth.signed_duration_since(chrono::NaiveDate::from_ymd(1970, 1, 1)).num_days() as i32).collect();
- if let parquet::column::writer::ColumnWriter::Int32ColumnWriter(ref mut typed) = column_writer {
- typed.write_batch(&vals[..], None, None) ?;
- } else {
- panic!("Schema and struct disagree on type for {}" , stringify!{ henceforth })
- }
- }
- }).to_string());
-
- let maybe_happened = Field::from(&fields[1]);
- assert_eq!(maybe_happened.writer_snippet().to_string(),(quote!{
- {
- let definition_levels : Vec<i16> = self.iter().map(|rec| if rec.maybe_happened.is_some() { 1 } else { 0 }).collect();
- let vals : Vec<_> = records.iter().filter_map(|rec| {
- rec.maybe_happened.map(|inner| { inner.signed_duration_since(chrono::NaiveDate::from_ymd(1970, 1, 1)).num_days() as i32 })
- }).collect();
-
- if let parquet::column::writer::ColumnWriter::Int32ColumnWriter(ref mut typed) = column_writer {
- typed.write_batch(&vals[..], Some(&definition_levels[..]), None) ?;
- } else {
- panic!("Schema and struct disagree on type for {}" , stringify!{ maybe_happened })
- }
- }
- }).to_string());
- }
-
- #[test]
- #[cfg(feature = "uuid")]
- fn test_uuid() {
- let snippet: proc_macro2::TokenStream = quote! {
- struct ATimestampStruct {
- unique_id: uuid::Uuid,
- maybe_unique_id: Option<&uuid::Uuid>,
- }
- };
-
- let fields = extract_fields(snippet);
- let when = Field::from(&fields[0]);
- assert_eq!(when.writer_snippet().to_string(),(quote!{
- {
- let vals : Vec<_> = records.iter().map(|rec| (&rec.unique_id.to_string()[..]).into() ).collect();
- if let parquet::column::writer::ColumnWriter::ByteArrayColumnWriter(ref mut typed) = column_writer {
- typed.write_batch(&vals[..], None, None) ?;
- } else {
- panic!("Schema and struct disagree on type for {}" , stringify!{ unique_id })
- }
- }
- }).to_string());
-
- let maybe_happened = Field::from(&fields[1]);
- assert_eq!(maybe_happened.writer_snippet().to_string(),(quote!{
- {
- let definition_levels : Vec<i16> = self.iter().map(|rec| if rec.maybe_unique_id.is_some() { 1 } else { 0 }).collect();
- let vals : Vec<_> = records.iter().filter_map(|rec| {
- rec.maybe_unique_id.map(|ref inner| { (&inner.to_string()[..]).into() })
- }).collect();
-
- if let parquet::column::writer::ColumnWriter::ByteArrayColumnWriter(ref mut typed) = column_writer {
- typed.write_batch(&vals[..], Some(&definition_levels[..]), None) ?;
- } else {
- panic!("Schema and struct disagree on type for {}" , stringify!{ maybe_unique_id })
- }
- }
- }).to_string());
- }
-}
diff --git a/parquet_derive_test/Cargo.toml b/parquet_derive_test/Cargo.toml
deleted file mode 100644
index 25c814b..0000000
--- a/parquet_derive_test/Cargo.toml
+++ /dev/null
@@ -1,32 +0,0 @@
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-
-[package]
-name = "parquet_derive_test"
-version = "5.0.0-SNAPSHOT"
-license = "Apache-2.0"
-description = "Integration test package for parquet-derive"
-homepage = "https://github.com/apache/arrow-rs"
-repository = "https://github.com/apache/arrow-rs"
-authors = ["Apache Arrow <dev@arrow.apache.org>"]
-keywords = [ "parquet" ]
-edition = "2018"
-publish = false
-
-[dependencies]
-parquet = { path = "../parquet", version = "5.0.0-SNAPSHOT" }
-parquet_derive = { path = "../parquet_derive", version = "5.0.0-SNAPSHOT" }
diff --git a/parquet_derive_test/src/lib.rs b/parquet_derive_test/src/lib.rs
deleted file mode 100644
index b4bfc42..0000000
--- a/parquet_derive_test/src/lib.rs
+++ /dev/null
@@ -1,131 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-#![allow(clippy::approx_constant)]
-
-extern crate parquet;
-
-#[macro_use]
-extern crate parquet_derive;
-
-use parquet::record::RecordWriter;
-
-#[derive(ParquetRecordWriter)]
-struct ACompleteRecord<'a> {
- pub a_bool: bool,
- pub a_str: &'a str,
- pub a_string: String,
- pub a_borrowed_string: &'a String,
- pub maybe_a_str: Option<&'a str>,
- pub maybe_a_string: Option<String>,
- pub magic_number: i32,
- pub low_quality_pi: f32,
- pub high_quality_pi: f64,
- pub maybe_pi: Option<f32>,
- pub maybe_best_pi: Option<f64>,
- pub borrowed_maybe_a_string: &'a Option<String>,
- pub borrowed_maybe_a_str: &'a Option<&'a str>,
-}
-
-#[cfg(test)]
-mod tests {
- use super::*;
-
- use parquet::{
- file::{
- properties::WriterProperties,
- writer::{FileWriter, SerializedFileWriter},
- },
- schema::parser::parse_message_type,
- };
- use std::{env, fs, io::Write, sync::Arc};
-
- #[test]
- fn test_parquet_derive_hello() {
- let file = get_temp_file("test_parquet_derive_hello", &[]);
- let schema_str = "message schema {
- REQUIRED boolean a_bool;
- REQUIRED BINARY a_str (UTF8);
- REQUIRED BINARY a_string (UTF8);
- REQUIRED BINARY a_borrowed_string (UTF8);
- OPTIONAL BINARY a_maybe_str (UTF8);
- OPTIONAL BINARY a_maybe_string (UTF8);
- REQUIRED INT32 magic_number;
- REQUIRED FLOAT low_quality_pi;
- REQUIRED DOUBLE high_quality_pi;
- OPTIONAL FLOAT maybe_pi;
- OPTIONAL DOUBLE maybe_best_pi;
- OPTIONAL BINARY borrowed_maybe_a_string (UTF8);
- OPTIONAL BINARY borrowed_maybe_a_str (UTF8);
- }";
-
- let schema = Arc::new(parse_message_type(schema_str).unwrap());
-
- let props = Arc::new(WriterProperties::builder().build());
- let mut writer = SerializedFileWriter::new(file, schema, props).unwrap();
-
- let a_str = "hello mother".to_owned();
- let a_borrowed_string = "cool news".to_owned();
- let maybe_a_string = Some("it's true, I'm a string".to_owned());
- let maybe_a_str = Some(&a_str[..]);
-
- let drs: Vec<ACompleteRecord> = vec![ACompleteRecord {
- a_bool: true,
- a_str: &a_str[..],
- a_string: "hello father".into(),
- a_borrowed_string: &a_borrowed_string,
- maybe_a_str: Some(&a_str[..]),
- maybe_a_string: Some(a_str.clone()),
- magic_number: 100,
- low_quality_pi: 3.14,
- high_quality_pi: 3.1415,
- maybe_pi: Some(3.14),
- maybe_best_pi: Some(3.1415),
- borrowed_maybe_a_string: &maybe_a_string,
- borrowed_maybe_a_str: &maybe_a_str,
- }];
-
- let mut row_group = writer.next_row_group().unwrap();
- drs.as_slice().write_to_row_group(&mut row_group).unwrap();
- writer.close_row_group(row_group).unwrap();
- writer.close().unwrap();
- }
-
- /// Returns file handle for a temp file in 'target' directory with a provided content
- pub fn get_temp_file(file_name: &str, content: &[u8]) -> fs::File {
- // build tmp path to a file in "target/debug/testdata"
- let mut path_buf = env::current_dir().unwrap();
- path_buf.push("target");
- path_buf.push("debug");
- path_buf.push("testdata");
- fs::create_dir_all(&path_buf).unwrap();
- path_buf.push(file_name);
-
- // write file content
- let mut tmp_file = fs::File::create(path_buf.as_path()).unwrap();
- tmp_file.write_all(content).unwrap();
- tmp_file.sync_all().unwrap();
-
- // return file handle for both read and write
- let file = fs::OpenOptions::new()
- .read(true)
- .write(true)
- .open(path_buf.as_path());
- assert!(file.is_ok());
- file.unwrap()
- }
-}
diff --git a/pre-commit.sh b/pre-commit.sh
deleted file mode 100755
index 5ce0807..0000000
--- a/pre-commit.sh
+++ /dev/null
@@ -1,88 +0,0 @@
-#!/bin/bash
-
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-
-# This file is git pre-commit hook.
-#
-# Soft link it as git hook under top dir of apache arrow git repository:
-# $ ln -s ../../rust/pre-commit.sh .git/hooks/pre-commit
-#
-# This file be run directly:
-# $ ./pre-commit.sh
-
-function RED() {
- echo "\033[0;31m$@\033[0m"
-}
-
-function GREEN() {
- echo "\033[0;32m$@\033[0m"
-}
-
-function BYELLOW() {
- echo "\033[1;33m$@\033[0m"
-}
-
-RUST_DIR="rust"
-
-# env GIT_DIR is set by git when run a pre-commit hook.
-if [ -z "${GIT_DIR}" ]; then
- GIT_DIR=$(git rev-parse --show-toplevel)
-fi
-
-cd ${GIT_DIR}/${RUST_DIR}
-
-NUM_CHANGES=$(git diff --cached --name-only . |
- grep -e ".*/*.rs$" |
- awk '{print $1}' |
- wc -l)
-
-if [ ${NUM_CHANGES} -eq 0 ]; then
- echo -e "$(GREEN INFO): no staged changes in *.rs, $(GREEN skip cargo fmt/clippy)"
- exit 0
-fi
-
-# 1. cargo clippy
-
-echo -e "$(GREEN INFO): cargo clippy ..."
-
-# Cargo clippy always return exit code 0, and `tee` doesn't work.
-# So let's just run cargo clippy.
-cargo clippy
-echo -e "$(GREEN INFO): cargo clippy done"
-
-# 2. cargo fmt: format with nightly and stable.
-
-CHANGED_BY_CARGO_FMT=false
-echo -e "$(GREEN INFO): cargo fmt with nightly and stable ..."
-
-for version in nightly stable; do
- CMD="cargo +${version} fmt"
- ${CMD} --all -q -- --check 2>/dev/null
- if [ $? -ne 0 ]; then
- ${CMD} --all
- echo -e "$(BYELLOW WARN): ${CMD} changed some files"
- CHANGED_BY_CARGO_FMT=true
- fi
-done
-
-if ${CHANGED_BY_CARGO_FMT}; then
- echo -e "$(RED FAIL): git commit $(RED ABORTED), please have a look and run git add/commit again"
- exit 1
-fi
-
-exit 0
diff --git a/rustfmt.toml b/rustfmt.toml
deleted file mode 100644
index c49cccd..0000000
--- a/rustfmt.toml
+++ /dev/null
@@ -1,24 +0,0 @@
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-
-edition = "2018"
-max_width = 90
-
-# ignore generated files
-# ignore = [
-# "arrow/src/ipc/gen",
-#]
diff --git a/testing b/testing
deleted file mode 160000
index b658b08..0000000
--- a/testing
+++ /dev/null
@@ -1 +0,0 @@
-Subproject commit b658b087767b041b2081766814655b4dd5a9a439
