diff --git a/.asf.yaml b/.asf.yaml index 56980884a4e2..910b13f90931 100644 --- a/.asf.yaml +++ b/.asf.yaml @@ -41,15 +41,47 @@ github: protected_branches: main: required_status_checks: - # require branches to be up-to-date before merging - strict: true - # don't require any jobs to pass - contexts: [] + # The merge queue tests each PR against the latest main. + strict: false + # This job aggregates all suites in .github/workflows/ci.yml. + contexts: + - "Required Checks" required_pull_request_reviews: required_approving_review_count: 1 + # Raw GitHub Rulesets API payload, as in apache/datafusion-comet. + # Rulesets and classic branch protection apply together. + rulesets: + - name: Merge Queue + target: branch + enforcement: active + bypass_actors: + # apache/root (ASF Infra) can recover a blocked queue. + - actor_id: 118420 + actor_type: Team + bypass_mode: always + conditions: + ref_name: + include: + - "~DEFAULT_BRANCH" + exclude: [] + rules: + - type: merge_queue + parameters: + merge_method: SQUASH + grouping_strategy: ALLGREEN + # Limit concurrent CI builds while allowing speculative testing. + max_entries_to_build: 2 + # Merge up to five already-green entries together. Each PR still + # gets its own squash commit; this does not batch CI builds. + min_entries_to_merge: 1 + max_entries_to_merge: 5 + min_entries_to_merge_wait_minutes: 0 + # Allow for the full test suite and ASF runner scheduling delays. + check_response_timeout_minutes: 300 pull_requests: # enable updating head branches of pull requests allow_update_branch: true + allow_auto_merge: true # publishes the content of the `asf-site` branch to # https://arrow.apache.org/rust/ diff --git a/.github/ci/paths.yaml b/.github/ci/paths.yaml new file mode 100644 index 000000000000..f89ca2fe2ce4 --- /dev/null +++ b/.github/ci/paths.yaml @@ -0,0 +1,106 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +# Shared build inputs select all suites on pull requests. +common: + - .asf.yaml + - .github/** + - .gitmodules + - .config/** + - Cargo.toml + - Cargo.lock + - rust-toolchain.toml + - testing + - parquet-testing + - format/** +suites: + arrow: + - arrow-arith/** + - arrow-array/** + - arrow-buffer/** + - arrow-cast/** + - arrow-cmp/** + - arrow-csv/** + - arrow-data/** + - arrow-integration-test/** + - arrow-integration-testing/** + - arrow-ipc/** + - arrow-json/** + - arrow-avro/** + - arrow-ord/** + - arrow-row/** + - arrow-schema/** + - arrow-select/** + - arrow-string/** + - arrow/** + arrow-flight: + - arrow-array/** + - arrow-buffer/** + - arrow-cast/** + - arrow-data/** + - arrow-flight/** + - arrow-ipc/** + - arrow-schema/** + - arrow-select/** + parquet: + - arrow/** + - arrow-arith/** + - arrow-array/** + - arrow-buffer/** + - arrow-cast/** + - arrow-cmp/** + - arrow-csv/** + - arrow-data/** + - arrow-ipc/** + - arrow-json/** + - arrow-ord/** + - arrow-row/** + - arrow-schema/** + - arrow-select/** + - arrow-string/** + - parquet/** + - parquet_derive/** + - parquet_derive_test/** + - parquet-geospatial/** + - parquet-variant/** + - parquet-variant-compute/** + - parquet-variant-json/** + integration: + - arrow-arith/** + - arrow-array/** + - arrow-buffer/** + - arrow-cast/** + - arrow-cmp/** + - arrow-csv/** + - arrow-data/** + - arrow-flight/** + - arrow-integration-test/** + - arrow-integration-testing/** + - arrow-ipc/** + - arrow-json/** + - arrow-ord/** + - arrow-pyarrow-integration-testing/** + - arrow-pyarrow-testing/** + - arrow-pyarrow/** + - arrow-row/** + - arrow-schema/** + - arrow-select/** + - arrow-string/** + - arrow/** + audit: + - '**/Cargo.toml' + - '**/Cargo.lock' diff --git a/.github/ci/scripts/check-ci-config.py b/.github/ci/scripts/check-ci-config.py new file mode 100644 index 000000000000..2648251b4fb4 --- /dev/null +++ b/.github/ci/scripts/check-ci-config.py @@ -0,0 +1,270 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +"""Check CI configuration and routing policy, following Comet's structure. + +Run from the repository root: + python3 .github/ci/scripts/check-ci-config.py +""" + +import contextlib +import importlib.util +import io +import json +import os +from pathlib import Path +import re +import shutil +import subprocess +import sys +import tempfile +import tomllib +import unittest + +spec = importlib.util.spec_from_file_location( + "compute_changes", Path(__file__).with_name("compute-changes.py") +) +changes = importlib.util.module_from_spec(spec) +spec.loader.exec_module(changes) + +ROOT = changes.ROOT +ALWAYS = changes.ALWAYS +load_yaml = changes.load_yaml + +# Only these jobs and shell steps should run when warming main's caches. +CACHE_REFRESH_STEPS = { + "integration": { + "pyarrow-integration-test": { + "Setup Rust toolchain", "Upgrade pip and setuptools", + "Create virtualenv and install dependencies", "Build Python extension", + }, + }, + "parquet": {"pyspark-integration-test": {"Install Python dependencies"}}, +} +FULL_RUN = "${{ !inputs.cache-refresh-only }}" + + +def validate_cache_refresh(ci, suite, workflow): + expected = "${{ needs.changes.outputs.cache-refresh-only == 'true' }}" + assert ci["jobs"][suite].get("with", {}).get("cache-refresh-only") == expected, f"{suite}: missing cache mode input" + cache_jobs = CACHE_REFRESH_STEPS[suite] + assert cache_jobs.keys() <= workflow["jobs"].keys(), f"{suite}: missing cache writer" + for name, job in workflow["jobs"].items(): + if name not in cache_jobs: + assert job.get("if") == FULL_RUN, f"{suite}/{name}: runs during cache refresh" + continue + assert "if" not in job, f"{suite}/{name}: cache writer must run in both modes" + run_steps = {step["name"]: step for step in job["steps"] if "run" in step} + assert cache_jobs[name] <= run_steps.keys(), f"{suite}/{name}: missing cache setup step" + for step in job["steps"]: + if "run" not in step or step["name"] in cache_jobs[name]: + assert "if" not in step, f"{suite}/{name}: cache setup must run in both modes" + else: + assert step.get("if") == FULL_RUN, f"{suite}/{name}/{step['name']}: runs during cache refresh" + + +def validate_package_paths(root, suite, workflow, filters): + """Check literal Cargo package arguments, without interpreting shell scripts.""" + packages = {} + for manifest in root.glob("*/Cargo.toml"): + package = tomllib.loads(manifest.read_text()).get("package") + if package: + packages[package["name"]] = manifest.parent.relative_to(root).as_posix() + for job in workflow["jobs"].values(): + for step in job.get("steps", []): + for command in re.findall(r"\bcargo\b[^\n]*", step.get("run", "")): + for name in re.findall(r"(?:^|\s)(?:-p|--package)(?:\s+|=)[\"']?([\w-]+)", command): + assert name in packages, f"{suite}: unknown local package {name}" + path = f"{packages[name]}/src/lib.rs" + selected = changes.select_suites("pull_request", [path], filters) + assert suite in selected, f"{suite}: changes to {name} do not select its tests" + + +def validate_config(root=ROOT): + """Catch stale check names, omitted suites, and workflow-level path filters.""" + workflows = root / ".github/workflows" + ci = load_yaml(workflows / "ci.yml") + filters = load_yaml(root / ".github/ci/paths.yaml") + jobs = ci["jobs"] + suites = set(filters["suites"]) | ALWAYS | {"miri"} + assert not (set(filters["suites"]) & (ALWAYS | {"miri"})), "Unexpected filtered suite" + assert set(jobs) == suites | {"changes", "required-checks"}, "CI suites and jobs differ" + # PyYAML's YAML 1.1 loader treats the key 'on' as the boolean True. + events = ci[True] + assert jobs["changes"]["outputs"].get("cache-refresh-only") == "${{ steps.select.outputs.cache-refresh-only }}", "Missing cache mode output" + assert set(events) == {"pull_request", "push", "merge_group"}, "Unexpected CI triggers" + assert events["pull_request"] is None, "Required CI must run on every PR" + assert events["merge_group"] == {"types": ["checks_requested"]}, "Missing queue trigger" + assert events["push"] == {"branches": ["main"]}, "Only main pushes should run CI" + + gate = jobs["required-checks"] + assert set(gate["needs"]) == set(jobs) - {"required-checks"}, "Required Checks omits a job" + assert gate["if"] == "always()", "Required Checks must run after failures and cancellations" + asf = load_yaml(root / ".asf.yaml")["github"] + checks = asf["protected_branches"]["main"]["required_status_checks"] + assert checks["contexts"] == [gate["name"]], "Required check name does not match .asf.yaml" + assert checks["strict"] is False, "The queue handles up-to-date checks" + + called = set() + for suite in suites: + job = jobs[suite] + assert job["needs"] == "changes", f"{suite}: missing suite selection dependency" + condition = f"contains(fromJSON(needs.changes.outputs.suites), '{suite}')" + assert job["if"] == condition, f"{suite}: unexpected routing condition" + path = root / job["uses"] + workflow = load_yaml(path) + call = None + if suite in CACHE_REFRESH_STEPS: + call = {"inputs": {"cache-refresh-only": { + "description": "Populate caches without running tests or linters", + "type": "boolean", "default": False, + }}} + validate_cache_refresh(ci, suite, workflow) + assert workflow[True] == {"workflow_call": call}, f"{suite}: duplicate or missing triggers/inputs" + assert "concurrency" not in workflow, f"{suite}: concurrency is owned by ci.yml" + if suite in filters["suites"]: + validate_package_paths(root, suite, workflow, filters) + called.add(path.resolve()) + reusable = { + path.resolve() + for path in workflows.iterdir() + if path.suffix in {".yml", ".yaml"} + and "workflow_call" in load_yaml(path).get(True, {}) + } + assert called == reusable, "A reusable workflow is missing from Required Checks" + miri = load_yaml(workflows / "miri.yaml")["jobs"]["miri-checks"] + assert "if" not in miri, "Miri event routing belongs in compute-changes.py" + print("CI configuration is valid") + + +class RoutingTests(unittest.TestCase): + def setUp(self): + self.filters = changes.load_yaml(changes.ROOT / ".github/ci/paths.yaml") + self.always = {"dev", "rust", "docs"} + self.all_suites = self.always | { + "arrow", "arrow-flight", "parquet", "integration", "audit", + } + + def select(self, event, paths): + return set(changes.select_suites(event, paths, self.filters)) + + def test_event_policy(self): + for paths in [[], ["README.md"], ["arrow-buffer/src/lib.rs"]]: + with self.subTest(paths=paths): + self.assertEqual(self.select("merge_group", paths), self.all_suites | {"miri"}) + self.assertEqual(self.select("push", paths), {"docs", "integration", "parquet"}) + self.assertNotIn("miri", self.select("pull_request", paths)) + with self.assertRaises(ValueError): + self.select("pull_request_target", []) + + def test_selector_outputs(self): + for event, expected in [("push", {"docs", "integration", "parquet"}), + ("merge_group", self.all_suites | {"miri"}), + ("pull_request", self.always)]: + with self.subTest(event=event), tempfile.TemporaryDirectory() as directory: + output = Path(directory) / "output" + env = dict(os.environ, GITHUB_EVENT_NAME=event, GITHUB_OUTPUT=str(output), + PR_BASE_SHA="HEAD", PR_HEAD_SHA="HEAD") + subprocess.run([sys.executable, str(ROOT / ".github/ci/scripts/compute-changes.py")], + env=env, check=True, capture_output=True) + values = dict(line.split("=", 1) for line in output.read_text().splitlines()) + self.assertEqual(set(json.loads(values["suites"])), expected) + self.assertEqual(json.loads(values["cache-refresh-only"]), event == "push") + + def test_pr_path_routing(self): + cases = [ + ([], set()), + (["README.md"], set()), + (["arrow-buffer/src/lib.rs"], {"arrow", "arrow-flight", "parquet", "integration"}), + (["arrow-avro/src/lib.rs"], {"arrow"}), + (["arrow-flight/src/lib.rs"], {"arrow-flight", "integration"}), + (["arrow-pyarrow/src/lib.rs"], {"integration"}), + (["arrow-integration-testing/src/lib.rs"], {"arrow", "integration"}), + (["arrow-pyarrow-testing/src/lib.rs"], {"integration"}), + (["parquet/src/lib.rs"], {"parquet"}), + (["parquet_derive/src/lib.rs"], {"parquet"}), + (["parquet_derive_test/src/lib.rs"], {"parquet"}), + (["parquet-variant-json/src/lib.rs"], {"parquet"}), + (["parquet-geospatial/src/lib.rs"], {"parquet"}), + (["arrow-row/Cargo.toml"], {"arrow", "parquet", "integration", "audit"}), + (["some/nested/Cargo.lock"], {"audit"}), + (["parquet-variant/src/lib.rs"], {"parquet"}), + (["parquet-variant-compute/src/lib.rs"], {"parquet"}), + (["arrow-arith/src/lib.rs"], {"arrow", "parquet", "integration"}), + (["arrow-ord/src/lib.rs"], {"arrow", "parquet", "integration"}), + (["arrow-string/src/lib.rs"], {"arrow", "parquet", "integration"}), + # A move or deletion must still select the affected suites. + (["arrow-cmp/old.rs", "parquet-geospatial/new.rs"], {"arrow", "parquet", "integration"}), + (["arrow-cmp/file\nwith-newline.rs"], {"arrow", "parquet", "integration"}), + ] + for paths, expected in cases: + with self.subTest(paths=paths): + self.assertEqual(self.select("pull_request", paths), self.always | expected) + + def test_shared_build_inputs(self): + for path in [ + ".asf.yaml", ".github/workflows/ci.yml", ".github/ci/paths.yaml", + ".github/actions/setup-builder/action.yaml", ".gitmodules", + "Cargo.toml", "Cargo.lock", "rust-toolchain.toml", + "testing", "parquet-testing", "format/Flight.proto", "format/FlightSql.proto", + ".config/nextest.toml", + ]: + with self.subTest(path=path): + self.assertEqual(self.select("pull_request", [path]), self.all_suites) + + +class ConfigurationTests(unittest.TestCase): + def test_current_configuration(self): + with contextlib.redirect_stdout(io.StringIO()): + validate_config() + + def test_configuration_regressions(self): + cases = [ + (".asf.yaml", '"Required Checks"', '"Missing Check"'), + (".github/workflows/ci.yml", " - miri\n", ""), + (".github/workflows/ci.yml", " pull_request:\n", " pull_request:\n paths: ['arrow/**']\n"), + (".github/workflows/ci.yml", " if: always()", " if: success()"), + (".github/workflows/miri.yaml", " name: MIRI\n", " name: MIRI\n if: github.event_name == 'merge_group'\n"), + (".github/workflows/arrow.yml", " workflow_call:", " workflow_call:\n pull_request:"), + (".github/ci/paths.yaml", " - arrow-integration-testing/**\n", ""), + (".github/workflows/ci.yml", " cache-refresh-only: ${{ steps.select.outputs.cache-refresh-only }}\n", ""), + (".github/workflows/ci.yml", " with:\n cache-refresh-only: ${{ needs.changes.outputs.cache-refresh-only == 'true' }}\n", ""), + (".github/workflows/parquet.yml", "\n if: ${{ !inputs.cache-refresh-only }}\n", "\n"), + (".github/workflows/integration.yml", "\n if: ${{ !inputs.cache-refresh-only }}\n", "\n"), + (".github/workflows/parquet.yml", " if: ${{ !inputs.cache-refresh-only }}\n", ""), + (".github/workflows/integration.yml", " if: ${{ !inputs.cache-refresh-only }}\n", ""), + (".github/workflows/integration.yml", " - name: Build Python extension\n", " - name: Build Python extension\n if: ${{ !inputs.cache-refresh-only }}\n"), + ] + for filename, before, after in cases: + with self.subTest(filename=filename, before=before), tempfile.TemporaryDirectory() as directory: + root = Path(directory) + shutil.copytree(changes.ROOT / ".github", root / ".github") + shutil.copy(changes.ROOT / ".asf.yaml", root / ".asf.yaml") + for manifest in changes.ROOT.glob("*/Cargo.toml"): + target = root / manifest.relative_to(changes.ROOT) + target.parent.mkdir(parents=True, exist_ok=True) + shutil.copy(manifest, target) + path = root / filename + original = path.read_text() + self.assertIn(before, original) + path.write_text(original.replace(before, after)) + with self.assertRaises(AssertionError): + validate_config(root) + + +if __name__ == "__main__": + unittest.main() diff --git a/.github/ci/scripts/compute-changes.py b/.github/ci/scripts/compute-changes.py new file mode 100644 index 000000000000..0b8706206c5d --- /dev/null +++ b/.github/ci/scripts/compute-changes.py @@ -0,0 +1,76 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +"""Select CI suites from the event and changed PR paths.""" + +import fnmatch +import json +import os +from pathlib import Path +import subprocess + +ROOT = Path(__file__).resolve().parents[3] +ALWAYS = {"dev", "rust", "docs"} + + +def load_yaml(path): + import yaml + + return yaml.safe_load(path.read_text()) + + +def select_suites(event, paths, filters): + suites = set(filters["suites"]) | ALWAYS + if event == "merge_group": + return sorted(suites | {"miri"}) + if event == "push": + # The queue already tested the merge result. Publish docs and populate + # main's caches through the reusable workflows' cache-refresh-only mode. + return ["docs", "integration", "parquet"] + if event != "pull_request": + raise ValueError(f"Unsupported CI event: {event}") + + def matches(patterns): + return any(fnmatch.fnmatchcase(path, pattern) for path in paths for pattern in patterns) + + if matches(filters["common"]): + return sorted(suites) + return sorted(ALWAYS | {name for name, patterns in filters["suites"].items() if matches(patterns)}) + + +def main(): + event = os.environ["GITHUB_EVENT_NAME"] + paths = [] + if event == "pull_request": + base = os.environ["PR_BASE_SHA"] + head = os.environ["PR_HEAD_SHA"] + # Full checkout history includes both PR parents. Disable rename + # detection so moving a file selects both its old and new suites. + changed = subprocess.check_output( + ["git", "diff", "--no-renames", "--name-only", "-z", f"{base}...{head}", "--"], + cwd=ROOT, + ) + paths = changed.decode("utf-8", errors="surrogateescape").split("\0")[:-1] + suites = select_suites(event, paths, load_yaml(ROOT / ".github/ci/paths.yaml")) + with open(os.environ["GITHUB_OUTPUT"], "a") as output: + output.write(f"suites={json.dumps(suites)}\n") + output.write(f"cache-refresh-only={json.dumps(event == 'push')}\n") + print(f"Selected suites: {', '.join(suites)}") + + +if __name__ == "__main__": + main() diff --git a/.github/pull_request_template.md b/.github/pull_request_template.md index 7daafd5b38df..c3720f6cc704 100644 --- a/.github/pull_request_template.md +++ b/.github/pull_request_template.md @@ -1,22 +1,32 @@ # Which issue does this PR close? -- Closes #NNN. +- Closes #NNN # Rationale for this change # What changes are included in this PR? # Are these changes tested? @@ -26,15 +36,27 @@ We typically require tests for all PRs in order to: 1. Prevent the code from being accidentally broken by subsequent changes 2. Serve as another way to document the expected behavior of the code -If tests are not included in your PR, please explain why (for example, are they covered by existing tests)? +If tests are not included in your PR, please explain why (for example, are they +covered by existing tests, if so identify these tests)? -If this PR claims a performance improvement, please include evidence such as benchmark results. +If this PR claims a performance improvement, please include evidence such as +benchmark results. + +DO NOT simply paste a list of commands ran to verify the code; this is what our +CI is for and it pollutes the PR. ONLY explain if testing was included with this +PR, or why it wasn't. --> -# Are there any user-facing changes? +# Are there any significant user-facing changes? diff --git a/.github/workflows/README.md b/.github/workflows/README.md index cfdd6d143c5b..d5c5d915c6cf 100644 --- a/.github/workflows/README.md +++ b/.github/workflows/README.md @@ -17,10 +17,7 @@ under the License. --> -The CI is structured so most tests are run in specific workflows: -`arrow.yml` for `arrow`, `parquet.yml` for `parquet` and the three -`parquet-variant*` crates, and so on. - -The basic idea is to run all tests on pushes to main (to ensure we -keep main green) but run only the individual workflows on PRs that -change files that could affect them. +CI runs a subset of checks on pull requests, selected by the changed files. +The merge queue runs the full test suite against the proposed merge result. +Pushes to `main` populate caches and publish documentation without repeating tests. +A single required status check combines the results and blocks merging on failure. diff --git a/.github/workflows/arrow.yml b/.github/workflows/arrow.yml index 71d4bda10f9f..a23ac9c21406 100644 --- a/.github/workflows/arrow.yml +++ b/.github/workflows/arrow.yml @@ -18,36 +18,9 @@ # tests for arrow crate name: arrow -concurrency: - group: ${{ github.repository }}-${{ github.head_ref || github.sha }}-${{ github.workflow }} - cancel-in-progress: true - +# Triggered by ci.yml, which owns event routing and concurrency. on: - # always trigger - push: - branches: - - main - pull_request: - paths: - - .github/** - - rust-toolchain.toml - - arrow-arith/** - - arrow-array/** - - arrow-buffer/** - - arrow-cast/** - - arrow-cmp/** - - arrow-csv/** - - arrow-data/** - - arrow-integration-test/** - - arrow-ipc/** - - arrow-json/** - - arrow-avro/** - - arrow-ord/** - - arrow-row/** - - arrow-schema/** - - arrow-select/** - - arrow-string/** - - arrow/** + workflow_call: permissions: contents: read diff --git a/.github/workflows/arrow_flight.yml b/.github/workflows/arrow_flight.yml index 1e158dabe302..f11b0d1b4138 100644 --- a/.github/workflows/arrow_flight.yml +++ b/.github/workflows/arrow_flight.yml @@ -19,27 +19,9 @@ # tests for arrow_flight crate name: arrow_flight -concurrency: - group: ${{ github.repository }}-${{ github.head_ref || github.sha }}-${{ github.workflow }} - cancel-in-progress: true - -# trigger for all PRs that touch certain files and changes to main +# Triggered by ci.yml, which owns event routing and concurrency. on: - push: - branches: - - main - pull_request: - paths: - - rust-toolchain.toml - - arrow-array/** - - arrow-buffer/** - - arrow-cast/** - - arrow-data/** - - arrow-flight/** - - arrow-ipc/** - - arrow-schema/** - - arrow-select/** - - .github/** + workflow_call: permissions: contents: read diff --git a/.github/workflows/audit.yml b/.github/workflows/audit.yml index 1f32fbfe4645..7a6b183d219d 100644 --- a/.github/workflows/audit.yml +++ b/.github/workflows/audit.yml @@ -17,20 +17,9 @@ name: audit -concurrency: - group: ${{ github.repository }}-${{ github.head_ref || github.sha }}-${{ github.workflow }} - cancel-in-progress: true - -# trigger for all PRs that touch certain files and changes to main +# Triggered by ci.yml, which owns event routing and concurrency. on: - push: - branches: - - main - pull_request: - paths: - - '**/Cargo.toml' - - '**/Cargo.lock' - - rust-toolchain.toml + workflow_call: permissions: contents: read @@ -42,7 +31,7 @@ jobs: steps: - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 - name: Install cargo-audit - uses: taiki-e/install-action@9114bf4d891761788c546334fd37538eae1bf8b3 # v2.87.16 + uses: taiki-e/install-action@4cef1412cce204788f482e778a0b9187f9626a29 # v2.87.21 with: tool: cargo-audit - name: Run audit check diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml new file mode 100644 index 000000000000..dcd624568c49 --- /dev/null +++ b/.github/workflows/ci.yml @@ -0,0 +1,147 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT 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: CI + +on: + pull_request: + merge_group: + types: [checks_requested] + push: + branches: [main] + +concurrency: + group: ${{ github.repository }}-${{ github.event_name }}-${{ github.event.pull_request.number || github.sha }}-${{ github.workflow }} + cancel-in-progress: true + +permissions: + contents: read + +jobs: + changes: + name: Select and validate CI suites + runs-on: ubuntu-latest + outputs: + suites: ${{ steps.select.outputs.suites }} + cache-refresh-only: ${{ steps.select.outputs.cache-refresh-only }} + steps: + - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 + with: + # PR routing uses the merge base, including changes from all commits. + fetch-depth: 0 + persist-credentials: false + - uses: actions/setup-python@5fda3b95a4ea91299a34e894583c3862153e4b97 # v7.0.0 + with: + python-version: '3.12' + - name: Install CI configuration dependency + run: | + for attempt in 1 2 3; do + python3 -m pip install PyYAML==6.0.3 && break + if [ "$attempt" -eq 3 ]; then exit 1; fi + sleep 5 + done + - name: Validate CI configuration and routing + run: python3 .github/ci/scripts/check-ci-config.py + - name: Select suites + id: select + env: + PR_BASE_SHA: ${{ github.event.pull_request.base.sha }} + PR_HEAD_SHA: ${{ github.event.pull_request.head.sha }} + run: python3 .github/ci/scripts/compute-changes.py + + arrow: + needs: changes + if: contains(fromJSON(needs.changes.outputs.suites), 'arrow') + uses: ./.github/workflows/arrow.yml + + arrow-flight: + needs: changes + if: contains(fromJSON(needs.changes.outputs.suites), 'arrow-flight') + uses: ./.github/workflows/arrow_flight.yml + + parquet: + needs: changes + if: contains(fromJSON(needs.changes.outputs.suites), 'parquet') + uses: ./.github/workflows/parquet.yml + with: + cache-refresh-only: ${{ needs.changes.outputs.cache-refresh-only == 'true' }} + + integration: + needs: changes + if: contains(fromJSON(needs.changes.outputs.suites), 'integration') + uses: ./.github/workflows/integration.yml + with: + cache-refresh-only: ${{ needs.changes.outputs.cache-refresh-only == 'true' }} + + audit: + needs: changes + if: contains(fromJSON(needs.changes.outputs.suites), 'audit') + uses: ./.github/workflows/audit.yml + + dev: + needs: changes + if: contains(fromJSON(needs.changes.outputs.suites), 'dev') + uses: ./.github/workflows/dev.yml + + rust: + needs: changes + if: contains(fromJSON(needs.changes.outputs.suites), 'rust') + uses: ./.github/workflows/rust.yml + + docs: + needs: changes + if: contains(fromJSON(needs.changes.outputs.suites), 'docs') + # Only the deploy job in docs.yml requests write access, on push to main. + permissions: + contents: write + uses: ./.github/workflows/docs.yml + + miri: + needs: changes + if: contains(fromJSON(needs.changes.outputs.suites), 'miri') + uses: ./.github/workflows/miri.yaml + + required-checks: + # Like datafusion-comet: one stable check name even when a reusable + # workflow is skipped (its individual check names would be different). + name: Required Checks + if: always() + permissions: {} + needs: + - changes + - arrow + - arrow-flight + - parquet + - integration + - audit + - dev + - rust + - docs + - miri + runs-on: ubuntu-slim + steps: + - name: Summarize suite results + env: + NEEDS: ${{ toJSON(needs) }} + run: echo "$NEEDS" + # Path-filtered suites and Miri outside the queue are intentionally + # skipped. A failed/cancelled selector or suite must still block merging. + - name: Fail if any suite failed or was cancelled + if: needs.changes.result != 'success' || contains(needs.*.result, 'failure') || contains(needs.*.result, 'cancelled') + run: | + echo "::error::One or more CI suites did not succeed. See the results above." + exit 1 diff --git a/.github/workflows/codeql.yml b/.github/workflows/codeql.yml index e8d9d3ea2dc5..f00e353b0298 100644 --- a/.github/workflows/codeql.yml +++ b/.github/workflows/codeql.yml @@ -45,11 +45,11 @@ jobs: persist-credentials: false - name: Initialize CodeQL - uses: github/codeql-action/init@1c5b675653bb5c22dbe9b12b556ec555138e09fd # v4.38.1 + uses: github/codeql-action/init@2892aa5e19bbd11bc0cff5427e3b750a04d9e3c2 # v4.38.2 with: languages: actions - name: Perform CodeQL Analysis - uses: github/codeql-action/analyze@1c5b675653bb5c22dbe9b12b556ec555138e09fd # v4.38.1 + uses: github/codeql-action/analyze@2892aa5e19bbd11bc0cff5427e3b750a04d9e3c2 # v4.38.2 with: category: "/language:actions" diff --git a/.github/workflows/dev.yml b/.github/workflows/dev.yml index f4c6fe24d9ca..786f3ded905a 100644 --- a/.github/workflows/dev.yml +++ b/.github/workflows/dev.yml @@ -17,16 +17,9 @@ name: dev -concurrency: - group: ${{ github.repository }}-${{ github.head_ref || github.sha }}-${{ github.workflow }} - cancel-in-progress: true - -# trigger for all PRs and changes to main +# Triggered by ci.yml, which owns event routing and concurrency. on: - push: - branches: - - main - pull_request: + workflow_call: permissions: contents: read @@ -68,7 +61,7 @@ jobs: with: persist-credentials: false - name: Install typos - uses: taiki-e/install-action@9114bf4d891761788c546334fd37538eae1bf8b3 # v2.87.16 + uses: taiki-e/install-action@4cef1412cce204788f482e778a0b9187f9626a29 # v2.87.21 with: tool: typos@1.37.0 - name: Run typos check diff --git a/.github/workflows/docs.yml b/.github/workflows/docs.yml index 54f5387aa376..8b8753b15b7b 100644 --- a/.github/workflows/docs.yml +++ b/.github/workflows/docs.yml @@ -17,16 +17,9 @@ name: docs -concurrency: - group: ${{ github.repository }}-${{ github.head_ref || github.sha }}-${{ github.workflow }} - cancel-in-progress: true - -# trigger for all PRs and changes to main +# Triggered by ci.yml, which owns event routing and concurrency. on: - push: - branches: - - main - pull_request: + workflow_call: jobs: # test doc links still work diff --git a/.github/workflows/integration.yml b/.github/workflows/integration.yml index 62a4b19f582e..658081db8985 100644 --- a/.github/workflows/integration.yml +++ b/.github/workflows/integration.yml @@ -17,45 +17,21 @@ name: integration -concurrency: - group: ${{ github.repository }}-${{ github.head_ref || github.sha }}-${{ github.workflow }} - cancel-in-progress: true - -# trigger for all PRs that touch certain files and changes to main +# Triggered by ci.yml, which owns event routing and concurrency. on: - push: - branches: - - main - pull_request: - paths: - - .github/** - - rust-toolchain.toml - - arrow-arith/** - - arrow-array/** - - arrow-buffer/** - - arrow-cast/** - - arrow-cmp/** - - arrow-csv/** - - arrow-data/** - - arrow-flight/** - - arrow-integration-test/** - - arrow-integration-testing/** - - arrow-ipc/** - - arrow-json/** - - arrow-ord/** - - arrow-pyarrow-integration-testing/** - - arrow-pyarrow/** - - arrow-row/** - - arrow-schema/** - - arrow-select/** - - arrow-string/** - - arrow/** + workflow_call: + inputs: + cache-refresh-only: + description: Populate caches without running tests or linters + type: boolean + default: false permissions: contents: read jobs: integration: + if: ${{ !inputs.cache-refresh-only }} name: Archery test With other arrows runs-on: ubuntu-latest container: @@ -198,7 +174,8 @@ jobs: matrix: rust: [stable] # PyArrow 15 was the first version to introduce StringView/BinaryView support - pyarrow: ["15", "16", "17"] + # The cache keys are shared across these versions; one build warms them. + pyarrow: ${{ inputs.cache-refresh-only && fromJSON('["17"]') || fromJSON('["15", "16", "17"]') }} steps: - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 with: @@ -212,13 +189,15 @@ jobs: uses: actions/cache@55cc8345863c7cc4c66a329aec7e433d2d1c52a9 # v6.1.0 with: path: /home/runner/.cargo - key: cargo-maturin-cache- + key: cargo-maturin-cache-${{ hashFiles('**/Cargo.toml', 'rust-toolchain.toml') }} + restore-keys: cargo-maturin-cache- - name: Cache Rust dependencies uses: actions/cache@55cc8345863c7cc4c66a329aec7e433d2d1c52a9 # v6.1.0 with: - path: /home/runner/target + path: arrow-pyarrow-integration-testing/target # this key is not equal because maturin uses different compilation flags. - key: ${{ runner.os }}-${{ matrix.arch }}-target-maturin-cache-${{ matrix.rust }}- + key: ${{ runner.os }}-${{ runner.arch }}-target-maturin-cache-${{ matrix.rust }}-${{ hashFiles('**/Cargo.toml', 'rust-toolchain.toml') }} + restore-keys: ${{ runner.os }}-${{ runner.arch }}-target-maturin-cache-${{ matrix.rust }}- - uses: actions/setup-python@5fda3b95a4ea91299a34e894583c3862153e4b97 # v7.0.0 with: python-version: '3.8' @@ -230,6 +209,7 @@ jobs: source venv/bin/activate pip install maturin toml pytest pytz pyarrow==${{ matrix.pyarrow }} - name: Run Rust tests + if: ${{ !inputs.cache-refresh-only }} run: | source venv/bin/activate # arrow-pyarrow's own tests are feature-gated, so the workspace-wide `cargo test` in @@ -239,6 +219,7 @@ jobs: cd arrow-pyarrow-testing cargo test - name: Clippy arrow-pyarrow + if: ${{ !inputs.cache-refresh-only }} run: | source venv/bin/activate # Here rather than in arrow.yml's clippy job: that one runs in a @@ -248,9 +229,14 @@ jobs: # default features, so the --no-default-features variant arrow.yml # also runs would be this same command. cargo clippy -p arrow-pyarrow --all-features -- -D unused_crate_dependencies - - name: Run Python tests + - name: Build Python extension run: | source venv/bin/activate cd arrow-pyarrow-integration-testing maturin develop + - name: Run Python tests + if: ${{ !inputs.cache-refresh-only }} + run: | + source venv/bin/activate + cd arrow-pyarrow-integration-testing pytest -v . diff --git a/.github/workflows/miri.yaml b/.github/workflows/miri.yaml index 5f8483443cc4..3e57aa142d6e 100644 --- a/.github/workflows/miri.yaml +++ b/.github/workflows/miri.yaml @@ -17,39 +17,9 @@ name: miri -concurrency: - group: ${{ github.repository }}-${{ github.head_ref || github.sha }}-${{ github.workflow }} - cancel-in-progress: true - -# trigger for all PRs that touch certain files and changes to main +# Triggered by ci.yml, which owns event routing and concurrency. on: - push: - branches: - - main - pull_request: - paths: - - .config/nextest.toml - - Cargo.toml - - Cargo.lock - - .github/** - - rust-toolchain.toml - - arrow/** - - arrow-arith/** - - arrow-array/** - - arrow-avro/** - - arrow-buffer/** - - arrow-cast/** - - arrow-csv/** - - arrow-data/** - - arrow-ipc/** - - arrow-json/** - - arrow-cmp/** - - arrow-ord/** - - arrow-row/** - - arrow-schema/** - - arrow-select/** - - arrow-string/** - - parquet/** + workflow_call: permissions: contents: read @@ -72,7 +42,7 @@ jobs: rustup override set nightly cargo miri setup - name: Set up nextest - uses: taiki-e/install-action@9114bf4d891761788c546334fd37538eae1bf8b3 # v2.87.16 + uses: taiki-e/install-action@4cef1412cce204788f482e778a0b9187f9626a29 # v2.87.21 with: tool: nextest - name: Run Miri Checks diff --git a/.github/workflows/parquet.yml b/.github/workflows/parquet.yml index c1c8c6704c59..380f456f4d2d 100644 --- a/.github/workflows/parquet.yml +++ b/.github/workflows/parquet.yml @@ -19,46 +19,14 @@ # tests for parquet and related crates name: "parquet" -concurrency: - group: ${{ github.repository }}-${{ github.head_ref || github.sha }}-${{ github.workflow }} - cancel-in-progress: true - -# trigger for all PRs that touch certain files and changes to main +# Triggered by ci.yml, which owns event routing and concurrency. on: - push: - branches: - - main - pull_request: - paths: - - Cargo.toml - - Cargo.lock - - rust-toolchain.toml - - .gitmodules - - testing - - parquet-testing - - arrow/** - - arrow-arith/** - - arrow-array/** - - arrow-buffer/** - - arrow-cast/** - - arrow-cmp/** - - arrow-csv/** - - arrow-data/** - - arrow-ipc/** - - arrow-json/** - - arrow-ord/** - - arrow-row/** - - arrow-schema/** - - arrow-select/** - - arrow-string/** - - parquet/** - - parquet_derive/** - - parquet_derive_test/** - - parquet-geospatial/** - - parquet-variant/** - - parquet-variant-compute/** - - parquet-variant-json/** - - .github/** + workflow_call: + inputs: + cache-refresh-only: + description: Populate caches without running tests or linters + type: boolean + default: false permissions: contents: read @@ -66,6 +34,7 @@ permissions: jobs: # test the crate linux-test: + if: ${{ !inputs.cache-refresh-only }} name: Test runs-on: ubuntu-latest container: @@ -100,6 +69,7 @@ jobs: # test compilation linux-features: + if: ${{ !inputs.cache-refresh-only }} name: Check Compilation runs-on: ubuntu-latest container: @@ -163,6 +133,7 @@ jobs: # test the parquet crate builds against wasm32 in stable rust wasm32-build: + if: ${{ !inputs.cache-refresh-only }} name: Build wasm32 runs-on: ubuntu-latest container: @@ -202,24 +173,29 @@ jobs: cd parquet/pytest pip install -r requirements.txt - name: Black check the test files + if: ${{ !inputs.cache-refresh-only }} run: | cd parquet/pytest black --check *.py --verbose - name: Setup Rust toolchain + if: ${{ !inputs.cache-refresh-only }} run: | rustup toolchain install ${{ matrix.rust }} rustup default ${{ matrix.rust }} - name: Install binary for checking + if: ${{ !inputs.cache-refresh-only }} run: | cargo install --path parquet --bin parquet-show-bloom-filter --features=cli cargo install --path parquet --bin parquet-fromcsv --features=arrow,cli cargo install --path parquet --bin parquet-index --features=cli - name: Run pytest + if: ${{ !inputs.cache-refresh-only }} run: | cd parquet/pytest pytest -v clippy: + if: ${{ !inputs.cache-refresh-only }} name: Clippy runs-on: ubuntu-latest container: diff --git a/.github/workflows/rust.yml b/.github/workflows/rust.yml index 0937bcf565d8..9b1b5d3872a6 100644 --- a/.github/workflows/rust.yml +++ b/.github/workflows/rust.yml @@ -18,16 +18,9 @@ # workspace wide tests name: rust -concurrency: - group: ${{ github.repository }}-${{ github.head_ref || github.sha }}-${{ github.workflow }} - cancel-in-progress: true - -# trigger for all PRs and changes to main +# Triggered by ci.yml, which owns event routing and concurrency. on: - push: - branches: - - main - pull_request: + workflow_call: permissions: contents: read @@ -112,7 +105,7 @@ jobs: - name: Setup Rust toolchain uses: ./.github/actions/setup-builder - name: Install cargo-msrv - uses: taiki-e/install-action@9114bf4d891761788c546334fd37538eae1bf8b3 # v2.87.16 + uses: taiki-e/install-action@4cef1412cce204788f482e778a0b9187f9626a29 # v2.87.21 with: tool: cargo-msrv - name: Check all packages diff --git a/Cargo.lock b/Cargo.lock index dfe9ed658eab..3e8e960406b8 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -989,9 +989,9 @@ dependencies = [ [[package]] name = "comfy-table" -version = "8.0.0" +version = "8.0.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "136c8c4c3823846e8ba6d4bda011b4e5d5827b8bb0ac26c31f9ab31abe9f54f2" +checksum = "b96a46ad395a5e00676eddaca4a45ed78b29c08be5a878d35e4b071c625637dd" dependencies = [ "unicode-segmentation", "unicode-width", @@ -1702,16 +1702,17 @@ dependencies = [ [[package]] name = "hyper-util" -version = "0.1.20" +version = "0.1.21" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "96547c2556ec9d12fb1578c4eaf448b04993e7fb79cbaad930a656880a6bdfa0" +checksum = "ddc03d96684f9226b8a787cdb71488417b53ab5ea8fdb1dac946cb9431cc8bff" dependencies = [ - "base64 0.22.1", + "base64 0.23.1", "bytes", "futures-channel", "futures-util", "http", "http-body", + "httparse", "hyper", "ipnet", "libc", diff --git a/Cargo.toml b/Cargo.toml index 2c9855b6f589..1bdb54db6274 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -66,7 +66,6 @@ exclude = [ [workspace.package] version = "60.0.0" -homepage = "https://github.com/apache/arrow-rs" repository = "https://github.com/apache/arrow-rs" authors = ["Apache Arrow "] license = "Apache-2.0" @@ -220,6 +219,7 @@ wildcard_dependencies = "warn" # Lints in `pedantic` that we currently allow. # Some of these we should consider enabling in the future. # Counts come from `cargo clippy --all-features --all-targets`. +assert_is_empty = "allow" # ~37 violations in ~19 files cast_lossless = "allow" # ~640 violations in ~100 files cast_possible_truncation = "allow" # ~890 violations in ~190 files cast_possible_wrap = "allow" # ~570 violations in ~120 files diff --git a/arrow-arith/Cargo.toml b/arrow-arith/Cargo.toml index b5e1fcc25725..8bb02b1d8c28 100644 --- a/arrow-arith/Cargo.toml +++ b/arrow-arith/Cargo.toml @@ -19,7 +19,6 @@ name = "arrow-arith" version = { workspace = true } description = "Arrow arithmetic kernels" -homepage = { workspace = true } repository = { workspace = true } authors = { workspace = true } license = { workspace = true } @@ -50,5 +49,9 @@ criterion = { workspace = true } name = "decimal_arithmetic" harness = false +[[bench]] +name = "temporal" +harness = false + [lints] workspace = true diff --git a/arrow-arith/benches/temporal.rs b/arrow-arith/benches/temporal.rs new file mode 100644 index 000000000000..e2861a70979c --- /dev/null +++ b/arrow-arith/benches/temporal.rs @@ -0,0 +1,91 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT 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::hint::black_box; + +use arrow_arith::temporal::{DatePart, date_part}; +use arrow_array::PrimitiveArray; +use arrow_array::types::*; +use criterion::{Criterion, Throughput, criterion_group, criterion_main}; + +const SIZE: usize = 8192; + +const TIME_PARTS: [DatePart; 6] = [ + DatePart::Hour, + DatePart::Minute, + DatePart::Second, + DatePart::Millisecond, + DatePart::Microsecond, + DatePart::Nanosecond, +]; + +fn benchmark(c: &mut Criterion, unit: &str, units_per_second: i64) { + // A deterministic mix of dates before and after the epoch, including subseconds. + let values: Vec = (0..SIZE) + .map(|i| { + let seconds = (i as i64 * 1_000_003) % 4_000_000_000 - 2_000_000_000; + seconds * units_per_second + (i as i64 * 7919) % units_per_second + }) + .collect(); + let with_nulls = PrimitiveArray::::from_iter( + values + .iter() + .enumerate() + .map(|(i, &v)| (i % 5 != 0).then_some(v)), + ); + let array = PrimitiveArray::::new(values.into(), None); + let mut group = c.benchmark_group(format!("timestamp_{unit}")); + group.throughput(Throughput::Elements(SIZE as u64)); + // Second timestamps have no fractional part, so their subsecond parts are + // constant zero and not worth measuring. + let parts = if units_per_second == 1 { + &TIME_PARTS[..3] + } else { + &TIME_PARTS[..] + }; + for &part in parts { + group.bench_function(format!("{part}/no_nulls"), |b| { + b.iter(|| black_box(date_part(black_box(&array), part).unwrap())) + }); + } + // Validity is not consulted per element, so one nullable case suffices. + group.bench_function("Hour/mixed_nulls", |b| { + b.iter(|| black_box(date_part(black_box(&with_nulls), DatePart::Hour).unwrap())) + }); + // Reference for the calendar conversion path. + group.bench_function("Year/control", |b| { + b.iter(|| black_box(date_part(black_box(&array), DatePart::Year).unwrap())) + }); + // Timestamps with a timezone always take the calendar conversion path. + let with_timezone = array.with_timezone("+05:45"); + for part in [DatePart::Minute, parts[parts.len() - 1]] { + group.bench_function(format!("{part}/timezone_control"), |b| { + b.iter(|| black_box(date_part(black_box(&with_timezone), part).unwrap())) + }); + } + group.finish(); +} + +fn temporal(c: &mut Criterion) { + benchmark::(c, "s", 1); + benchmark::(c, "ms", 1_000); + benchmark::(c, "us", 1_000_000); + benchmark::(c, "ns", 1_000_000_000); +} + +criterion_group!(benches, temporal); +criterion_main!(benches); diff --git a/arrow-arith/src/aggregate.rs b/arrow-arith/src/aggregate.rs index d059f2b28711..99615f760125 100644 --- a/arrow-arith/src/aggregate.rs +++ b/arrow-arith/src/aggregate.rs @@ -700,10 +700,7 @@ mod ree { where F: FnMut(V::Native, V::Native, usize) -> Result, { - let run_ends = array.run_ends(); - let logical_start = run_ends.offset(); - let logical_end = run_ends.offset() + run_ends.len(); - let run_ends = run_ends.sliced_values(); + let run_ends = array.run_ends().sliced_values(); let values_slice = array.run_array().values_slice(); let values = values_slice @@ -717,7 +714,7 @@ mod ree { let mut has_non_null_value = false; for (run_end, value) in run_ends.zip(values) { - let current_run_end = run_end.as_usize().clamp(logical_start, logical_end); + let current_run_end = run_end.as_usize(); let run_length = current_run_end - prev_end; if let Some(value) = value { @@ -726,9 +723,6 @@ mod ree { } prev_end = current_run_end; - if current_run_end == logical_end { - break; - } } Ok(if has_non_null_value { Some(acc) } else { None }) @@ -2083,6 +2077,108 @@ mod tests { assert_eq!(result, Some(100)); } + #[test] + fn test_ree_sum_array_sliced_across_runs() { + let run_ends = Int16Array::from(vec![4, 8]); + let values = Int32Array::from(vec![10, 100]); + let array = RunArray::::try_new(&run_ends, &values).unwrap(); + let sliced = array.slice(3, 4); + let typed_array = sliced.downcast::().unwrap(); + + assert_eq!(sum_array::(typed_array), Some(310)); + } + + #[test] + fn test_ree_sum_array_all_slices() { + fn check() { + let values = [ + Some(2), + Some(2), + None, + None, + Some(-3), + Some(-3), + Some(-3), + Some(4), + ]; + let run_array = make_run_array::(&values); + let plain_array = Int32Array::from(values.to_vec()); + + for offset in 0..=values.len() { + for len in 0..=values.len() - offset { + let sliced = run_array.slice(offset, len); + let typed = sliced.downcast::().unwrap(); + let expected = plain_array.slice(offset, len); + + assert_eq!( + sum_array::(typed), + sum_array::(&expected), + "offset={offset}, len={len}" + ); + assert_eq!( + sum_array_checked::(typed).unwrap(), + sum_array_checked::(&expected).unwrap(), + "offset={offset}, len={len}" + ); + } + } + + let nested = run_array.slice(1, 7).slice(2, 4); + let typed = nested.downcast::().unwrap(); + let expected = plain_array.slice(1, 7).slice(2, 4); + assert_eq!( + sum_array::(typed), + sum_array::(&expected) + ); + assert_eq!( + sum_array_checked::(typed).unwrap(), + sum_array_checked::(&expected).unwrap() + ); + } + + check::(); + check::(); + check::(); + } + + #[test] + fn test_ree_sum_array_checked_sliced_overflow() { + let values = [50, 50, 50, 50, 1, 1, 1, 1]; + let run_array = make_run_array::(&values); + let sliced = run_array.slice(3, 4); + let typed = sliced.downcast::().unwrap(); + assert_eq!(sum_array_checked::(typed).unwrap(), Some(53)); + + let values = [1, 1, 1, 1, 50, 50, 50, 50]; + let run_array = make_run_array::(&values); + let sliced = run_array.slice(3, 4); + let typed = sliced.downcast::().unwrap(); + assert!(sum_array_checked::(typed).is_err()); + } + + #[test] + fn test_ree_sum_array_sliced_infinity() { + let values = [ + 1.0, + 1.0, + 1.0, + 1.0, + 2.0, + 2.0, + 2.0, + 2.0, + f64::INFINITY, + f64::INFINITY, + f64::INFINITY, + f64::INFINITY, + ]; + let run_array = make_run_array::(&values); + let sliced = run_array.slice(5, 4); + let typed = sliced.downcast::().unwrap(); + + assert_eq!(sum_array::(typed), Some(f64::INFINITY)); + } + #[test] fn test_ree_min_max_array_basic() { let run_array = make_run_array::(&[30, 30, 10, 20, 20]); diff --git a/arrow-arith/src/numeric.rs b/arrow-arith/src/numeric.rs index f3a87a74d36d..170d013d6d83 100644 --- a/arrow-arith/src/numeric.rs +++ b/arrow-arith/src/numeric.rs @@ -1042,14 +1042,13 @@ fn decimal_op( // https://cwiki.apache.org/confluence/download/attachments/27362075/Hive_Decimal_Precision_Scale_Support.pdf let array: PrimitiveArray = match op { Op::Add | Op::AddWrapping | Op::Sub | Op::SubWrapping => { + let (p1, s1, p2, s2) = (*p1 as i16, *s1 as i16, *p2 as i16, *s2 as i16); // max(s1, s2) - let result_scale = *s1.max(s2); + let result_scale = s1.max(s2); // max(s1, s2) + max(p1-s1, p2-s2) + 1 let result_precision = - (result_scale.saturating_add((*p1 as i8 - s1).max(*p2 as i8 - s2)) as u8) - .saturating_add(1) - .min(T::MAX_PRECISION); + (result_scale + (p1 - s1).max(p2 - s2) + 1).min(T::MAX_PRECISION as i16) as u8; let l_mul = T::Native::usize_as(10).pow_checked((result_scale - s1) as _)?; let r_mul = T::Native::usize_as(10).pow_checked((result_scale - s2) as _)?; @@ -1082,12 +1081,12 @@ fn decimal_op( } _ => unreachable!(), } - .with_precision_and_scale(result_precision, result_scale)? + .with_precision_and_scale(result_precision, result_scale as i8)? } Op::Mul | Op::MulWrapping => { let result_precision = p1.saturating_add(p2 + 1).min(T::MAX_PRECISION); - let result_scale = s1.saturating_add(*s2); - if result_scale > T::MAX_SCALE { + let result_scale = *s1 as i16 + *s2 as i16; + if result_scale > T::MAX_SCALE as i16 { // SQL standard says that if the resulting scale of a multiply operation goes // beyond the maximum, rounding is not acceptable and thus an error occurs return Err(ArrowError::InvalidArgumentError(format!( @@ -1097,6 +1096,14 @@ fn decimal_op( T::MAX_SCALE ))); } + let result_scale = i8::try_from(result_scale).map_err(|_| { + ArrowError::InvalidArgumentError(format!( + "Output scale of {} {op} {} would be less than min scale of {}", + l.data_type(), + r.data_type(), + i8::MIN + )) + })?; try_op!(l, l_s, r, r_s, l.mul_checked(r)) .with_precision_and_scale(result_precision, result_scale)? @@ -1137,15 +1144,15 @@ fn decimal_op( } Op::Rem => { + let (p1, s1, p2, s2) = (*p1 as i16, *s1 as i16, *p2 as i16, *s2 as i16); // max(s1, s2) - let result_scale = *s1.max(s2); + let result_scale = s1.max(s2); // min(p1-s1, p2 -s2) + max( s1,s2 ) let result_precision = - (result_scale.saturating_add((*p1 as i8 - s1).min(*p2 as i8 - s2)) as u8) - .min(T::MAX_PRECISION); + (result_scale + (p1 - s1).min(p2 - s2)).min(T::MAX_PRECISION as i16) as u8; - let l_mul = T::Native::usize_as(10).pow_wrapping((result_scale - s1) as _); - let r_mul = T::Native::usize_as(10).pow_wrapping((result_scale - s2) as _); + let l_mul = T::Native::usize_as(10).pow_checked((result_scale - s1) as _)?; + let r_mul = T::Native::usize_as(10).pow_checked((result_scale - s2) as _)?; try_op!( l, @@ -1154,7 +1161,7 @@ fn decimal_op( r_s, l.mul_checked(l_mul)?.mod_checked(r.mul_checked(r_mul)?) ) - .with_precision_and_scale(result_precision, result_scale)? + .with_precision_and_scale(result_precision, result_scale as i8)? } }; @@ -1534,6 +1541,115 @@ mod tests { assert_eq!(err, "Divide by zero error"); } + #[test] + fn test_decimal256_add_sub_negative_scale_metadata() { + let a = Decimal256Array::from(vec![Some(i256::ONE), Some(i256::MINUS_ONE), None]) + .with_precision_and_scale(76, -52) + .unwrap(); + let b = Decimal256Array::from(vec![i256::from_i128(2), i256::ONE, i256::ONE]) + .with_precision_and_scale(76, -52) + .unwrap(); + let expected = + Decimal256Array::from(vec![Some(i256::from_i128(3)), Some(i256::ZERO), None]) + .with_precision_and_scale(76, -52) + .unwrap(); + assert_eq!(add(&a, &b).unwrap().as_ref(), &expected); + + let expected = + Decimal256Array::from(vec![Some(i256::MINUS_ONE), Some(i256::from_i128(-2)), None]) + .with_precision_and_scale(76, -52) + .unwrap(); + assert_eq!(sub(&a, &b).unwrap().as_ref(), &expected); + } + + #[test] + fn test_decimal256_add_sub_minimum_scale() { + let a = Decimal256Array::from(vec![i256::ONE]) + .with_precision_and_scale(20, i8::MIN) + .unwrap(); + let b = Decimal256Array::from(vec![i256::from_i128(2)]) + .with_precision_and_scale(20, i8::MIN) + .unwrap(); + let expected = Decimal256Array::from(vec![i256::from_i128(3)]) + .with_precision_and_scale(21, i8::MIN) + .unwrap(); + assert_eq!(add(&a, &b).unwrap().as_ref(), &expected); + + let expected = Decimal256Array::from(vec![i256::MINUS_ONE]) + .with_precision_and_scale(21, i8::MIN) + .unwrap(); + assert_eq!(sub(&a, &b).unwrap().as_ref(), &expected); + } + + #[test] + fn test_decimal256_remainder_negative_scale_metadata() { + let a = Decimal256Array::from(vec![Some(i256::ONE), Some(i256::MINUS_ONE), None]) + .with_precision_and_scale(76, -52) + .unwrap(); + let b = Decimal256Array::from(vec![i256::from_i128(2), i256::ONE, i256::ONE]) + .with_precision_and_scale(76, -52) + .unwrap(); + let expected = Decimal256Array::from(vec![Some(i256::ONE), Some(i256::ZERO), None]) + .with_precision_and_scale(76, -52) + .unwrap(); + assert_eq!(rem(&a, &b).unwrap().as_ref(), &expected); + } + + #[test] + fn test_decimal256_adjacent_minimum_scales() { + let a = Decimal256Array::from(vec![i256::ONE]) + .with_precision_and_scale(20, i8::MIN) + .unwrap(); + let b = Decimal256Array::from(vec![i256::from_i128(2)]) + .with_precision_and_scale(20, i8::MIN + 1) + .unwrap(); + let expected = Decimal256Array::from(vec![i256::from_i128(12)]) + .with_precision_and_scale(22, -127) + .unwrap(); + assert_eq!(add(&a, &b).unwrap().as_ref(), &expected); + + let expected = Decimal256Array::from(vec![i256::from_i128(8)]) + .with_precision_and_scale(22, -127) + .unwrap(); + assert_eq!(sub(&a, &b).unwrap().as_ref(), &expected); + + let expected = Decimal256Array::from(vec![i256::ZERO]) + .with_precision_and_scale(20, -127) + .unwrap(); + assert_eq!(rem(&a, &b).unwrap().as_ref(), &expected); + } + + #[test] + fn test_decimal256_extreme_scale_difference_overflow() { + let a = Decimal256Array::from(vec![i256::ONE]) + .with_precision_and_scale(20, i8::MIN) + .unwrap(); + let b = Decimal256Array::from(vec![i256::from_i128(2)]) + .with_precision_and_scale(76, 76) + .unwrap(); + let expected = "Arithmetic overflow: Overflow happened on: 10 ^ 204"; + assert_eq!(add(&a, &b).unwrap_err().to_string(), expected); + assert_eq!(sub(&a, &b).unwrap_err().to_string(), expected); + assert_eq!(rem(&a, &b).unwrap_err().to_string(), expected); + } + + #[test] + fn test_decimal256_multiply_minimum_scale() { + let a = Decimal256Array::from(vec![i256::ONE]) + .with_precision_and_scale(76, -64) + .unwrap(); + let expected = Decimal256Array::from(vec![i256::ONE]) + .with_precision_and_scale(76, i8::MIN) + .unwrap(); + assert_eq!(mul(&a, &a).unwrap().as_ref(), &expected); + + let b = a.clone().with_precision_and_scale(76, -65).unwrap(); + assert_eq!( + mul(&a, &b).unwrap_err().to_string(), + "Invalid argument error: Output scale of Decimal256(76, -64) * Decimal256(76, -65) would be less than min scale of -128" + ); + } + #[test] fn test_decimal256_div_wide_intermediate() { // Dividing two scale-37 values needs l * 10^41, which is 79 digits and does not diff --git a/arrow-arith/src/temporal.rs b/arrow-arith/src/temporal.rs index b4653c2fe78e..f42d99d4c3a4 100644 --- a/arrow-arith/src/temporal.rs +++ b/arrow-arith/src/temporal.rs @@ -184,6 +184,10 @@ where /// Returns an [`Int32Array`] unless input was a dictionary type, in which case returns /// the dictionary but with this function applied onto its values. /// +/// Null inputs produce null outputs. A timestamp outside the supported calendar +/// range also produces null, except that for timestamps without a timezone the +/// time parts (`Hour` through `Nanosecond`) are defined for every value. +/// /// If array passed in is not of the above listed types (or is a dictionary array where the /// values array isn't of the above listed types), then this function will return an error. /// @@ -415,6 +419,62 @@ impl ExtractDatePartExt for PrimitiveArray { } } +/// Number of ticks in one second for timestamps of type `T`. +const fn units_per_second() -> i64 { + match T::UNIT { + TimeUnit::Second => 1, + TimeUnit::Millisecond => MILLISECONDS, + TimeUnit::Microsecond => MICROSECONDS, + TimeUnit::Nanosecond => NANOSECONDS, + } +} + +/// Extracts a time-of-day part from a timestamp array without a timezone using +/// integer arithmetic. Returns `None` for parts that require a calendar +/// conversion. +/// +/// `div_euclid` rounds toward negative infinity and `rem_euclid` yields a +/// nonnegative remainder, so pre-epoch values produce in-range components: +/// -1 ms is 1969-12-31 23:59:59.999, whose second is 59 and millisecond is 999. +fn timestamp_time_part( + array: &PrimitiveArray, + part: DatePart, +) -> Option { + Some(match part { + DatePart::Hour => { + array.unary(|v| v.div_euclid(units_per_second::() * 3_600).rem_euclid(24) as i32) + } + DatePart::Minute => { + array.unary(|v| v.div_euclid(units_per_second::() * 60).rem_euclid(60) as i32) + } + DatePart::Second => { + array.unary(|v| v.div_euclid(units_per_second::()).rem_euclid(60) as i32) + } + DatePart::Millisecond => { + array.unary(|v| subsecond(v, units_per_second::(), MILLISECONDS)) + } + DatePart::Microsecond => { + array.unary(|v| subsecond(v, units_per_second::(), MICROSECONDS)) + } + DatePart::Nanosecond => array.unary(|v| subsecond(v, units_per_second::(), NANOSECONDS)), + _ => return None, + }) +} + +/// Fractional second of `v`, counted in `units` per second, rescaled to `target` +/// per second. Both scales divide one second evenly, so the conversion is exact +/// when scaling up and truncating when scaling down. +#[inline] +fn subsecond(v: i64, units: i64, target: i64) -> i32 { + let fraction = v.rem_euclid(units); + let scaled = if target >= units { + fraction * (target / units) + } else { + fraction / (units / target) + }; + scaled as i32 +} + impl ExtractDatePartExt for PrimitiveArray { fn date_part(&self, part: DatePart) -> Result { // TimestampSecond only encodes number of seconds, so these will always be 0 @@ -428,6 +488,8 @@ impl ExtractDatePartExt for PrimitiveArray { .map(|c| Utc.from_utc_datetime(&c).with_timezone(&tz)) .map(map_func) }) + } else if let Some(array) = timestamp_time_part(self, part) { + array } else { let map_func = get_date_time_part_extract_fn(part); self.unary_opt(|d| timestamp_s_to_datetime(d).map(map_func)) @@ -445,6 +507,8 @@ impl ExtractDatePartExt for PrimitiveArray { .map(|c| Utc.from_utc_datetime(&c).with_timezone(&tz)) .map(map_func) }) + } else if let Some(array) = timestamp_time_part(self, part) { + array } else { let map_func = get_date_time_part_extract_fn(part); self.unary_opt(|d| timestamp_ms_to_datetime(d).map(map_func)) @@ -462,6 +526,8 @@ impl ExtractDatePartExt for PrimitiveArray { .map(|c| Utc.from_utc_datetime(&c).with_timezone(&tz)) .map(map_func) }) + } else if let Some(array) = timestamp_time_part(self, part) { + array } else { let map_func = get_date_time_part_extract_fn(part); self.unary_opt(|d| timestamp_us_to_datetime(d).map(map_func)) @@ -479,6 +545,8 @@ impl ExtractDatePartExt for PrimitiveArray { .map(|c| Utc.from_utc_datetime(&c).with_timezone(&tz)) .map(map_func) }) + } else if let Some(array) = timestamp_time_part(self, part) { + array } else { let map_func = get_date_time_part_extract_fn(part); self.unary_opt(|d| timestamp_ns_to_datetime(d).map(map_func)) @@ -752,6 +820,9 @@ impl ChronoDateExt for T { mod tests { use super::*; + use arrow_array::temporal_conversions::as_datetime; + use chrono::NaiveDate; + /// Used to integrate new [`date_part()`] method with deprecated shims such as /// [`hour()`] and [`week()`]. fn date_part_primitive( @@ -762,6 +833,119 @@ mod tests { Ok(array.as_primitive::().to_owned()) } + /// Time-of-day parts of a timestamp. + const TIME_PARTS: [DatePart; 6] = [ + DatePart::Hour, + DatePart::Minute, + DatePart::Second, + DatePart::Millisecond, + DatePart::Microsecond, + DatePart::Nanosecond, + ]; + + /// Checks time part extraction against the calendar conversion for values within + /// its supported range. Covers component boundaries, negative timestamps, nulls + /// and slices, plus timezone dispatch and invalid timezone metadata. + fn check_timestamp_time_parts(units_per_second: i64) { + let mut values = vec![i64::MIN, i64::MAX]; + for boundary in [1_000, 1_000_000] { + for delta in [-1, 0, 1] { + values.extend([boundary + delta, -boundary + delta]); + } + } + for base in [0, 60, 3_600, 86_400] { + for delta in [-1, 0, 1] { + for seconds in [base + delta, -base + delta] { + for fraction in [-1, 0, 1, units_per_second - 1] { + let value = i128::from(seconds) * i128::from(units_per_second) + + i128::from(fraction); + if let Ok(value) = i64::try_from(value) { + values.push(value); + } + } + } + } + } + // Sample both the full i64 domain and dates near the epoch: most full-range + // second and millisecond timestamps are outside the supported calendar range. + let mut seed = 42_u64; + for _ in 0..256 { + seed = seed.wrapping_mul(6364136223846793005).wrapping_add(1); + values.push(seed as i64); + let seconds = (seed % 4_000_000_000) as i64 - 2_000_000_000; + values.push(seconds * units_per_second + (seed % units_per_second as u64) as i64); + } + // The calendar conversion provides expected values only within its range. + // Behavior at the i64 extremes is checked separately. + values.retain(|&v| as_datetime::(v).is_some()); + let no_nulls = PrimitiveArray::::new(values.clone().into(), None); + let mixed = PrimitiveArray::::from_iter( + values + .iter() + .enumerate() + .map(|(i, &v)| (i % 5 != 0).then_some(v)), + ); + for array in [no_nulls, mixed.slice(1, mixed.len() - 2)] { + for part in TIME_PARTS { + let map_func = get_date_time_part_extract_fn(part); + let expected: Int32Array = array.unary_opt(|v| as_datetime::(v).map(map_func)); + let actual = date_part(&array, part).unwrap(); + assert_eq!( + actual.as_primitive::(), + &expected, + "{} {part}", + T::DATA_TYPE + ); + } + } + + // An offset that changes both hour and minute checks that timezone metadata + // still routes through the calendar conversion. + let zoned = PrimitiveArray::::from_iter_values([0]).with_timezone("+05:45"); + for (part, expected) in [(DatePart::Hour, 5), (DatePart::Minute, 45)] { + assert_eq!( + date_part_primitive(&zoned, part).unwrap(), + Int32Array::from(vec![expected]) + ); + } + let invalid = PrimitiveArray::::from_iter_values([0]).with_timezone("invalid timezone"); + assert!(date_part(&invalid, DatePart::Second).is_err()); + } + + #[test] + fn test_timestamp_time_parts_match_chrono() { + check_timestamp_time_parts::(1); + check_timestamp_time_parts::(MILLISECONDS); + check_timestamp_time_parts::(MICROSECONDS); + check_timestamp_time_parts::(NANOSECONDS); + } + + #[test] + fn test_timestamp_time_parts_extremes() { + /// `(hour, minute, second, nanosecond)` expected at `i64::MIN` and `i64::MAX`. + fn check(min: (u32, u32, u32, u32), max: (u32, u32, u32, u32)) { + let time = |(h, m, s, ns)| NaiveDate::default().and_hms_nano_opt(h, m, s, ns).unwrap(); + let expected = [Some(time(min)), None, Some(time(max))]; + let array = PrimitiveArray::::from_iter([Some(i64::MIN), None, Some(i64::MAX)]); + for part in TIME_PARTS { + let map_func = get_date_time_part_extract_fn(part); + assert_eq!( + date_part_primitive(&array, part).unwrap(), + Int32Array::from_iter(expected.iter().map(|t| t.map(map_func))), + "{} {part}", + T::DATA_TYPE + ); + } + } + + // These values lie outside the supported calendar range for these units, so + // the expectations cannot be derived from the calendar conversion. Nanosecond + // extremes are within range and are covered by the comparison against it. + check::((8, 29, 52, 0), (15, 30, 7, 0)); + check::((16, 47, 4, 192_000_000), (7, 12, 55, 807_000_000)); + check::((19, 59, 5, 224_192_000), (4, 0, 54, 775_807_000)); + } + #[test] fn test_temporal_array_date64_hour() { let a: PrimitiveArray = diff --git a/arrow-array/Cargo.toml b/arrow-array/Cargo.toml index b7c7cfabc6ba..c9ade8ef240d 100644 --- a/arrow-array/Cargo.toml +++ b/arrow-array/Cargo.toml @@ -19,7 +19,6 @@ name = "arrow-array" version = { workspace = true } description = "Array abstractions for Apache Arrow" -homepage = { workspace = true } repository = { workspace = true } authors = { workspace = true } license = "Apache-2.0 AND MIT" diff --git a/arrow-array/src/array/byte_array.rs b/arrow-array/src/array/byte_array.rs index b3f2025322fb..049d1f47cfcb 100644 --- a/arrow-array/src/array/byte_array.rs +++ b/arrow-array/src/array/byte_array.rs @@ -282,13 +282,18 @@ impl GenericByteArray { /// Returns the values of this array /// /// Unlike [`Self::value_data`] this returns the [`Buffer`] - /// allowing for zero-copy cloning + /// allowing for zero-copy cloning. Like [`Self::value_data`], it can + /// include bytes that are not part of any value. #[inline] pub fn values(&self) -> &Buffer { &self.value_data } /// Returns the raw value data + /// + /// This can include bytes before the first offset or after the last one, + /// for example when the array is a slice of a larger array. Those bytes are + /// not part of any value, and in a string array they need not be valid UTF-8. pub fn value_data(&self) -> &[u8] { self.value_data.as_slice() } @@ -378,87 +383,56 @@ impl GenericByteArray { } } - /// Returns `GenericByteBuilder` of this byte array for mutating its values if the underlying - /// offset and data buffers are not shared by others. + /// Converts this array into a [`GenericByteBuilder`] that reuses its buffers, + /// if possible. + /// + /// Returns `Err(self)` if the offsets or values cannot be reused, for + /// example because they are shared with another array, or because this is a + /// sliced array with a non-zero first offset. + /// + /// The null buffer is copied if it cannot be reused. pub fn into_builder(self) -> Result, Self> { - let len = self.len(); - let value_len = T::Offset::as_usize(self.value_offsets()[len] - self.value_offsets()[0]); - - let data = self.into_data(); - let null_bit_buffer = data.nulls().map(|b| b.inner().sliced()); - - let element_len = std::mem::size_of::(); - let offset_buffer = data.buffers()[0] - .slice_with_length(data.offset() * element_len, (len + 1) * element_len); - - let element_len = std::mem::size_of::(); - let value_buffer = data.buffers()[1] - .slice_with_length(data.offset() * element_len, value_len * element_len); - - drop(data); + let (offsets, values, nulls) = self.into_parts(); - let try_mutable_null_buffer = match null_bit_buffer { - None => Ok(None), - Some(null_buffer) => { - // Null buffer exists, tries to make it mutable - null_buffer.into_mutable().map(Some) + // Puts the parts back together into the original array + let rebuild = |offsets: Buffer, values: Buffer, nulls: Option| { + // SAFETY: the parts come unchanged from `self` + unsafe { + Self::new_unchecked(OffsetBuffer::new_unchecked(offsets.into()), values, nulls) } }; - let try_mutable_buffers = match try_mutable_null_buffer { - Ok(mutable_null_buffer) => { - // Got mutable null buffer, tries to get mutable value buffer - let try_mutable_offset_buffer = offset_buffer.into_mutable(); - let try_mutable_value_buffer = value_buffer.into_mutable(); - - // try_mutable_offset_buffer.map(...).map_err(...) doesn't work as the compiler complains - // mutable_null_buffer is moved into map closure. - match (try_mutable_offset_buffer, try_mutable_value_buffer) { - (Ok(mutable_offset_buffer), Ok(mutable_value_buffer)) => unsafe { - Ok(GenericByteBuilder::::new_from_buffer( - mutable_offset_buffer, - mutable_value_buffer, - mutable_null_buffer, - )) - }, - (Ok(mutable_offset_buffer), Err(value_buffer)) => Err(( - mutable_offset_buffer.into(), - value_buffer, - mutable_null_buffer.map(|b| b.into()), - )), - (Err(offset_buffer), Ok(mutable_value_buffer)) => Err(( - offset_buffer, - mutable_value_buffer.into(), - mutable_null_buffer.map(|b| b.into()), - )), - (Err(offset_buffer), Err(value_buffer)) => Err(( - offset_buffer, - value_buffer, - mutable_null_buffer.map(|b| b.into()), - )), - } - } - Err(mutable_null_buffer) => { - // Unable to get mutable null buffer - Err((offset_buffer, value_buffer, Some(mutable_null_buffer))) - } + // A buffer can be reused only if this array is its only owner and it + // starts at the beginning of its allocation. Otherwise, the original + // array is returned. + let offsets = match offsets.into_inner().into_inner().into_mutable() { + Ok(offsets) => offsets, + Err(offsets) => return Err(rebuild(offsets, values, nulls)), + }; + let values = match values.into_mutable() { + Ok(values) => values, + Err(values) => return Err(rebuild(offsets.into(), values, nulls)), }; - match try_mutable_buffers { - Ok(builder) => Ok(builder), - Err((offset_buffer, value_buffer, null_bit_buffer)) => { - let builder = ArrayData::builder(T::DATA_TYPE) - .len(len) - .add_buffer(offset_buffer) - .add_buffer(value_buffer) - .null_bit_buffer(null_bit_buffer); - - let array_data = unsafe { builder.build_unchecked() }; - let array = GenericByteArray::::from(array_data); - - Err(array) + // The builder needs a null buffer whose first bit is for the first row. + // That is the case when the bit offset is 0, and then the buffer is reused + // if this array is its only owner. Otherwise, the bits are copied, which is + // cheap as there is one bit per row. + let nulls = nulls.map(|nulls| { + let bits = nulls.into_inner(); + if bits.offset() == 0 { + let len = bits.len(); + bits.into_inner() + .into_mutable() + .unwrap_or_else(|buffer| buffer[..len.div_ceil(8)].to_vec().into()) + } else { + // `sliced` returns the bits starting at the first row + bits.sliced().to_vec().into() } - } + }); + + // SAFETY: the buffers come from a valid array + Ok(unsafe { GenericByteBuilder::new_from_buffer(offsets, values, nulls) }) } } diff --git a/arrow-array/src/array/dictionary_array.rs b/arrow-array/src/array/dictionary_array.rs index 6bd3e64bef21..c561d3aa7894 100644 --- a/arrow-array/src/array/dictionary_array.rs +++ b/arrow-array/src/array/dictionary_array.rs @@ -1068,7 +1068,7 @@ impl AnyDictionaryArray for DictionaryArray { mod tests { use super::*; use crate::cast::as_dictionary_array; - use crate::{Int8Array, Int16Array, Int32Array, RunArray, UInt8Array}; + use crate::{Decimal128Array, Int8Array, Int16Array, Int32Array, RunArray, UInt8Array}; use arrow_buffer::{Buffer, ToByteSlice}; #[test] @@ -1455,6 +1455,26 @@ mod tests { assert_eq!(expected, returned); } + #[test] + fn test_unary_mut_keeps_value_type() { + let values = || { + Decimal128Array::from(vec![123, 456]) + .with_precision_and_scale(5, 2) + .unwrap() + }; + let keys = Int8Array::from_iter_values([0, 1, 0]); + let dict = DictionaryArray::new(keys, Arc::new(values())); + let data_type = dict.data_type().clone(); + let updated = dict.unary_mut::<_, Decimal128Type>(|x| x + 1).unwrap(); + assert_eq!(updated.data_type(), &data_type); + + // When the keys are shared, the original array is returned + let keys = Int8Array::from_iter_values([0, 1, 0]); + let dict = DictionaryArray::new(keys.clone(), Arc::new(values())); + let returned = dict.unary_mut::<_, Decimal128Type>(|x| x + 1).unwrap_err(); + assert_eq!(returned, DictionaryArray::new(keys, Arc::new(values()))); + } + #[test] fn test_occupancy() { let keys = Int32Array::new((100..200).collect(), None); diff --git a/arrow-array/src/array/primitive_array.rs b/arrow-array/src/array/primitive_array.rs index 30d8f95f805d..b88089f52783 100644 --- a/arrow-array/src/array/primitive_array.rs +++ b/arrow-array/src/array/primitive_array.rs @@ -1150,6 +1150,7 @@ impl PrimitiveArray { /// `Err(self)` pub fn into_builder(self) -> Result, Self> { let len = self.len(); + let data_type = self.data_type.clone(); let data = self.into_data(); let null_bit_buffer = data.nulls().map(|b| b.inner().sliced()); @@ -1189,9 +1190,9 @@ impl PrimitiveArray { }; match try_mutable_buffers { - Ok(builder) => Ok(builder), + Ok(builder) => Ok(builder.with_data_type(data_type)), Err((buffer, null_bit_buffer)) => { - let builder = ArrayData::builder(T::DATA_TYPE) + let builder = ArrayData::builder(data_type) .len(len) .add_buffer(buffer) .null_bit_buffer(null_bit_buffer); @@ -2935,6 +2936,25 @@ mod tests { } } + #[test] + fn test_into_builder_keeps_data_type() { + let array = Decimal128Array::from(vec![1, 2]) + .with_precision_and_scale(10, 2) + .unwrap(); + let data_type = array.data_type().clone(); + + // Shared buffers: the array is returned + let shared = array.clone(); + let array = array.into_builder().unwrap_err(); + assert_eq!(array.data_type(), &data_type); + drop(shared); + + // Unshared buffers: the buffers are reused by the builder + let mut builder = array.into_builder().unwrap(); + builder.append_value(3); + assert_eq!(builder.finish().data_type(), &data_type); + } + #[test] fn test_into_builder_on_sliced_array() { let array: Int32Array = vec![1, 2, 3].into_iter().map(Some).collect(); diff --git a/arrow-array/src/array/string_array.rs b/arrow-array/src/array/string_array.rs index b9f06c2a80a0..39427181d288 100644 --- a/arrow-array/src/array/string_array.rs +++ b/arrow-array/src/array/string_array.rs @@ -159,7 +159,7 @@ mod tests { use crate::Array; use crate::builder::{ListBuilder, PrimitiveBuilder, StringBuilder}; use crate::types::UInt8Type; - use arrow_buffer::Buffer; + use arrow_buffer::{Buffer, NullBuffer, OffsetBuffer}; use arrow_data::ArrayData; use arrow_schema::{DataType, Field}; use std::sync::Arc; @@ -504,6 +504,29 @@ mod tests { let _ = StringArray::from(list); } + #[test] + fn test_string_array_invalid_bytes_outside_offsets() { + // Only the bytes that the offsets span need to be valid UTF-8, as in a + // slice of a larger array + let values = Buffer::from_slice_ref(b"\xFFa\xC3\xA9\xFF"); + let offsets = OffsetBuffer::new(vec![1, 2, 4].into()); + let string = StringArray::try_new(offsets, values.clone(), None).unwrap(); + assert_eq!(string, StringArray::from(vec!["a", "é"])); + + let offsets = OffsetBuffer::new(vec![1, 3, 4].into()); + let err = StringArray::try_new(offsets, values, None).unwrap_err(); + assert_eq!( + err.to_string(), + "Invalid argument error: Split UTF-8 codepoint at offset 3" + ); + + // An empty value whose offset is inside a character, in a values buffer + // that is valid UTF-8 as a whole. ArrayData validation must agree. + let offsets = OffsetBuffer::new(vec![1, 1].into()); + let string = StringArray::try_new(offsets, Buffer::from_slice_ref("é"), None).unwrap(); + string.to_data().validate_full().unwrap(); + } + #[test] fn test_empty_offsets() { let string = StringArray::from( @@ -550,6 +573,66 @@ mod tests { assert_eq!(&err_return, &shared_array); } + #[test] + fn test_into_builder_non_zero_first_offset() { + // Sliced, and the only owner of its buffers + let array: StringArray = vec!["abcde", "fgh", "ij"].into(); + let sliced = array.slice(1, 2); + drop(array); + let err_return = sliced.into_builder().unwrap_err(); + assert_eq!(err_return, StringArray::from(vec!["fgh", "ij"])); + + // Not sliced, but there are unused bytes before the first offset and + // after the last offset + let array = StringArray::new( + OffsetBuffer::new(vec![3, 5, 8].into()), + Buffer::from(b"xxxabcdeyy"), + None, + ); + let mut builder = array.into_builder().unwrap(); + builder.append_value("fg"); + assert_eq!(builder.finish(), StringArray::from(vec!["ab", "cde", "fg"])); + } + + #[test] + fn test_into_builder_err_returns_original_buffers() { + // The offsets could be reused, but the values buffer is shared + let values = Buffer::from(b"abxyz"); + let array = StringArray::new( + OffsetBuffer::new(vec![0, 2, 2].into()), + values.clone(), + None, + ); + let offsets_ptr = array.offsets().as_ptr(); + + let err_return = array.into_builder().unwrap_err(); + assert_eq!(err_return.offsets().as_ptr(), offsets_ptr); + assert_eq!(err_return.values(), &values); + } + + #[test] + fn test_into_builder_nulls() { + // Reused when this array is the only owner + let array = StringArray::from(vec![Some("ab"), None]); + let nulls_ptr = array.nulls().unwrap().buffer().as_ptr(); + let builder = array.into_builder().unwrap(); + assert_eq!(builder.validity_slice().unwrap().as_ptr(), nulls_ptr); + + // Copied when shared + let nulls = NullBuffer::from(vec![true, false]); + let array = StringArray::new( + OffsetBuffer::new(vec![0, 2, 2].into()), + Buffer::from(b"ab"), + Some(nulls.clone()), + ); + let mut builder = array.into_builder().unwrap(); + builder.append_value("c"); + assert_eq!( + builder.finish(), + StringArray::from(vec![Some("ab"), None, Some("c")]) + ); + } + #[test] fn test_non_null_string_array_equal() { let a = StringArray::from(vec![Some("ab"), Some("c")]); diff --git a/arrow-array/src/builder/generic_bytes_builder.rs b/arrow-array/src/builder/generic_bytes_builder.rs index af8950fb70b0..e823946fc9cc 100644 --- a/arrow-array/src/builder/generic_bytes_builder.rs +++ b/arrow-array/src/builder/generic_bytes_builder.rs @@ -58,17 +58,24 @@ impl GenericByteBuilder { /// Creates a new [`GenericByteBuilder`] from buffers. /// + /// Any bytes in `value_buffer` after the last offset are dropped. The first + /// bit of `null_buffer` is for the first row. + /// /// # Safety /// /// This doesn't verify buffer contents as it assumes the buffers are from /// existing and valid [`GenericByteArray`]. pub unsafe fn new_from_buffer( offsets_buffer: MutableBuffer, - value_buffer: MutableBuffer, + mut value_buffer: MutableBuffer, null_buffer: Option, ) -> Self { let offsets_builder: Vec = ScalarBuffer::::from(offsets_buffer).into(); + + // New values are appended at the end of the values buffer, so it must end + // at the last offset + value_buffer.truncate(offsets_builder.last().unwrap().as_usize()); let value_builder: Vec = ScalarBuffer::::from(value_buffer).into(); let null_buffer_builder = null_buffer diff --git a/arrow-array/src/builder/primitive_dictionary_builder.rs b/arrow-array/src/builder/primitive_dictionary_builder.rs index bc5d87850057..f2ab6a49d32a 100644 --- a/arrow-array/src/builder/primitive_dictionary_builder.rs +++ b/arrow-array/src/builder/primitive_dictionary_builder.rs @@ -459,7 +459,8 @@ where let values = self.values_builder.finish_cloned(); let keys = self.keys_builder.finish_cloned(); - let data_type = DataType::Dictionary(Box::new(K::DATA_TYPE), Box::new(V::DATA_TYPE)); + let data_type = + DataType::Dictionary(Box::new(K::DATA_TYPE), Box::new(values.data_type().clone())); let builder = keys .into_data() @@ -492,7 +493,8 @@ where let values = self.values_builder.finish_cloned(); let keys = self.keys_builder.finish(); - let data_type = DataType::Dictionary(Box::new(K::DATA_TYPE), Box::new(V::DATA_TYPE)); + let data_type = + DataType::Dictionary(Box::new(K::DATA_TYPE), Box::new(values.data_type().clone())); let builder = keys .into_data() @@ -604,15 +606,16 @@ mod tests { keys_builder, values_builder, ); + let expected = DataType::Dictionary( + Box::new(DataType::Int32), + Box::new(DataType::Decimal128(1, 2)), + ); + assert_eq!(builder.finish_cloned().data_type(), &expected); + assert_eq!(builder.finish_preserve_values().data_type(), &expected); + let dict_array = builder.finish(); assert_eq!(dict_array.value_type(), DataType::Decimal128(1, 2)); - assert_eq!( - dict_array.data_type(), - &DataType::Dictionary( - Box::new(DataType::Int32), - Box::new(DataType::Decimal128(1, 2)), - ) - ); + assert_eq!(dict_array.data_type(), &expected); } #[test] diff --git a/arrow-array/src/builder/primitive_run_builder.rs b/arrow-array/src/builder/primitive_run_builder.rs index 71fcf874759f..ec2f7febc50f 100644 --- a/arrow-array/src/builder/primitive_run_builder.rs +++ b/arrow-array/src/builder/primitive_run_builder.rs @@ -200,9 +200,10 @@ where // write the last run end to the array. self.append_run_end(); - // reset the run index to zero. + // reset the run indexes to zero. self.current_value = None; self.current_run_end_index = 0; + self.prev_run_end_index = 0; // build the run encoded array by adding run_ends and values array as its children. let run_ends_array = self.run_ends_builder.finish(); @@ -333,6 +334,23 @@ mod tests { ); } + #[test] + fn test_reuse_after_finish() { + let mut builder = PrimitiveRunBuilder::::new(); + builder.extend([1, 2, 3].into_iter().map(Some)); + builder.finish(); + + // A single run with the same length as the previous batch + builder.extend([7, 7, 7].into_iter().map(Some)); + let array = builder.finish(); + assert_eq!(array.len(), 3); + assert_eq!(array.run_ends().values(), &[3]); + assert_eq!(array.values().as_primitive::().values(), &[7]); + + // Nothing appended since the last finish + assert_eq!(builder.finish().len(), 0); + } + #[test] #[should_panic(expected = "incompatible data type for builder")] fn test_override_data_type_invalid() { @@ -366,4 +384,15 @@ mod tests { let values = array.values(); assert_eq!(values.timezone(), Some("Europe/Paris")); } + + #[test] + fn test_finish_cloned_keeps_data_type() { + let data_type = DataType::Decimal128(5, 2); + let mut builder = PrimitiveRunBuilder::::new() + .with_data_type(data_type.clone()); + // Leave the run open, so `finish_cloned` has to add it to a copy of the values + builder.append_value(123); + let array = builder.finish_cloned(); + assert_eq!(array.values().data_type(), &data_type); + } } diff --git a/arrow-array/src/types.rs b/arrow-array/src/types.rs index ba0d88ae644c..c7f2d5198959 100644 --- a/arrow-array/src/types.rs +++ b/arrow-array/src/types.rs @@ -1686,23 +1686,26 @@ impl ByteArrayType for GenericStringType { }; fn validate(offsets: &OffsetBuffer, values: &Buffer) -> Result<(), ArrowError> { - // Verify that the slice as a whole is valid UTF-8 - let validated = std::str::from_utf8(values).map_err(|e| { + let start = offsets.first().as_usize(); + let end = offsets.last().as_usize(); + if end > values.len() { + return Err(ArrowError::InvalidArgumentError(format!( + "Offset of {end} exceeds length of values {}", + values.len() + ))); + } + + // Verify that the bytes the offsets span are valid UTF-8 + let validated = std::str::from_utf8(&values[start..end]).map_err(|e| { ArrowError::InvalidArgumentError(format!("Encountered non UTF-8 data: {e}")) })?; - // Verify each offset is at a valid character boundary in this UTF-8 array + // Verify each offset is at a valid character boundary for offset in offsets.iter() { let o = offset.as_usize(); - if !validated.is_char_boundary(o) { - if o < validated.len() { - return Err(ArrowError::InvalidArgumentError(format!( - "Split UTF-8 codepoint at offset {o}" - ))); - } + if !validated.is_char_boundary(o - start) { return Err(ArrowError::InvalidArgumentError(format!( - "Offset of {o} exceeds length of values {}", - validated.len() + "Split UTF-8 codepoint at offset {o}" ))); } } diff --git a/arrow-avro/Cargo.toml b/arrow-avro/Cargo.toml index 9e6581f199db..f78ba505455d 100644 --- a/arrow-avro/Cargo.toml +++ b/arrow-avro/Cargo.toml @@ -19,7 +19,6 @@ name = "arrow-avro" version = { workspace = true } description = "Support for parsing Avro format into the Arrow format" -homepage = { workspace = true } repository = { workspace = true } authors = { workspace = true } license = { workspace = true } diff --git a/arrow-avro/src/reader/mod.rs b/arrow-avro/src/reader/mod.rs index 499fabf42491..2faa90c83e25 100644 --- a/arrow-avro/src/reader/mod.rs +++ b/arrow-avro/src/reader/mod.rs @@ -34,15 +34,15 @@ //! //! This module exposes three layers of the API surface, from highest to lowest-level: //! -//! * [`ReaderBuilder`](crate::reader::ReaderBuilder): configures how Avro is read (batch size, strict union handling, +//! * [`ReaderBuilder`]: configures how Avro is read (batch size, strict union handling, //! string representation, reader schema, etc.) and produces either: //! * a `Reader` for **Avro Object Container Files (OCF)** read from any `BufRead`, or //! * a low-level `Decoder` for **unframed Avro datums**, **single‑object encoded** Avro //! bytes, and Confluent **Schema Registry** framed messages. -//! * [`Reader`](crate::reader::Reader): a convenient, synchronous iterator over `RecordBatch` decoded from an OCF +//! * [`Reader`]: a convenient, synchronous iterator over `RecordBatch` decoded from an OCF //! input. Implements [`Iterator>`] and //! `RecordBatchReader`. -//! * [`Decoder`](crate::reader::Decoder): a push‑based row decoder that consumes unframed or +//! * [`Decoder`]: a push‑based row decoder that consumes unframed or //! framed Avro bytes and yields ready `RecordBatch` values when batches fill. This is suitable //! for integrating with async byte streams, network protocols, or other custom data sources. //! @@ -55,7 +55,7 @@ //! * **Unframed binary datums**: Bare Avro records without an OCF header, schema fingerprint, //! or schema-registry prefix. Register the known writer schema in a `SchemaStore`, select it //! with [`ReaderBuilder::with_active_fingerprint`](crate::reader::ReaderBuilder::with_active_fingerprint), -//! configure [`DecoderMode::UnframedDatum`](crate::reader::DecoderMode::UnframedDatum) with +//! configure [`DecoderMode::UnframedDatum`] with //! [`ReaderBuilder::with_decoder_mode`](crate::reader::ReaderBuilder::with_decoder_mode), and //! call [`Decoder::decode`](crate::reader::Decoder::decode) once per record. This supports bare //! Kafka messages and consecutive records in one buffer. diff --git a/arrow-avro/src/writer/mod.rs b/arrow-avro/src/writer/mod.rs index 93835ecb0180..ebd654b0c4bb 100644 --- a/arrow-avro/src/writer/mod.rs +++ b/arrow-avro/src/writer/mod.rs @@ -63,7 +63,7 @@ //! //! By default, the writer converts your Arrow schema to Avro (including a top‑level record //! name). If you already have an Avro schema JSON you want to use verbatim, put it into the -//! Arrow schema metadata under the [`SCHEMA_METADATA_KEY`](crate::schema::SCHEMA_METADATA_KEY) +//! Arrow schema metadata under the [`SCHEMA_METADATA_KEY`] //! key before constructing the writer. The builder will use that schema instead of generating //! a new one. //! diff --git a/arrow-buffer/Cargo.toml b/arrow-buffer/Cargo.toml index 6f36d9a0fae7..465f4eee4898 100644 --- a/arrow-buffer/Cargo.toml +++ b/arrow-buffer/Cargo.toml @@ -19,7 +19,6 @@ name = "arrow-buffer" version = { workspace = true } description = "Buffer abstractions for Apache Arrow" -homepage = { workspace = true } repository = { workspace = true } authors = { workspace = true } license = { workspace = true } diff --git a/arrow-buffer/src/buffer/immutable.rs b/arrow-buffer/src/buffer/immutable.rs index 5d2a1f80f405..bd6837aa6af7 100644 --- a/arrow-buffer/src/buffer/immutable.rs +++ b/arrow-buffer/src/buffer/immutable.rs @@ -789,7 +789,7 @@ mod tests { } #[test] - #[expect(clippy::float_cmp)] + #[expect(clippy::float_cmp_const)] fn test_as_typed_data() { check_as_typed_data!(&[1i8, 3i8, 6i8], i8); check_as_typed_data!(&[1u8, 3u8, 6u8], u8); diff --git a/arrow-buffer/src/util/bit_chunk_iterator.rs b/arrow-buffer/src/util/bit_chunk_iterator.rs index 77ebe83045c1..15bdafc589c1 100644 --- a/arrow-buffer/src/util/bit_chunk_iterator.rs +++ b/arrow-buffer/src/util/bit_chunk_iterator.rs @@ -330,6 +330,18 @@ impl<'a> BitChunks<'a> { /// # Panics /// /// Panics if `index >= self.chunk_len()` + /// + /// # Example + /// + /// ``` + /// # use arrow_buffer::Buffer; + /// let buffer = Buffer::from((0..17_u8).collect::>()); + /// // Start at bit 8, so chunk 1 holds bytes 9 through 16 + /// let chunks = buffer.bit_chunks(8, 128); + /// assert_eq!(chunks.chunk_len(), 2); + /// assert_eq!(chunks.chunk(1), 0x100f0e0d0c0b0a09); + /// assert_eq!(chunks.iter().nth(1), Some(chunks.chunk(1))); + /// ``` #[inline] pub fn chunk(&self, index: usize) -> u64 { assert!(index < self.chunk_len, "chunk index out of bounds"); diff --git a/arrow-buffer/src/util/bit_util.rs b/arrow-buffer/src/util/bit_util.rs index ff19f4465b9f..e50e111bef9d 100644 --- a/arrow-buffer/src/util/bit_util.rs +++ b/arrow-buffer/src/util/bit_util.rs @@ -81,6 +81,78 @@ pub fn compress(value: u64, mask: u64) -> u64 { } } +/// Parallel bit deposit: scatter the lowest `mask.count_ones()` bits of +/// `value` into the set positions of `mask`, preserving their order. +/// All other bits in the result are zero; excess input bits are ignored. +/// +/// This is the inverse of [`compress`] on the selected bits: +/// `expand(compress(value, mask), mask) == value & mask`. +/// +/// Equivalent to the x86 BMI2 `PDEP` instruction. When compiled with the +/// `bmi2` target feature enabled (for example `-C target-cpu=x86-64-v3`) +/// this lowers to the hardware `pdep` instruction; otherwise it falls back +/// to a portable scalar loop that visits whichever is fewer: unset or set +/// bits in `mask`. +/// +/// # Functional Example +/// +/// Using 8 bits for brevity (the function operates on all 64). The low bits +/// of `value` are scattered into the set positions of `mask`: +/// +/// ```text +/// bit: 7 6 5 4 3 2 1 0 +/// value: 0 0 0 b c e f h +/// mask: 0 1 1 0 1 1 0 1 +/// result: 0 b c 0 e f 0 h +/// ``` +/// +/// # Code Example +/// +/// ``` +/// # use arrow_buffer::bit_util::{compress, expand}; +/// assert_eq!(expand(0b0000_1010, 0b0110_1101), 0b0010_0100); +/// let value = 0b1011_0100; +/// let mask = 0b0110_1101; +/// assert_eq!(expand(compress(value, mask), mask), value & mask); +/// ``` +#[inline] +pub fn expand(value: u64, mask: u64) -> u64 { + #[cfg(all(target_arch = "x86_64", target_feature = "bmi2"))] + { + // SAFETY: the `bmi2` target feature is statically enabled for this + // build, so the `pdep` instruction is guaranteed to be available. + unsafe { std::arch::x86_64::_pdep_u64(value, mask) } + } + + #[cfg(not(all(target_arch = "x86_64", target_feature = "bmi2")))] + { + let mut value = value; + if value == 0 { + return 0; + } + + let mut zeros = !mask; + if zeros.count_ones() <= 32 { + // Insert zeros from low to high; excess input bits shift out. + while zeros != 0 { + let lower = (1_u64 << zeros.trailing_zeros()) - 1; + value = (value & lower) | ((value & !lower) << 1); + zeros &= zeros - 1; + } + value + } else { + let mut output = 0; + let mut ones = mask; + while ones != 0 { + output |= (value & 1) << ones.trailing_zeros(); + value >>= 1; + ones &= ones - 1; + } + output + } + } +} + /// Returns the nearest number that is `>=` than `num` and is a multiple of 64 /// /// # Panics @@ -976,6 +1048,52 @@ mod tests { } } + #[test] + fn test_expand() { + fn reference(values: u64, mask: u64) -> u64 { + let mut expected = 0; + let mut input_idx = 0; + for output_idx in 0..64 { + if mask & (1 << output_idx) != 0 { + expected |= ((values >> input_idx) & 1) << output_idx; + input_idx += 1; + } + } + expected + } + + assert_eq!(expand(0b1010, 0b1111), 0b1010); + assert_eq!(expand(0b11, 0b1010), 0b1010); + assert_eq!(expand(u64::MAX, 0), 0); + assert_eq!(expand(0, u64::MAX), 0); + assert_eq!(expand(0, 0x5555_5555_5555_5555), 0); + assert_eq!(expand(u64::MAX, u64::MAX), u64::MAX); + + let mut rng = StdRng::seed_from_u64(0x2b7e_1516_28ae_d2a6); + + // Masks with at most one unset bit or at most one set bit + for bit in 0..64 { + for mask in [u64::MAX, !(1 << bit), 1 << bit, 0] { + for _ in 0..16 { + let values = rng.random::(); + assert_eq!(expand(values, mask), reference(values, mask), "{mask:#x}"); + } + } + } + + // Masks across the full density range, exercising the hardware `pdep` + // path on a `bmi2` build or both portable loops otherwise. + for _ in 0..20_000 { + let density = rng.random_range(0.0..=1.0); + let mask = (0..64).fold(0_u64, |mask, bit| { + mask | ((rng.random_bool(density) as u64) << bit) + }); + let values = rng.random::(); + assert_eq!(expand(compress(values, mask), mask), values & mask); + assert_eq!(expand(values, mask), reference(values, mask), "{mask:#x}"); + } + } + #[test] fn test_round_upto_multiple_of_64() { assert_eq!(0, round_upto_multiple_of_64(0)); diff --git a/arrow-cast/Cargo.toml b/arrow-cast/Cargo.toml index 9e4a43a8d978..dfe10803c67a 100644 --- a/arrow-cast/Cargo.toml +++ b/arrow-cast/Cargo.toml @@ -19,7 +19,6 @@ name = "arrow-cast" version = { workspace = true } description = "Cast kernel and utilities for Apache Arrow" -homepage = { workspace = true } repository = { workspace = true } authors = { workspace = true } license = { workspace = true } @@ -55,6 +54,9 @@ base64 = "0.23" ryu = "1.0.16" [dev-dependencies] +# Enable `chrono-tz` so that tests can exercise IANA (named) timezones such as +# `America/New_York`, which is where DST transitions can be observed. +arrow-array = { workspace = true, features = ["chrono-tz"] } criterion = { workspace = true, default-features = false } half = { version = "2.1", default-features = false } insta = { workspace = true } @@ -85,9 +87,17 @@ harness = false name = "format_temporal" harness = false +[[bench]] +name = "cast_timestamp" +harness = false + [lints] workspace = true [[bench]] name = "format_decimal" harness = false + +[[bench]] +name = "format_boolean" +harness = false diff --git a/arrow-cast/benches/cast_timestamp.rs b/arrow-cast/benches/cast_timestamp.rs new file mode 100644 index 000000000000..5e516aae8b78 --- /dev/null +++ b/arrow-cast/benches/cast_timestamp.rs @@ -0,0 +1,75 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::hint::black_box; + +use arrow_array::{PrimitiveArray, types::*}; +use arrow_cast::cast; +use arrow_schema::{DataType, TimeUnit}; +use criterion::{Criterion, Throughput, criterion_group, criterion_main}; +use rand::rngs::StdRng; +use rand::{RngExt, SeedableRng}; + +const SIZE: usize = 8192; + +fn benchmark(c: &mut Criterion, unit: &str, units_per_second: i64) { + // Timestamps spread over roughly 63 years either side of the epoch. + let mut rng = StdRng::seed_from_u64(42); + let range = -2_000_000_000 * units_per_second..2_000_000_000 * units_per_second; + let values: Vec = (0..SIZE).map(|_| rng.random_range(range.clone())).collect(); + let mixed_nulls = PrimitiveArray::::from_iter( + values + .iter() + .enumerate() + .map(|(i, &v)| (i % 5 != 0).then_some(v)), + ); + let no_nulls = PrimitiveArray::::new(values.into(), None); + // Timezone-aware casts go through Chrono and serve as a control. + let zoned = no_nulls.clone().with_timezone("+05:45"); + let inputs = [ + ("no_nulls", &no_nulls), + ("mixed_nulls", &mixed_nulls), + ("timezone_control", &zoned), + ]; + let targets = [ + ("date32", DataType::Date32), + ("time32_s", DataType::Time32(TimeUnit::Second)), + ("time32_ms", DataType::Time32(TimeUnit::Millisecond)), + ("time64_us", DataType::Time64(TimeUnit::Microsecond)), + ("time64_ns", DataType::Time64(TimeUnit::Nanosecond)), + ]; + let mut group = c.benchmark_group(format!("timestamp_{unit}")); + group.throughput(Throughput::Elements(SIZE as u64)); + for (name, target) in &targets { + for (input, array) in inputs { + group.bench_function(format!("{name}/{input}"), |b| { + b.iter(|| black_box(cast(black_box(array), target).unwrap())) + }); + } + } + group.finish(); +} + +fn cast_timestamp(c: &mut Criterion) { + benchmark::(c, "s", 1); + benchmark::(c, "ms", 1_000); + benchmark::(c, "us", 1_000_000); + benchmark::(c, "ns", 1_000_000_000); +} + +criterion_group!(benches, cast_timestamp); +criterion_main!(benches); diff --git a/arrow-cast/benches/format_boolean.rs b/arrow-cast/benches/format_boolean.rs new file mode 100644 index 000000000000..61c0da7dde04 --- /dev/null +++ b/arrow-cast/benches/format_boolean.rs @@ -0,0 +1,48 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. +use std::hint::black_box; + +use arrow_array::BooleanArray; +use arrow_cast::display::{ArrayFormatter, FormatOptions}; +use criterion::{Criterion, Throughput, criterion_group, criterion_main}; +use rand::rngs::StdRng; +use rand::{RngExt, SeedableRng}; + +const ARRAY_LEN: usize = 8192; + +fn format_array(c: &mut Criterion) { + let mut rng = StdRng::seed_from_u64(42); + let array = BooleanArray::from((0..ARRAY_LEN).map(|_| rng.random()).collect::>()); + let formatter = ArrayFormatter::try_new(&array, &FormatOptions::new()).unwrap(); + let mut output = String::with_capacity(8); + + let mut group = c.benchmark_group("format_boolean"); + group.throughput(Throughput::Elements(ARRAY_LEN as u64)); + group.bench_function("random", |b| { + b.iter(|| { + for idx in 0..array.len() { + output.clear(); + formatter.value(idx).write(&mut output).unwrap(); + } + black_box(&output); + }) + }); + group.finish(); +} + +criterion_group!(benches, format_array); +criterion_main!(benches); diff --git a/arrow-cast/src/cast/binary.rs b/arrow-cast/src/cast/binary.rs new file mode 100644 index 000000000000..72ce5ae1a946 --- /dev/null +++ b/arrow-cast/src/cast/binary.rs @@ -0,0 +1,198 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Cast support for binary, fixed-size binary and byte-view arrays. + +use arrow_array::{builder::*, cast::*, types::*, *}; +use arrow_buffer::{Buffer, OffsetBuffer}; +use arrow_data::{ArrayData, ByteView}; +use arrow_schema::{ArrowError, DataType}; +use num_traits::{NumCast, ToPrimitive}; +use std::sync::Arc; + +use super::CastOptions; + +pub(crate) fn cast_numeric_to_binary( + array: &dyn Array, +) -> Result { + let array = array.as_primitive::(); + let size = std::mem::size_of::(); + let offsets = OffsetBuffer::from_repeated_length(size, array.len()); + Ok(Arc::new(GenericBinaryArray::::try_new( + offsets, + array.values().inner().clone(), + array.nulls().cloned(), + )?)) +} + +/// Helper function to cast from one `BinaryArray` or 'LargeBinaryArray' to 'FixedSizeBinaryArray'. +pub(crate) fn cast_binary_to_fixed_size_binary( + array: &dyn Array, + byte_width: i32, + cast_options: &CastOptions, +) -> Result { + let array = array.as_binary::(); + let mut builder = FixedSizeBinaryBuilder::with_capacity(array.len(), byte_width); + + for i in 0..array.len() { + if array.is_null(i) { + builder.append_null(); + } else { + match builder.append_value(array.value(i)) { + Ok(()) => {} + Err(e) => match cast_options.safe { + true => builder.append_null(), + false => return Err(e), + }, + } + } + } + + Ok(Arc::new(builder.finish())) +} + +/// Helper function to cast from 'FixedSizeBinaryArray' to one `BinaryArray` or 'LargeBinaryArray'. +/// If the target one is too large for the source array it will return an Error. +pub(crate) fn cast_fixed_size_binary_to_binary( + array: &dyn Array, + byte_width: i32, +) -> Result { + let array = array + .as_any() + .downcast_ref::() + .unwrap(); + + let offsets: i128 = byte_width as i128 * array.len() as i128; + + let is_binary = matches!(GenericBinaryType::::DATA_TYPE, DataType::Binary); + if is_binary && offsets > i32::MAX as i128 { + return Err(ArrowError::ComputeError( + "FixedSizeBinary array too large to cast to Binary array".to_string(), + )); + } else if !is_binary && offsets > i64::MAX as i128 { + return Err(ArrowError::ComputeError( + "FixedSizeBinary array too large to cast to LargeBinary array".to_string(), + )); + } + + let mut builder = GenericBinaryBuilder::::with_capacity(array.len(), array.len()); + + for i in 0..array.len() { + if array.is_null(i) { + builder.append_null(); + } else { + builder.append_value(array.value(i)); + } + } + + Ok(Arc::new(builder.finish())) +} + +pub(crate) fn cast_fixed_size_binary_to_binary_view( + array: &dyn Array, + _byte_width: i32, +) -> Result { + let array = array + .as_any() + .downcast_ref::() + .unwrap(); + + let mut builder = BinaryViewBuilder::with_capacity(array.len()); + for i in 0..array.len() { + if array.is_null(i) { + builder.append_null(); + } else { + builder.append_value(array.value(i)); + } + } + + Ok(Arc::new(builder.finish())) +} + +/// Helper function to cast from one `ByteArrayType` to another and vice versa. +/// If the target one (e.g., `LargeUtf8`) is too large for the source array it will return an Error. +pub(crate) fn cast_byte_container(array: &dyn Array) -> Result +where + FROM: ByteArrayType, + TO: ByteArrayType, + FROM::Offset: OffsetSizeTrait + ToPrimitive, + TO::Offset: OffsetSizeTrait + NumCast, +{ + let data = array.to_data(); + assert_eq!(data.data_type(), &FROM::DATA_TYPE); + let str_values_buf = data.buffers()[1].clone(); + let offsets = data.buffers()[0].typed_data::(); + + let mut cast_offsets = Vec::::with_capacity(offsets.len()); + offsets + .iter() + .try_for_each::<_, Result<_, ArrowError>>(|offset| { + let offset = + <::Offset as NumCast>::from(*offset).ok_or_else(|| { + ArrowError::ComputeError(format!( + "{}{} array too large to cast to {}{} array", + FROM::Offset::PREFIX, + FROM::PREFIX, + TO::Offset::PREFIX, + TO::PREFIX + )) + })?; + cast_offsets.push(offset); + Ok(()) + })?; + + let offset_buffer = Buffer::from_vec(cast_offsets); + + let dtype = TO::DATA_TYPE; + + let builder = ArrayData::builder(dtype) + .offset(array.offset()) + .len(array.len()) + .add_buffer(offset_buffer) + .add_buffer(str_values_buf) + .nulls(data.nulls().cloned()); + + let array_data = unsafe { builder.build_unchecked() }; + + Ok(Arc::new(GenericByteArray::::from(array_data))) +} + +/// Helper function to cast from one `ByteViewType` array to `ByteArrayType` array. +pub(crate) fn cast_view_to_byte(array: &dyn Array) -> Result +where + FROM: ByteViewType, + TO: ByteArrayType, + FROM::Native: AsRef, +{ + let data = array.to_data(); + let view_array = GenericByteViewArray::::from(data); + + let len = view_array.len(); + let bytes = view_array + .views() + .iter() + .map(|v| ByteView::from(*v).length as usize) + .sum::(); + + let mut byte_array_builder = GenericByteBuilder::::with_capacity(len, bytes); + + for val in &view_array { + byte_array_builder.append_option(val); + } + + Ok(Arc::new(byte_array_builder.finish())) +} diff --git a/arrow-cast/src/cast/decimal.rs b/arrow-cast/src/cast/decimal.rs index a5d704bba8d5..be98b4b4dfe7 100644 --- a/arrow-cast/src/cast/decimal.rs +++ b/arrow-cast/src/cast/decimal.rs @@ -172,7 +172,7 @@ where I::Native: DecimalCast + ArrowNativeTypeOp, O::Native: DecimalCast + ArrowNativeTypeOp, { - let delta_scale = output_scale - input_scale; + let delta_scale = output_scale as i16 - input_scale as i16; // O::MAX_FOR_EACH_PRECISION[k] stores 10^k - 1 (e.g., 9, 99, 999, ...). // Adding 1 yields exactly 10^k without computing a power at runtime. @@ -188,7 +188,7 @@ where // then an increase of scale by 3 will have the following effect on the representation: // [xxxxx] -> [xxxxx000], so for the cast to be infallible, the output type // needs to provide at least 8 digits precision - let is_infallible_cast = (input_precision as i8) + delta_scale <= (output_precision as i8); + let is_infallible_cast = (input_precision as i16) + delta_scale <= (output_precision as i16); let f_infallible = is_infallible_cast .then_some(move |x| O::Native::from_decimal(x).unwrap().mul_wrapping(mul)); Some((f_fallible, f_infallible)) @@ -221,7 +221,7 @@ where I::Native: DecimalCast + ArrowNativeTypeOp, O::Native: DecimalCast + ArrowNativeTypeOp, { - let delta_scale = input_scale - output_scale; + let delta_scale = input_scale as i16 - output_scale as i16; // delta_scale is guaranteed to be > 0, but may also be larger than I::MAX_PRECISION. If so, the // scale change divides out more digits than the input has precision and the result of the cast @@ -259,7 +259,7 @@ where // the output type needs to have at least 3 digits of precision. // e.g. Decimal(5, 3) 99.999 to Decimal(3, 0) will result in 100: // [99999] -> [99] + 1 = [100], a cast to Decimal(2, 0) would not be possible - let is_infallible_cast = (input_precision as i8) - delta_scale < (output_precision as i8); + let is_infallible_cast = (input_precision as i16) - delta_scale < (output_precision as i16); let f_infallible = is_infallible_cast.then_some(move |x| f_fallible(x).unwrap()); Some((f_fallible, f_infallible)) } @@ -910,6 +910,26 @@ mod tests { assert_eq!(result, None); } + #[test] + fn test_rescale_decimal256_extreme_scales() { + assert_eq!( + rescale_decimal::(i256::ONE, 76, -76, 76, 0), + None + ); + assert_eq!( + rescale_decimal::(i256::ZERO, 76, -76, 76, 0), + Some(i256::ZERO) + ); + assert_eq!( + rescale_decimal::(i256::ONE, 76, i8::MIN, 76, 76), + None + ); + assert_eq!( + rescale_decimal::(i256::ONE, 76, 76, 76, i8::MIN), + Some(i256::ZERO) + ); + } + #[test] fn test_rescale_decimal_invalid_input_precision_scale_returns_none() { let result = rescale_decimal::(123_i128, 39, 39, 38, 38); diff --git a/arrow-cast/src/cast/list.rs b/arrow-cast/src/cast/list.rs index ee35b1a9f39e..e34f808b69d0 100644 --- a/arrow-cast/src/cast/list.rs +++ b/arrow-cast/src/cast/list.rs @@ -16,6 +16,7 @@ // under the License. use crate::cast::*; +use arrow_buffer::ScalarBuffer; /// Converts a non-list array to a list array where every element is a single element /// list. `NULL`s in the original array become `[NULL]` (i.e. output list array @@ -282,15 +283,72 @@ pub(crate) fn cast_list_values( cast_options: &CastOptions, ) -> Result { let list = array.as_list::(); - let values = cast_with_options(list.values(), to.data_type(), cast_options)?; + let (offsets, values) = list_visible_values(list); + let values = cast_with_options(&values, to.data_type(), cast_options)?; Ok(Arc::new(GenericListArray::::try_new( to.clone(), - list.offsets().clone(), + offsets, values, list.nulls().cloned(), )?)) } +/// Returns the part of `list.values()` that the rows use, with the offsets +/// adjusted to start at 0. For a sliced list, this drops the values outside +/// the slice. +fn list_visible_values( + list: &GenericListArray, +) -> (OffsetBuffer, ArrayRef) { + let offsets = list.offsets(); + let first = offsets.first(); + let values = list + .values() + .slice(first.as_usize(), (offsets.last() - first).as_usize()); + (offsets.clone().subtract(first), values) +} + +/// Returns the values of `list_view` from the lowest offset to the highest end of +/// its rows, and the offsets of the rows into those values. As rows can be in any +/// order, these values can include some that no row uses. +fn list_view_visible_values( + list_view: &GenericListViewArray, +) -> (ScalarBuffer, ArrayRef) { + let offsets = list_view.offsets(); + let sizes = list_view.sizes(); + let values = list_view.values(); + + let mut start = usize::MAX; + let mut end = 0; + for (offset, size) in offsets.iter().zip(sizes.iter()) { + if size.as_usize() > 0 { + start = start.min(offset.as_usize()); + end = end.max(offset.as_usize() + size.as_usize()); + } + } + if end == 0 { + // All rows are empty + start = 0; + } + + if start == 0 && end == values.len() { + return (offsets.clone(), Arc::clone(values)); + } + + // The offsets of empty rows can be outside the new range, so set them to 0 + let offsets = offsets + .iter() + .zip(sizes.iter()) + .map(|(offset, size)| { + if size.as_usize() > 0 { + O::usize_as(offset.as_usize() - start) + } else { + O::usize_as(0) + } + }) + .collect(); + (offsets, values.slice(start, end - start)) +} + /// Casting between list view arrays of same offset size; we cast only the inner type. pub(crate) fn cast_list_view_values( array: &dyn Array, @@ -298,10 +356,11 @@ pub(crate) fn cast_list_view_values( cast_options: &CastOptions, ) -> Result { let list = array.as_list_view::(); - let values = cast_with_options(list.values(), to.data_type(), cast_options)?; + let (offsets, values) = list_view_visible_values(list); + let values = cast_with_options(&values, to.data_type(), cast_options)?; Ok(Arc::new(GenericListViewArray::::try_new( to.clone(), - list.offsets().clone(), + offsets, list.sizes().clone(), values, list.nulls().cloned(), @@ -315,8 +374,7 @@ pub(crate) fn cast_list( cast_options: &CastOptions, ) -> Result { let list = array.as_list::(); - let values = list.values(); - let offsets = list.offsets(); + let (offsets, values) = list_visible_values(list); let nulls = list.nulls().cloned(); if offsets.last().as_usize() > O::MAX_OFFSET { @@ -328,7 +386,7 @@ pub(crate) fn cast_list( } // Recursively cast values - let values = cast_with_options(values, field.data_type(), cast_options)?; + let values = cast_with_options(&values, field.data_type(), cast_options)?; let offsets: Vec<_> = offsets.iter().map(|x| O::usize_as(x.as_usize())).collect(); // Safety: valid offsets and checked for overflow @@ -409,12 +467,12 @@ pub(crate) fn cast_list_view( cast_options: &CastOptions, ) -> Result { let list_view = array.as_list_view::(); + let (offsets, values) = list_view_visible_values(list_view); // Recursively cast values - let values = cast_with_options(list_view.values(), to_field.data_type(), cast_options)?; + let values = cast_with_options(&values, to_field.data_type(), cast_options)?; - let offsets = list_view - .offsets() + let offsets = offsets .iter() .map(|offset| { let offset = offset.as_usize(); @@ -459,7 +517,8 @@ pub(crate) fn cast_list_to_list_view( cast_options: &CastOptions, ) -> Result { let list = array.as_list::(); - let (_field, offsets, values, nulls) = list.clone().into_parts(); + let (offsets, values) = list_visible_values(list); + let nulls = list.nulls().cloned(); let len = offsets.len() - 1; let mut sizes = Vec::with_capacity(len); @@ -480,12 +539,12 @@ pub(crate) fn cast_list_to_list_view( sizes.push(O::usize_as(size)); } let values = cast_with_options(&values, to_field.data_type(), cast_options)?; - let array = GenericListViewArray::::new( + let array = GenericListViewArray::::try_new( to_field.clone(), view_offsets.into(), sizes.into(), values, nulls, - ); + )?; Ok(Arc::new(array)) } diff --git a/arrow-cast/src/cast/map.rs b/arrow-cast/src/cast/map.rs index 7395d32efe5f..d4a3ae161141 100644 --- a/arrow-cast/src/cast/map.rs +++ b/arrow-cast/src/cast/map.rs @@ -37,16 +37,22 @@ pub(crate) fn cast_map_values( "map is missing value field".to_string(), ))?; - let key_array = cast_with_options(from.keys(), key_field.data_type(), cast_options)?; - let value_array = cast_with_options(from.values(), value_field.data_type(), cast_options)?; + let offsets = from.offsets(); + let first = offsets.first(); + // Only the entries that the offsets refer to, which excludes entries outside a slice + let entries = from + .entries() + .slice(first.as_usize(), (offsets.last() - first).as_usize()); + let key_array = cast_with_options(entries.column(0), key_field.data_type(), cast_options)?; + let value_array = cast_with_options(entries.column(1), value_field.data_type(), cast_options)?; Ok(Arc::new(MapArray::try_new( entries_field.clone(), - from.offsets().clone(), + offsets.clone().subtract(first), StructArray::try_new( Fields::from(vec![key_field, value_field]), vec![key_array, value_array], - from.entries().nulls().cloned(), + entries.nulls().cloned(), )?, from.nulls().cloned(), to_ordered, diff --git a/arrow-cast/src/cast/mod.rs b/arrow-cast/src/cast/mod.rs index 6eed69976e1a..f008bbcae4d9 100644 --- a/arrow-cast/src/cast/mod.rs +++ b/arrow-cast/src/cast/mod.rs @@ -37,6 +37,7 @@ //! assert_eq!(7.0, c.value(2)); //! ``` +mod binary; mod decimal; mod dictionary; mod list; @@ -46,6 +47,7 @@ mod string; mod structs; mod union; +use crate::cast::binary::*; use crate::cast::decimal::*; use crate::cast::dictionary::*; use crate::cast::list::*; @@ -56,8 +58,9 @@ use crate::cast::structs::*; pub use crate::cast::union::*; use arrow_buffer::IntervalMonthDayNano; -use arrow_data::ByteView; -use chrono::{NaiveTime, Offset, TimeZone, Utc}; +use chrono::{ + FixedOffset, LocalResult, NaiveDateTime, NaiveTime, Offset, TimeDelta, TimeZone, Utc, +}; use std::cmp::Ordering; use std::sync::Arc; @@ -67,7 +70,7 @@ use crate::parse::{ string_to_datetime, }; use arrow_array::{builder::*, cast::*, temporal_conversions::*, timezone::Tz, types::*, *}; -use arrow_buffer::{ArrowNativeType, Buffer, OffsetBuffer, i256}; +use arrow_buffer::{ArrowNativeType, OffsetBuffer, i256}; use arrow_data::ArrayData; use arrow_data::transform::MutableArrayData; use arrow_schema::*; @@ -359,6 +362,34 @@ where M::from_decimal(value.into()) } +/// Returns true if every value of `array` has at most `digits` decimal digits. +/// +/// If the source type cannot hold that many digits, every value fits; +/// otherwise the data is scanned, including null slots, since `unary` +/// converts those as well. +fn integers_fit_digits(array: &PrimitiveArray, digits: u32) -> bool +where + T: ArrowPrimitiveType, + T::Native: ArrowNativeTypeOp, +{ + let Ok(bound) = T::Native::usize_as(10).pow_checked(digits) else { + // The type's maximum is below the bound, and its minimum, one larger + // in magnitude, is a power of two that no power of ten can equal. + return true; + }; + let hi = bound.sub_wrapping(T::Native::ONE); + // Unsigned types have no negative values to bound. + let lo = T::Native::ZERO.sub_checked(hi).unwrap_or(T::Native::ZERO); + + // The scan works in chunks, so an array that does not fit is rejected after + // its first failing chunk. Each chunk is checked without branching, so it + // vectorizes. + array + .values() + .chunks(64) + .all(|chunk| chunk.iter().fold(true, |ok, &v| ok & (v >= lo) & (v <= hi))) +} + fn cast_integer_to_decimal< T: ArrowPrimitiveType, D: DecimalType + ArrowPrimitiveType, @@ -381,6 +412,23 @@ where )) }; + // An invalid precision or scale would let the bound below admit values the + // decimal native type cannot hold, so reject it before choosing a kernel. + validate_decimal_precision_and_scale::(precision, scale)?; + + // Scan for values that don't fit in the target type before choosing a + // kernel: in the common case that all values fit, the cast is infallible + // and we can use the `unary` kernel, which is very fast. Empirically, the + // cost of doing the out-of-range scan is small relative to the win from + // using a faster kernel. + // + // A value fits when it has at most `precision - scale` decimal digits: a + // nonnegative scale appends that many trailing zeros, and a negative scale + // removes digits first. Validation guarantees the difference is not + // negative. + let digits = (precision as i32 - scale as i32) as u32; + let all_fit = integers_fit_digits(array, digits); + let array = if scale < 0 { // Compute the scale factor once in the source type. Scaling before the // checked conversion permits values that only fit the decimal native @@ -390,6 +438,10 @@ where .ok(); match (scale_factor, cast_options.safe) { + (Some(scale_factor), _) if all_fit => array.unary::<_, D>(|v| { + integer_to_decimal_native::<_, M>(v.div_wrapping(scale_factor)) + .expect("value fits the decimal") + }), (Some(scale_factor), true) => array.unary_opt::<_, D>(|v| { let v = v .div_checked(scale_factor) @@ -421,18 +473,25 @@ where )) })?; - match cast_options.safe { - true => array.unary_opt::<_, D>(|v| { + if all_fit { + array.unary::<_, D>(|v| { + integer_to_decimal_native::<_, M>(v) + .expect("value fits the decimal") + .mul_wrapping(scale_factor) + }) + } else if cast_options.safe { + array.unary_opt::<_, D>(|v| { let v = integer_to_decimal_native::<_, M>(v) .and_then(|v| v.mul_checked(scale_factor).ok())?; (D::is_valid_decimal_precision(v, precision)).then_some(v) - }), - false => array.try_unary::<_, D, _>(|v| { + }) + } else { + array.try_unary::<_, D, _>(|v| { let v = integer_to_decimal_native::<_, M>(v) .ok_or_else(|| overflow(v)) .and_then(|v| v.mul_checked(scale_factor))?; D::validate_decimal_precision(v, precision, scale).map(|()| v) - })?, + })? } }; @@ -608,30 +667,128 @@ fn make_duration_array(array: &PrimitiveArray, unit: TimeUnit) -> Arr } } -fn as_time_res_with_timezone( - v: i64, - tz: Option, -) -> Result { - let time = match tz { - Some(tz) => as_datetime_with_timezone::(v, tz).map(|d| d.time()), - None => as_datetime::(v).map(|d| d.time()), - }; +/// Casts timestamps to the time of day in the unit of `to_type`. +fn cast_timestamp_to_time( + array: &dyn Array, + to_type: &DataType, + cast_options: &CastOptions, +) -> Result { + let array = array.as_primitive::(); + match to_type { + DataType::Time32(TimeUnit::Second) => { + timestamp_to_time::(array, cast_options) + } + DataType::Time32(TimeUnit::Millisecond) => { + timestamp_to_time::(array, cast_options) + } + DataType::Time64(TimeUnit::Microsecond) => { + timestamp_to_time::(array, cast_options) + } + DataType::Time64(TimeUnit::Nanosecond) => { + timestamp_to_time::(array, cast_options) + } + _ => Err(ArrowError::CastError(format!( + "Casting from {} to {to_type} not supported", + array.data_type() + ))), + } +} - time.ok_or_else(|| { - ArrowError::CastError(format!( - "Failed to create naive time with {} {}", - std::any::type_name::(), - v - )) - }) +/// A `Time32` or `Time64` type with a fixed unit and a Chrono conversion. +trait TimeType: ArrowTemporalType { + /// Number of units in one second. + const UNIT_MULTIPLE: i64; + + fn from_naive_time(time: NaiveTime) -> Self::Native; +} + +impl TimeType for Time32SecondType { + const UNIT_MULTIPLE: i64 = 1; + + fn from_naive_time(time: NaiveTime) -> Self::Native { + time_to_time32s(time) + } +} + +impl TimeType for Time32MillisecondType { + const UNIT_MULTIPLE: i64 = MILLISECONDS; + + fn from_naive_time(time: NaiveTime) -> Self::Native { + time_to_time32ms(time) + } +} + +impl TimeType for Time64MicrosecondType { + const UNIT_MULTIPLE: i64 = MICROSECONDS; + + fn from_naive_time(time: NaiveTime) -> Self::Native { + time_to_time64us(time) + } +} + +impl TimeType for Time64NanosecondType { + const UNIT_MULTIPLE: i64 = NANOSECONDS; + + fn from_naive_time(time: NaiveTime) -> Self::Native { + time_to_time64ns(time) + } +} + +/// Casts timestamps to the time of day in the unit of `O`. +fn timestamp_to_time( + array: &PrimitiveArray, + cast_options: &CastOptions, +) -> Result +where + T: ArrowTimestampType, + O: TimeType, + i64: AsPrimitive, +{ + // A time within one day fits its Time32 or Time64 representation. + let array = match array.timezone() { + Some(tz) => { + let tz: Tz = tz.parse()?; + let time = |v: i64| { + as_datetime_with_timezone::(v, tz).map(|d| O::from_naive_time(d.time())) + }; + if cast_options.safe { + array.unary_opt::<_, O>(time) + } else { + array.try_unary::<_, O, _>(|v| { + time(v).ok_or_else(|| { + ArrowError::CastError(format!( + "Failed to create naive time with {} {}", + std::any::type_name::(), + v + )) + }) + })? + } + } + None => array.unary::<_, O>(|v| { + // The remainder within a day is the time of day; `rem_euclid` keeps it + // nonnegative for timestamps before the epoch. The units are constants, + // so the branch and the divisions fold at compile time. + let from = time_unit_multiple(&T::UNIT); + let time = v.rem_euclid(SECONDS_IN_DAY * from); + let time = if from >= O::UNIT_MULTIPLE { + time / (from / O::UNIT_MULTIPLE) + } else { + time * (O::UNIT_MULTIPLE / from) + }; + time.as_() + }), + }; + Ok(Arc::new(array)) } fn timestamp_to_date32( array: &PrimitiveArray, + cast_options: &CastOptions, ) -> Result { let err = |x: i64| { ArrowError::CastError(format!( - "Cannot convert {} {x} to datetime", + "Cannot convert {} {x} to Date32", std::any::type_name::() )) }; @@ -639,17 +796,41 @@ fn timestamp_to_date32( let array: Date32Array = match array.timezone() { Some(tz) => { let tz: Tz = tz.parse()?; - array.try_unary(|x| { + let date = |x: i64| { as_datetime_with_timezone::(x, tz) - .ok_or_else(|| err(x)) .map(|d| Date32Type::from_naive_date(d.date_naive())) - })? + }; + if cast_options.safe { + array.unary_opt(date) + } else { + array.try_unary(|x| date(x).ok_or_else(|| err(x)))? + } + } + None => { + // Date32 stores days since the epoch. Round down so that a timestamp + // just before the epoch belongs to the preceding day. + let days = |x: i64| x.div_euclid(units_per_day::()); + let all_in_range = match T::UNIT { + // Every microsecond or nanosecond timestamp lies within the Date32 range. + TimeUnit::Microsecond | TimeUnit::Nanosecond => true, + // A branch-free scan lets the common case skip the per-value check. + _ => { + let day = units_per_day::(); + let (lo, hi) = (i32::MIN as i64 * day, (i32::MAX as i64 + 1) * day); + array + .values() + .iter() + .fold(true, |ok, &x| ok & (lo <= x) & (x < hi)) + } + }; + if all_in_range { + array.unary(|x| days(x) as i32) + } else if cast_options.safe { + array.unary_opt(|x| i32::try_from(days(x)).ok()) + } else { + array.try_unary(|x| i32::try_from(days(x)).map_err(|_| err(x)))? + } } - None => array.try_unary(|x| { - as_datetime::(x) - .ok_or_else(|| err(x)) - .map(|d| Date32Type::from_naive_date(d.date())) - })?, }; Ok(Arc::new(array)) } @@ -661,7 +842,7 @@ fn timestamp_to_date32( /// Accepts [`CastOptions`] to specify cast behavior. See also [`cast()`]. /// /// # Behavior -/// * `Boolean` to `Utf8`: `true` => '1', `false` => `0` +/// * `Boolean` to `Utf8`, `LargeUtf8`, or `Utf8View`: `true` => `"true"`, `false` => `"false"` /// * `Utf8` to `Boolean`: `true`, `yes`, `on`, `1` => `true`, `false`, `no`, `off`, `0` => `false`, /// short variants are accepted, other strings return null or error /// * `Utf8` to Numeric: strings that can't be parsed to numbers return null, float strings @@ -674,6 +855,10 @@ fn timestamp_to_date32( /// * `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 +/// * `Timestamp` to `Date32`, `Time32`, or `Time64`: timestamps without a timezone +/// support the full `i64` range; timestamps with a timezone must lie within Chrono's +/// date range. If safe is true, a timestamp outside that range or a `Date32` day +/// count that does not fit in `i32` becomes NULL, otherwise an error is returned /// * Temporal to/from backing Primitive: zero-copy with data type change /// * `Float16/Float32/Float64` to `Decimal(precision, scale)` rounds to the `scale` decimals /// (i.e. casting `6.4999` to `Decimal(10, 1)` becomes `6.5`). @@ -1270,9 +1455,9 @@ pub fn cast_with_options( Float16 => cast_bool_to_numeric::(array, cast_options), Float32 => cast_bool_to_numeric::(array, cast_options), Float64 => cast_bool_to_numeric::(array, cast_options), - Utf8View => value_to_string_view(array, cast_options), - Utf8 => value_to_string::(array, cast_options), - LargeUtf8 => value_to_string::(array, cast_options), + Utf8View => cast_bool_to_string::(array), + Utf8 => cast_bool_to_string::(array), + LargeUtf8 => cast_bool_to_string::(array), _ => Err(ArrowError::CastError(format!( "Casting from {from_type} to {to_type} not supported", ))), @@ -1921,6 +2106,11 @@ pub fn cast_with_options( // unchanged. // // i.e. Timestamp('2001-01-01T00:00', None) -> Timestamp('2001-01-01T00:00', '+0700') + // + // For IANA timezones a wall clock reading is not always a unique + // instant: a "fall back" DST transition makes an hour occur twice + // and a "spring forward" transition skips an hour entirely. See + // `resolve_local_offset` for how those readings are resolved. (None, Some(to_tz)) => { let to_tz: Tz = to_tz.parse()?; match to_unit { @@ -1950,18 +2140,21 @@ pub fn cast_with_options( }; Ok(make_timestamp_array(&adjusted, *to_unit, to_tz.clone())) } - (Timestamp(TimeUnit::Microsecond, _), Date32) => { - timestamp_to_date32(array.as_primitive::()) - } - (Timestamp(TimeUnit::Millisecond, _), Date32) => { - timestamp_to_date32(array.as_primitive::()) - } + (Timestamp(TimeUnit::Microsecond, _), Date32) => timestamp_to_date32( + array.as_primitive::(), + cast_options, + ), + (Timestamp(TimeUnit::Millisecond, _), Date32) => timestamp_to_date32( + array.as_primitive::(), + cast_options, + ), (Timestamp(TimeUnit::Second, _), Date32) => { - timestamp_to_date32(array.as_primitive::()) - } - (Timestamp(TimeUnit::Nanosecond, _), Date32) => { - timestamp_to_date32(array.as_primitive::()) + timestamp_to_date32(array.as_primitive::(), cast_options) } + (Timestamp(TimeUnit::Nanosecond, _), Date32) => timestamp_to_date32( + array.as_primitive::(), + cast_options, + ), (Timestamp(TimeUnit::Second, _), Date64) => Ok(Arc::new(match cast_options.safe { true => { // change error to None @@ -1986,197 +2179,17 @@ pub fn cast_with_options( .as_primitive::() .unary::<_, Date64Type>(|x| x / (NANOSECONDS / MILLISECONDS)), )), - (Timestamp(TimeUnit::Second, tz), Time64(TimeUnit::Microsecond)) => { - let tz = tz.as_ref().map(|tz| tz.parse()).transpose()?; - Ok(Arc::new( - array - .as_primitive::() - .try_unary::<_, Time64MicrosecondType, ArrowError>(|x| { - Ok(time_to_time64us(as_time_res_with_timezone::< - TimestampSecondType, - >(x, tz)?)) - })?, - )) - } - (Timestamp(TimeUnit::Second, tz), Time64(TimeUnit::Nanosecond)) => { - let tz = tz.as_ref().map(|tz| tz.parse()).transpose()?; - Ok(Arc::new( - array - .as_primitive::() - .try_unary::<_, Time64NanosecondType, ArrowError>(|x| { - Ok(time_to_time64ns(as_time_res_with_timezone::< - TimestampSecondType, - >(x, tz)?)) - })?, - )) - } - (Timestamp(TimeUnit::Millisecond, tz), Time64(TimeUnit::Microsecond)) => { - let tz = tz.as_ref().map(|tz| tz.parse()).transpose()?; - Ok(Arc::new( - array - .as_primitive::() - .try_unary::<_, Time64MicrosecondType, ArrowError>(|x| { - Ok(time_to_time64us(as_time_res_with_timezone::< - TimestampMillisecondType, - >(x, tz)?)) - })?, - )) - } - (Timestamp(TimeUnit::Millisecond, tz), Time64(TimeUnit::Nanosecond)) => { - let tz = tz.as_ref().map(|tz| tz.parse()).transpose()?; - Ok(Arc::new( - array - .as_primitive::() - .try_unary::<_, Time64NanosecondType, ArrowError>(|x| { - Ok(time_to_time64ns(as_time_res_with_timezone::< - TimestampMillisecondType, - >(x, tz)?)) - })?, - )) - } - (Timestamp(TimeUnit::Microsecond, tz), Time64(TimeUnit::Microsecond)) => { - let tz = tz.as_ref().map(|tz| tz.parse()).transpose()?; - Ok(Arc::new( - array - .as_primitive::() - .try_unary::<_, Time64MicrosecondType, ArrowError>(|x| { - Ok(time_to_time64us(as_time_res_with_timezone::< - TimestampMicrosecondType, - >(x, tz)?)) - })?, - )) + (Timestamp(TimeUnit::Second, _), Time32(_) | Time64(_)) => { + cast_timestamp_to_time::(array, to_type, cast_options) } - (Timestamp(TimeUnit::Microsecond, tz), Time64(TimeUnit::Nanosecond)) => { - let tz = tz.as_ref().map(|tz| tz.parse()).transpose()?; - Ok(Arc::new( - array - .as_primitive::() - .try_unary::<_, Time64NanosecondType, ArrowError>(|x| { - Ok(time_to_time64ns(as_time_res_with_timezone::< - TimestampMicrosecondType, - >(x, tz)?)) - })?, - )) + (Timestamp(TimeUnit::Millisecond, _), Time32(_) | Time64(_)) => { + cast_timestamp_to_time::(array, to_type, cast_options) } - (Timestamp(TimeUnit::Nanosecond, tz), Time64(TimeUnit::Microsecond)) => { - let tz = tz.as_ref().map(|tz| tz.parse()).transpose()?; - Ok(Arc::new( - array - .as_primitive::() - .try_unary::<_, Time64MicrosecondType, ArrowError>(|x| { - Ok(time_to_time64us(as_time_res_with_timezone::< - TimestampNanosecondType, - >(x, tz)?)) - })?, - )) + (Timestamp(TimeUnit::Microsecond, _), Time32(_) | Time64(_)) => { + cast_timestamp_to_time::(array, to_type, cast_options) } - (Timestamp(TimeUnit::Nanosecond, tz), Time64(TimeUnit::Nanosecond)) => { - let tz = tz.as_ref().map(|tz| tz.parse()).transpose()?; - Ok(Arc::new( - array - .as_primitive::() - .try_unary::<_, Time64NanosecondType, ArrowError>(|x| { - Ok(time_to_time64ns(as_time_res_with_timezone::< - TimestampNanosecondType, - >(x, tz)?)) - })?, - )) - } - (Timestamp(TimeUnit::Second, tz), Time32(TimeUnit::Second)) => { - let tz = tz.as_ref().map(|tz| tz.parse()).transpose()?; - Ok(Arc::new( - array - .as_primitive::() - .try_unary::<_, Time32SecondType, ArrowError>(|x| { - Ok(time_to_time32s(as_time_res_with_timezone::< - TimestampSecondType, - >(x, tz)?)) - })?, - )) - } - (Timestamp(TimeUnit::Second, tz), Time32(TimeUnit::Millisecond)) => { - let tz = tz.as_ref().map(|tz| tz.parse()).transpose()?; - Ok(Arc::new( - array - .as_primitive::() - .try_unary::<_, Time32MillisecondType, ArrowError>(|x| { - Ok(time_to_time32ms(as_time_res_with_timezone::< - TimestampSecondType, - >(x, tz)?)) - })?, - )) - } - (Timestamp(TimeUnit::Millisecond, tz), Time32(TimeUnit::Second)) => { - let tz = tz.as_ref().map(|tz| tz.parse()).transpose()?; - Ok(Arc::new( - array - .as_primitive::() - .try_unary::<_, Time32SecondType, ArrowError>(|x| { - Ok(time_to_time32s(as_time_res_with_timezone::< - TimestampMillisecondType, - >(x, tz)?)) - })?, - )) - } - (Timestamp(TimeUnit::Millisecond, tz), Time32(TimeUnit::Millisecond)) => { - let tz = tz.as_ref().map(|tz| tz.parse()).transpose()?; - Ok(Arc::new( - array - .as_primitive::() - .try_unary::<_, Time32MillisecondType, ArrowError>(|x| { - Ok(time_to_time32ms(as_time_res_with_timezone::< - TimestampMillisecondType, - >(x, tz)?)) - })?, - )) - } - (Timestamp(TimeUnit::Microsecond, tz), Time32(TimeUnit::Second)) => { - let tz = tz.as_ref().map(|tz| tz.parse()).transpose()?; - Ok(Arc::new( - array - .as_primitive::() - .try_unary::<_, Time32SecondType, ArrowError>(|x| { - Ok(time_to_time32s(as_time_res_with_timezone::< - TimestampMicrosecondType, - >(x, tz)?)) - })?, - )) - } - (Timestamp(TimeUnit::Microsecond, tz), Time32(TimeUnit::Millisecond)) => { - let tz = tz.as_ref().map(|tz| tz.parse()).transpose()?; - Ok(Arc::new( - array - .as_primitive::() - .try_unary::<_, Time32MillisecondType, ArrowError>(|x| { - Ok(time_to_time32ms(as_time_res_with_timezone::< - TimestampMicrosecondType, - >(x, tz)?)) - })?, - )) - } - (Timestamp(TimeUnit::Nanosecond, tz), Time32(TimeUnit::Second)) => { - let tz = tz.as_ref().map(|tz| tz.parse()).transpose()?; - Ok(Arc::new( - array - .as_primitive::() - .try_unary::<_, Time32SecondType, ArrowError>(|x| { - Ok(time_to_time32s(as_time_res_with_timezone::< - TimestampNanosecondType, - >(x, tz)?)) - })?, - )) - } - (Timestamp(TimeUnit::Nanosecond, tz), Time32(TimeUnit::Millisecond)) => { - let tz = tz.as_ref().map(|tz| tz.parse()).transpose()?; - Ok(Arc::new( - array - .as_primitive::() - .try_unary::<_, Time32MillisecondType, ArrowError>(|x| { - Ok(time_to_time32ms(as_time_res_with_timezone::< - TimestampNanosecondType, - >(x, tz)?)) - })?, - )) + (Timestamp(TimeUnit::Nanosecond, _), Time32(_) | Time64(_)) => { + cast_timestamp_to_time::(array, to_type, cast_options) } (Date64, Timestamp(TimeUnit::Second, _)) => { let array = array @@ -2483,6 +2496,11 @@ where } } +/// Number of `T` units in one day +const fn units_per_day() -> i64 { + SECONDS_IN_DAY * time_unit_multiple(&T::UNIT) +} + /// Get the time unit as a multiple of a second const fn time_unit_multiple(unit: &TimeUnit) -> i64 { match unit { @@ -2528,40 +2546,90 @@ fn cast_numeric_arrays( where FROM: ArrowPrimitiveType, TO: ArrowPrimitiveType, - FROM::Native: NumCast, - TO::Native: NumCast, + FROM::Native: NumCast + NumericNative, + TO::Native: NumCast + NumericNative, { - if cast_options.safe { - // If the value can't be casted to the `TO::Native`, return null - Ok(Arc::new(numeric_cast::( - from.as_primitive::(), - ))) + let from = from.as_primitive::(); + let array: PrimitiveArray = if const { + is_infallible_numeric_cast( + ::KIND, + ::KIND, + ) + } { + // This cast cannot fail, so the fastest kernel, `unary`, can be used. + from.unary(|v| num_cast(v).expect("numeric cast is infallible")) + } else if cast_options.safe { + // If the value can't be cast to the `TO::Native`, return null + from.unary_opt(num_cast) } else { - // If the value can't be casted to the `TO::Native`, return error - Ok(Arc::new(try_numeric_cast::( - from.as_primitive::(), - )?)) - } + // If the value can't be cast to the `TO::Native`, return error + from.try_unary(|v| { + num_cast(v).ok_or_else(|| { + ArrowError::CastError(format!("Can't cast value {v:?} to type {}", TO::DATA_TYPE)) + }) + })? + }; + Ok(Arc::new(array)) } -// Natural cast between numeric types -// If the value of T can't be casted to R, will throw error -fn try_numeric_cast(from: &PrimitiveArray) -> Result, ArrowError> -where - T: ArrowPrimitiveType, - R: ArrowPrimitiveType, - T::Native: NumCast, - R::Native: NumCast, -{ - from.try_unary(|value| { - num_cast::(value).ok_or_else(|| { - ArrowError::CastError(format!( - "Can't cast value {:?} to type {}", - value, - R::DATA_TYPE - )) - }) - }) +/// Properties of a numeric native type that decide whether [`num_cast`] to +/// another numeric type can fail. +#[derive(Clone, Copy)] +struct NumericKind { + float: bool, + signed: bool, + bits: u32, +} + +/// A native type with a [`NumericKind`]. +trait NumericNative { + const KIND: NumericKind; +} + +macro_rules! numeric_native { + ($($t:ty => $float:literal, $signed:literal;)*) => {$( + impl NumericNative for $t { + const KIND: NumericKind = NumericKind { + float: $float, + signed: $signed, + bits: 8 * std::mem::size_of::<$t>() as u32, + }; + } + )*}; +} + +numeric_native! { + u8 => false, false; + u16 => false, false; + u32 => false, false; + u64 => false, false; + i8 => false, true; + i16 => false, true; + i32 => false, true; + i64 => false, true; + half::f16 => true, true; + f32 => true, true; + f64 => true, true; +} + +/// Returns true if [`num_cast`] succeeds for every value of `from` when casting +/// to `to`: every cast to a float, and integer casts whose target holds every +/// source value. +/// +/// This must hold for every bit pattern of the source type, not only for the +/// valid values of a particular array, because [`PrimitiveArray::unary`] applies +/// the conversion to null slots as well, and their contents are arbitrary. +const fn is_infallible_numeric_cast(from: NumericKind, to: NumericKind) -> bool { + if to.float { + true + } else if from.float { + false + } else if from.signed == to.signed { + to.bits >= from.bits + } else { + // Only an unsigned source fits a strictly wider signed target. + !from.signed && to.bits > from.bits + } } /// Natural cast between numeric types @@ -2575,29 +2643,63 @@ where num_traits::cast::cast::(value) } -// Natural cast between numeric types -// If the value of T can't be casted to R, it will be converted to null -fn numeric_cast(from: &PrimitiveArray) -> PrimitiveArray -where - T: ArrowPrimitiveType, - R: ArrowPrimitiveType, - T::Native: NumCast, - R::Native: NumCast, -{ - from.unary_opt::<_, R>(num_cast::) -} - -fn cast_numeric_to_binary( - array: &dyn Array, -) -> Result { - let array = array.as_primitive::(); - let size = std::mem::size_of::(); - let offsets = OffsetBuffer::from_repeated_length(size, array.len()); - Ok(Arc::new(GenericBinaryArray::::try_new( - offsets, - array.values().inner().clone(), - array.nulls().cloned(), - )?)) +/// Returns the offset to use when interpreting `local` as a wall clock reading +/// in `tz`, or `None` if it cannot be resolved. +/// +/// `None` is not expected in practice. With the current timezone database no +/// reading reaches it, because every ambiguous or nonexistent reading resolves +/// as described below. The `None` path is a safeguard against a future +/// timezone database that breaks the assumptions of the gap handling. Callers +/// then apply their usual error or null handling. +/// +/// In an IANA timezone a wall clock reading does not always identify a unique +/// instant, and this function picks one following the same rules as PostgreSQL +/// and DuckDB: +/// +/// * **Ambiguous** -- when the clocks go back ("fall back") the same reading +/// occurs twice. The *later* instant is chosen, i.e. the offset in effect +/// after the transition. For example `2024-11-03T01:30:00` in +/// `America/New_York` is read as `-05:00` (EST), not `-04:00` (EDT). +/// * **Nonexistent** -- when the clocks go forward ("spring forward") the +/// reading never occurs. It is shifted forward by the length of the gap, +/// which is the same as reading it with the offset in effect *before* the +/// transition. For example `2024-03-10T02:30:00` in `America/New_York` is +/// read as `-05:00` (EST) and therefore denotes `2024-03-10T03:30:00-04:00`. +/// +/// Timezones with a fixed offset are never ambiguous and have no gaps. +/// +/// See for the PostgreSQL and +/// ICU (DuckDB) sources these rules are taken from. +fn resolve_local_offset(tz: &Tz, local: &NaiveDateTime) -> Option { + match tz.offset_from_local_datetime(local) { + LocalResult::Single(offset) => Some(offset.fix()), + LocalResult::Ambiguous(_earlier, later) => Some(later.fix()), + LocalResult::None => { + // The reading falls in a gap. Recover the offset in effect before + // the transition by probing 24 hours earlier. + // + // Two separate properties of the timezone database make this sound: + // + // 1. No local gap is longer than 24 hours, so the probe lands + // outside this gap and is itself resolvable. Seven zones have a + // gap of exactly 24 hours -- the dateline changes, such as + // `Pacific/Apia` in 2011 and `Pacific/Kiritimati` in 1994. At the + // last second of one of those the probe lands one second before + // the gap starts, so the true margin here is one second, not a + // comfortable one. + // 2. No two transitions are closer together than 24 hours, so the + // offset the probe finds is the one in effect immediately before + // this transition, and not some older offset. The smallest + // observed interval is 167 hours (`America/Boa_Vista`, 2000). + // + // Property 1 is what makes the probe resolvable; property 2 is what + // makes the answer correct. If the probe is still unresolvable, give + // up and let the caller apply the usual error / null handling. + tz.offset_from_local_datetime(&(*local - TimeDelta::hours(24))) + .earliest() + .map(|offset| offset.fix()) + } + } } fn adjust_timestamp_to_timezone( @@ -2607,8 +2709,8 @@ fn adjust_timestamp_to_timezone( ) -> Result, ArrowError> { let adjust = |o| { let local = as_datetime::(o)?; - let offset = to_tz.offset_from_local_datetime(&local).single()?; - T::from_naive_datetime(local - offset.fix(), None) + let offset = resolve_local_offset(to_tz, &local)?; + T::from_naive_datetime(local - offset, None) }; let adjusted = if cast_options.safe { array.unary_opt::<_, Int64Type>(adjust) @@ -2629,24 +2731,10 @@ fn cast_numeric_to_bool(from: &dyn Array) -> Result where FROM: ArrowPrimitiveType, { - numeric_to_bool_cast::(from.as_primitive::()).map(|to| Arc::new(to) as ArrayRef) -} - -fn numeric_to_bool_cast(from: &PrimitiveArray) -> Result -where - T: ArrowPrimitiveType, -{ - let mut b = BooleanBuilder::with_capacity(from.len()); - - for i in 0..from.len() { - if from.is_null(i) { - b.append_null(); - } else { - b.append_value(cast_num_to_bool::(from.value(i))); - } - } - - Ok(b.finish()) + Ok(Arc::new(BooleanArray::from_unary( + from.as_primitive::(), + cast_num_to_bool, + ))) } /// Cast numeric types to boolean @@ -2667,7 +2755,6 @@ fn cast_bool_to_numeric( ) -> Result where TO: ArrowPrimitiveType, - TO::Native: num_traits::cast::NumCast, { Ok(Arc::new(bool_to_numeric_cast::( from.as_any().downcast_ref::().unwrap(), @@ -2678,192 +2765,38 @@ where fn bool_to_numeric_cast(from: &BooleanArray, _cast_options: &CastOptions) -> PrimitiveArray where T: ArrowPrimitiveType, - T::Native: num_traits::NumCast, -{ - let iter = (0..from.len()).map(|i| { - if from.is_null(i) { - None - } else { - single_bool_to_numeric::(from.value(i)) - } - }); - // Benefit: - // 20% performance improvement - // Soundness: - // The iterator is trustedLen because it comes from a Range - unsafe { PrimitiveArray::::from_trusted_len_iter(iter) } -} - -/// Cast single bool value to numeric value. -#[inline] -pub fn single_bool_to_numeric(value: bool) -> Option -where - O: num_traits::NumCast + Default, { - if value { - // a workaround to cast a primitive to type O, infallible - num_traits::cast::cast(1) - } else { - Some(O::default()) - } -} - -/// Helper function to cast from one `BinaryArray` or 'LargeBinaryArray' to 'FixedSizeBinaryArray'. -fn cast_binary_to_fixed_size_binary( - array: &dyn Array, - byte_width: i32, - cast_options: &CastOptions, -) -> Result { - let array = array.as_binary::(); - let mut builder = FixedSizeBinaryBuilder::with_capacity(array.len(), byte_width); - - for i in 0..array.len() { - if array.is_null(i) { - builder.append_null(); - } else { - match builder.append_value(array.value(i)) { - Ok(()) => {} - Err(e) => match cast_options.safe { - true => builder.append_null(), - false => return Err(e), - }, - } - } - } - - Ok(Arc::new(builder.finish())) -} - -/// Helper function to cast from 'FixedSizeBinaryArray' to one `BinaryArray` or 'LargeBinaryArray'. -/// If the target one is too large for the source array it will return an Error. -fn cast_fixed_size_binary_to_binary( - array: &dyn Array, - byte_width: i32, -) -> Result { - let array = array - .as_any() - .downcast_ref::() - .unwrap(); - - let offsets: i128 = byte_width as i128 * array.len() as i128; - - let is_binary = matches!(GenericBinaryType::::DATA_TYPE, DataType::Binary); - if is_binary && offsets > i32::MAX as i128 { - return Err(ArrowError::ComputeError( - "FixedSizeBinary array too large to cast to Binary array".to_string(), - )); - } else if !is_binary && offsets > i64::MAX as i128 { - return Err(ArrowError::ComputeError( - "FixedSizeBinary array too large to cast to LargeBinary array".to_string(), - )); - } - - let mut builder = GenericBinaryBuilder::::with_capacity(array.len(), array.len()); - - for i in 0..array.len() { - if array.is_null(i) { - builder.append_null(); - } else { - builder.append_value(array.value(i)); - } - } - - Ok(Arc::new(builder.finish())) -} - -fn cast_fixed_size_binary_to_binary_view( - array: &dyn Array, - _byte_width: i32, -) -> Result { - let array = array - .as_any() - .downcast_ref::() - .unwrap(); - - let mut builder = BinaryViewBuilder::with_capacity(array.len()); - for i in 0..array.len() { - if array.is_null(i) { - builder.append_null(); - } else { - builder.append_value(array.value(i)); - } - } - - Ok(Arc::new(builder.finish())) -} - -/// Helper function to cast from one `ByteArrayType` to another and vice versa. -/// If the target one (e.g., `LargeUtf8`) is too large for the source array it will return an Error. -fn cast_byte_container(array: &dyn Array) -> Result -where - FROM: ByteArrayType, - TO: ByteArrayType, - FROM::Offset: OffsetSizeTrait + ToPrimitive, - TO::Offset: OffsetSizeTrait + NumCast, -{ - let data = array.to_data(); - assert_eq!(data.data_type(), &FROM::DATA_TYPE); - let str_values_buf = data.buffers()[1].clone(); - let offsets = data.buffers()[0].typed_data::(); - - let mut cast_offsets = Vec::::with_capacity(offsets.len()); - offsets - .iter() - .try_for_each::<_, Result<_, ArrowError>>(|offset| { - let offset = - <::Offset as NumCast>::from(*offset).ok_or_else(|| { - ArrowError::ComputeError(format!( - "{}{} array too large to cast to {}{} array", - FROM::Offset::PREFIX, - FROM::PREFIX, - TO::Offset::PREFIX, - TO::PREFIX - )) - })?; - cast_offsets.push(offset); - Ok(()) - })?; - - let offset_buffer = Buffer::from_vec(cast_offsets); - - let dtype = TO::DATA_TYPE; - - let builder = ArrayData::builder(dtype) - .offset(array.offset()) - .len(array.len()) - .add_buffer(offset_buffer) - .add_buffer(str_values_buf) - .nulls(data.nulls().cloned()); - - let array_data = unsafe { builder.build_unchecked() }; - - Ok(Arc::new(GenericByteArray::::from(array_data))) + let to_numeric = |bits: u64, i: usize| { + if bits & (1 << i) != 0 { + T::Native::ONE + } else { + T::Native::ZERO + } + }; + // Unpack a 64-bit word at a time: the fixed-size inner loop helps the + // compiler vectorize the conversion. + let chunks = from.values().bit_chunks(); + let mut values = Vec::with_capacity(from.len()); + for bits in &chunks { + values.extend((0..64).map(|i| to_numeric(bits, i))); + } + let bits = chunks.remainder_bits(); + values.extend((0..chunks.remainder_len()).map(|i| to_numeric(bits, i))); + PrimitiveArray::new(values.into(), from.nulls().cloned()) } -/// Helper function to cast from one `ByteViewType` array to `ByteArrayType` array. -fn cast_view_to_byte(array: &dyn Array) -> Result +/// Cast single bool value to numeric value. +#[inline] +pub fn single_bool_to_numeric(value: bool) -> Option where - FROM: ByteViewType, - TO: ByteArrayType, - FROM::Native: AsRef, + O: num_traits::NumCast + Default, { - let data = array.to_data(); - let view_array = GenericByteViewArray::::from(data); - - let len = view_array.len(); - let bytes = view_array - .views() - .iter() - .map(|v| ByteView::from(*v).length as usize) - .sum::(); - - let mut byte_array_builder = GenericByteBuilder::::with_capacity(len, bytes); - - for val in &view_array { - byte_array_builder.append_option(val); + if value { + // a workaround to cast a primitive to type O, infallible + num_traits::cast::cast(1) + } else { + Some(O::default()) } - - Ok(Arc::new(byte_array_builder.finish())) } #[cfg(test)] @@ -2872,10 +2805,10 @@ mod tests { use crate::parse::parse_decimal; use DataType::*; use arrow_array::{Int64Array, RunArray, StringArray}; - use arrow_buffer::{Buffer, IntervalDayTime, NullBuffer}; + use arrow_buffer::{BooleanBuffer, Buffer, IntervalDayTime, NullBuffer}; use arrow_buffer::{ScalarBuffer, i256}; use arrow_schema::{DataType, Field}; - use chrono::NaiveDate; + use chrono::{NaiveDate, NaiveTime}; use half::f16; use std::sync::Arc; @@ -4686,6 +4619,52 @@ mod tests { ); } + /// Checks the fast-path classification against `num_cast` for every pair of + /// numeric types. `num_cast` fails only on values outside the target's range, + /// so each type's extremes decide, plus NaN for floats. + #[test] + fn test_is_infallible_numeric_cast() { + macro_rules! check { + ($from:ty => $values:expr, [$($to:ty),*]) => {$( + let infallible = is_infallible_numeric_cast( + <<$from as ArrowPrimitiveType>::Native as NumericNative>::KIND, + <<$to as ArrowPrimitiveType>::Native as NumericNative>::KIND, + ); + let succeeds = $values + .iter() + .all(|&v| num_cast::<_, <$to as ArrowPrimitiveType>::Native>(v).is_some()); + assert_eq!( + infallible, + succeeds, + "{} to {}", + <$from as ArrowPrimitiveType>::DATA_TYPE, + <$to as ArrowPrimitiveType>::DATA_TYPE + ); + )*}; + } + macro_rules! check_all { + ($($from:ty => $values:expr),* $(,)?) => {$( + check!($from => $values, [ + UInt8Type, UInt16Type, UInt32Type, UInt64Type, Int8Type, Int16Type, + Int32Type, Int64Type, Float16Type, Float32Type, Float64Type + ]); + )*}; + } + check_all! { + UInt8Type => [u8::MIN, u8::MAX], + UInt16Type => [u16::MIN, u16::MAX], + UInt32Type => [u32::MIN, u32::MAX], + UInt64Type => [u64::MIN, u64::MAX], + Int8Type => [i8::MIN, i8::MAX], + Int16Type => [i16::MIN, i16::MAX], + Int32Type => [i32::MIN, i32::MAX], + Int64Type => [i64::MIN, i64::MAX], + Float16Type => [f16::MIN, f16::MAX, f16::NAN], + Float32Type => [f32::MIN, f32::MAX, f32::NAN], + Float64Type => [f64::MIN, f64::MAX, f64::NAN], + } + } + #[test] fn test_cast_i32_to_f64() { let array = Int32Array::from(vec![5, 6, 7, 8, 9]); @@ -5014,33 +4993,85 @@ mod tests { } #[test] - fn test_cast_bool_to_utf8view() { - let array = BooleanArray::from(vec![Some(true), Some(false), None]); - let b = cast(&array, &DataType::Utf8View).unwrap(); - let c = b.as_any().downcast_ref::().unwrap(); - assert_eq!("true", c.value(0)); - assert_eq!("false", c.value(1)); - assert!(!c.is_valid(2)); + fn test_cast_bool_numeric_sliced() { + for nulls in [ + None, + Some(NullBuffer::from( + [false, true, true, false, true, true].repeat(12), + )), + Some(NullBuffer::new_null(72)), + ] { + // Include true values under nulls and slices crossing a bitmap word boundary. + let booleans = BooleanArray::new( + BooleanBuffer::from([true, false, true, false, true, false].repeat(12)), + nulls.clone(), + ); + let numbers = Int32Array::new([1, 0, 1, 0, 1, 0].repeat(12).into(), nulls); + for data_type in [ + Int8, Int16, Int32, Int64, UInt8, UInt16, UInt32, UInt64, Float16, Float32, Float64, + ] { + let expected = cast(&numbers, &data_type).unwrap(); + for offset in [0, 3, 63, 72] { + let input = booleans.slice(offset, 72 - offset); + let expected = expected.slice(offset, 72 - offset); + let actual = cast(&input, &data_type).unwrap(); + assert_eq!(actual.as_ref(), expected.as_ref()); + + let actual = cast(expected.as_ref(), &Boolean).unwrap(); + assert_eq!(actual.as_ref(), &input); + } + } + } } #[test] - fn test_cast_bool_to_utf8() { - let array = BooleanArray::from(vec![Some(true), Some(false), None]); - let b = cast(&array, &DataType::Utf8).unwrap(); - let c = b.as_any().downcast_ref::().unwrap(); - assert_eq!("true", c.value(0)); - assert_eq!("false", c.value(1)); - assert!(!c.is_valid(2)); + fn test_cast_float_to_bool_special_values() { + let numbers = Float64Array::from(vec![ + 0.0, + -0.0, + 1.5, + -1.5, + f64::NAN, + f64::INFINITY, + f64::NEG_INFINITY, + ]); + let expected = BooleanArray::from(vec![false, false, true, true, true, true, true]); + for data_type in [Float16, Float32, Float64] { + let input = cast(&numbers, &data_type).unwrap(); + let actual = cast(input.as_ref(), &Boolean).unwrap(); + assert_eq!(actual.as_ref(), &expected); + } } #[test] - fn test_cast_bool_to_large_utf8() { - let array = BooleanArray::from(vec![Some(true), Some(false), None]); - let b = cast(&array, &DataType::LargeUtf8).unwrap(); - let c = b.as_any().downcast_ref::().unwrap(); - assert_eq!("true", c.value(0)); - assert_eq!("false", c.value(1)); - assert!(!c.is_valid(2)); + fn test_cast_bool_to_string() { + for nulls in [ + None, + Some(NullBuffer::from([false, false, true, true].repeat(18))), + Some(NullBuffer::new_null(72)), + ] { + // Include both Boolean values under nulls, an unaligned slice, and an empty slice. + let array = BooleanArray::new(BooleanBuffer::from([true, false].repeat(36)), nulls); + let expected: Vec<_> = ["true", "false"] + .repeat(36) + .into_iter() + .enumerate() + .map(|(i, value)| array.is_valid(i).then_some(value)) + .collect(); + for (offset, len) in [(0, 72), (3, 65), (72, 0)] { + let input = array.slice(offset, len); + for data_type in [Utf8, LargeUtf8, Utf8View] { + let actual = cast(&input, &data_type).unwrap(); + let values: Vec<_> = match data_type { + Utf8 => actual.as_string::().iter().collect(), + LargeUtf8 => actual.as_string::().iter().collect(), + Utf8View => actual.as_string_view().iter().collect(), + _ => unreachable!(), + }; + assert_eq!(values, expected[offset..offset + len]); + } + } + } } #[test] @@ -5437,6 +5468,17 @@ mod tests { assert_eq!(result.values(), &[247112596800]); } + #[test] + fn test_cast_string_to_date32_with_surrounding_whitespace() { + let array = StringArray::from(vec![" 2026-09-29 ", "2026-09-29"]); + + let result = cast(&array, &DataType::Date32).unwrap(); + let result = result.as_primitive::(); + + assert!(!result.is_null(0)); + assert_eq!(result.value(0), result.value(1)); + } + #[test] fn test_cast_string_to_date32() { let a0 = Arc::new(StringViewArray::from(vec![ @@ -6633,6 +6675,185 @@ mod tests { assert_eq!(17890, c.value(1)); } + /// Compares arithmetic casts with Chrono for representable dates, including + /// negative and fractional boundaries, nulls, and nonzero-offset slices. + fn check_timestamp_casts() { + use rand::rngs::StdRng; + use rand::{RngExt, SeedableRng}; + + let units = time_unit_multiple(&T::UNIT); + let around = |x: i64| [x - 1, x, x + 1, -x - 1, -x, -x + 1]; + let mut values = vec![i64::MIN, i64::MAX]; + for base in [0, 60, 3_600, 86_400] { + for seconds in around(base) { + for fraction in [-1, 0, 1, units - 1] { + values.push(seconds * units + fraction); + } + } + } + for boundary in [1_000, 1_000_000] { + values.extend(around(boundary)); + } + let mut rng = StdRng::seed_from_u64(42); + for _ in 0..256 { + values.push(rng.random()); + let seconds = rng.random_range(-2_000_000_000..2_000_000_000); + values.push(seconds * units + rng.random_range(0..units)); + } + values.retain(|&v| as_datetime::(v).is_some()); + let no_nulls = PrimitiveArray::::new(values.clone().into(), None); + let mixed = PrimitiveArray::::from_iter( + values + .iter() + .enumerate() + .map(|(i, &v)| (i % 5 != 0).then_some(v)), + ); + for array in [no_nulls, mixed.slice(1, mixed.len() - 2)] { + let dates: Vec<_> = array + .iter() + .map(|v| v.map(|v| as_datetime::(v).unwrap())) + .collect(); + let expected: [ArrayRef; 5] = [ + Arc::new(Date32Array::from_iter( + dates + .iter() + .map(|d| d.map(|d| Date32Type::from_naive_date(d.date()))), + )), + Arc::new(Time32SecondArray::from_iter( + dates.iter().map(|d| d.map(|d| time_to_time32s(d.time()))), + )), + Arc::new(Time32MillisecondArray::from_iter( + dates.iter().map(|d| d.map(|d| time_to_time32ms(d.time()))), + )), + Arc::new(Time64MicrosecondArray::from_iter( + dates.iter().map(|d| d.map(|d| time_to_time64us(d.time()))), + )), + Arc::new(Time64NanosecondArray::from_iter( + dates.iter().map(|d| d.map(|d| time_to_time64ns(d.time()))), + )), + ]; + for expected in expected { + let actual = cast(&array, expected.data_type()).unwrap(); + assert_eq!(actual.as_ref(), expected.as_ref(), "{}", T::DATA_TYPE); + } + } + } + + #[test] + fn test_cast_timestamp_without_timezone_matches_chrono() { + check_timestamp_casts::(); + check_timestamp_casts::(); + check_timestamp_casts::(); + check_timestamp_casts::(); + } + + #[test] + fn test_cast_timestamp_to_time_extremes() { + fn check(min: (u32, u32, u32, u32), max: (u32, u32, u32, u32)) { + let time = |(h, m, s, ns)| NaiveTime::from_hms_nano_opt(h, m, s, ns).unwrap(); + let times = [Some(time(min)), None, Some(time(max))]; + let array = PrimitiveArray::::from_iter([Some(i64::MIN), None, Some(i64::MAX)]); + let expected: [ArrayRef; 4] = [ + Arc::new(Time32SecondArray::from_iter( + times.map(|v| v.map(time_to_time32s)), + )), + Arc::new(Time32MillisecondArray::from_iter( + times.map(|v| v.map(time_to_time32ms)), + )), + Arc::new(Time64MicrosecondArray::from_iter( + times.map(|v| v.map(time_to_time64us)), + )), + Arc::new(Time64NanosecondArray::from_iter( + times.map(|v| v.map(time_to_time64ns)), + )), + ]; + for expected in expected { + assert_eq!( + cast(&array, expected.data_type()).unwrap().as_ref(), + expected.as_ref() + ); + } + } + check::((8, 29, 52, 0), (15, 30, 7, 0)); + check::((16, 47, 4, 192_000_000), (7, 12, 55, 807_000_000)); + check::((19, 59, 5, 224_192_000), (4, 0, 54, 775_807_000)); + // Nanosecond extremes are representable in Chrono and covered above. + } + + #[test] + fn test_cast_timestamp_to_date32_range() { + fn check() { + let day = SECONDS_IN_DAY * time_unit_multiple(&T::UNIT); + let first = i32::MIN as i64 * day; + let last = (i32::MAX as i64 + 1) * day - 1; + let array = PrimitiveArray::::new( + vec![first, i64::MAX, last].into(), + Some(NullBuffer::from(vec![true, false, true])), + ); + let expected = Date32Array::from(vec![Some(i32::MIN), None, Some(i32::MAX)]); + for safe in [true, false] { + let options = CastOptions { + safe, + ..Default::default() + }; + assert_eq!( + cast_with_options(&array, &Date32, &options) + .unwrap() + .as_ref(), + &expected + ); + } + // Day counts that do not fit in i32 become null when safe, otherwise an error. + let array = PrimitiveArray::::from_iter_values([ + first - 1, + last, + last + 1, + i64::MIN, + i64::MAX, + ]); + let expected = Date32Array::from(vec![None, Some(i32::MAX), None, None, None]); + assert_eq!(cast(&array, &Date32).unwrap().as_ref(), &expected); + let options = CastOptions { + safe: false, + ..Default::default() + }; + assert!(cast_with_options(&array, &Date32, &options).is_err()); + } + check::(); + check::(); + // The entire microsecond domain fits in Date32, including dates outside Chrono's range. + let array = TimestampMicrosecondArray::from(vec![Some(i64::MIN), None, Some(i64::MAX)]); + assert_eq!( + cast(&array, &Date32).unwrap().as_ref(), + &Date32Array::from(vec![Some(-106_751_992), None, Some(106_751_991)]), + ); + } + + #[test] + fn test_cast_timestamp_date_time_timezone_validation() { + let invalid = TimestampSecondArray::from(vec![0]).with_timezone("invalid timezone"); + let extreme = TimestampSecondArray::from(vec![0, i64::MAX]).with_timezone("+00:00"); + let options = CastOptions { + safe: false, + ..Default::default() + }; + for target in [ + Date32, + Time32(TimeUnit::Second), + Time32(TimeUnit::Millisecond), + Time64(TimeUnit::Microsecond), + Time64(TimeUnit::Nanosecond), + ] { + // A timezone that fails to parse is an error even when the cast is safe. + assert!(cast(&invalid, &target).is_err()); + // A timestamp outside Chrono's date range is null when safe, otherwise an error. + let actual = cast(&extreme, &target).unwrap(); + assert!(actual.is_valid(0)); + assert!(actual.is_null(1)); + assert!(cast_with_options(&extreme, &target, &options).is_err()); + } + } + #[test] fn test_cast_timestamp_to_date32() { let array = @@ -6740,29 +6961,39 @@ mod tests { assert_eq!(3601000000000, c.value(1)); assert!(c.is_null(2)); - // test timestamp nanoseconds - let a = TimestampNanosecondArray::from(vec![Some(86405000000000), Some(1000000000), None]) + // Test fractional timestamps on both sides of the epoch. + let a = TimestampNanosecondArray::from(vec![Some(86_405_123_456_789), Some(-1), None]) .with_timezone("+01:00".to_string()); let array = Arc::new(a) as ArrayRef; let b = cast(&array, &DataType::Time64(TimeUnit::Microsecond)).unwrap(); let c = b.as_primitive::(); - assert_eq!(3605000000, c.value(0)); - assert_eq!(3601000000, c.value(1)); + assert_eq!(3_605_123_456, c.value(0)); + assert_eq!(3_599_999_999, c.value(1)); assert!(c.is_null(2)); let b = cast(&array, &DataType::Time64(TimeUnit::Nanosecond)).unwrap(); let c = b.as_primitive::(); - assert_eq!(3605000000000, c.value(0)); - assert_eq!(3601000000000, c.value(1)); + assert_eq!(3_605_123_456_789, c.value(0)); + assert_eq!(3_599_999_999_999, c.value(1)); assert!(c.is_null(2)); - // test overflow + // test overflow, safe cast let a = TimestampSecondArray::from(vec![Some(i64::MAX)]).with_timezone("+01:00".to_string()); let array = Arc::new(a) as ArrayRef; - let b = cast(&array, &DataType::Time64(TimeUnit::Microsecond)); + let b = cast(&array, &DataType::Time64(TimeUnit::Microsecond)).unwrap(); + assert!(b.is_null(0)); + let b = cast(&array, &DataType::Time64(TimeUnit::Nanosecond)).unwrap(); + assert!(b.is_null(0)); + // test overflow, unsafe cast + let options = CastOptions { + safe: false, + ..Default::default() + }; + let b = cast_with_options(&array, &DataType::Time64(TimeUnit::Microsecond), &options); assert!(b.is_err()); - let b = cast(&array, &DataType::Time64(TimeUnit::Nanosecond)); + let b = cast_with_options(&array, &DataType::Time64(TimeUnit::Nanosecond), &options); assert!(b.is_err()); + // Time64 only supports microseconds and nanoseconds let b = cast(&array, &DataType::Time64(TimeUnit::Millisecond)); assert!(b.is_err()); } @@ -6814,28 +7045,37 @@ mod tests { assert_eq!(3601000, c.value(1)); assert!(c.is_null(2)); - // test timestamp nanoseconds - let a = TimestampNanosecondArray::from(vec![Some(86405000000000), Some(1000000000), None]) + // Test fractional timestamps on both sides of the epoch. + let a = TimestampNanosecondArray::from(vec![Some(86_405_123_456_789), Some(-1), None]) .with_timezone("+01:00".to_string()); let array = Arc::new(a) as ArrayRef; let b = cast(&array, &DataType::Time32(TimeUnit::Second)).unwrap(); let c = b.as_primitive::(); assert_eq!(3605, c.value(0)); - assert_eq!(3601, c.value(1)); + assert_eq!(3599, c.value(1)); assert!(c.is_null(2)); let b = cast(&array, &DataType::Time32(TimeUnit::Millisecond)).unwrap(); let c = b.as_primitive::(); - assert_eq!(3605000, c.value(0)); - assert_eq!(3601000, c.value(1)); + assert_eq!(3_605_123, c.value(0)); + assert_eq!(3_599_999, c.value(1)); assert!(c.is_null(2)); - // test overflow + // test overflow, safe cast let a = TimestampSecondArray::from(vec![Some(i64::MAX)]).with_timezone("+01:00".to_string()); let array = Arc::new(a) as ArrayRef; - let b = cast(&array, &DataType::Time32(TimeUnit::Second)); + let b = cast(&array, &DataType::Time32(TimeUnit::Second)).unwrap(); + assert!(b.is_null(0)); + let b = cast(&array, &DataType::Time32(TimeUnit::Millisecond)).unwrap(); + assert!(b.is_null(0)); + // test overflow, unsafe cast + let options = CastOptions { + safe: false, + ..Default::default() + }; + let b = cast_with_options(&array, &DataType::Time32(TimeUnit::Second), &options); assert!(b.is_err()); - let b = cast(&array, &DataType::Time32(TimeUnit::Millisecond)); + let b = cast_with_options(&array, &DataType::Time32(TimeUnit::Millisecond), &options); assert!(b.is_err()); } @@ -6850,70 +7090,255 @@ mod tests { let to_type = DataType::Timestamp(TimeUnit::Nanosecond, None); let timestamp_array = cast(&string_array, &to_type).unwrap(); - let to_type = DataType::Timestamp(TimeUnit::Microsecond, Some("+0700".into())); - let timestamp_array = cast(×tamp_array, &to_type).unwrap(); + let to_type = DataType::Timestamp(TimeUnit::Microsecond, Some("+0700".into())); + let timestamp_array = cast(×tamp_array, &to_type).unwrap(); + + let string_array = cast(×tamp_array, &DataType::Utf8).unwrap(); + let result = string_array.as_string::(); + assert_eq!("2000-01-01T00:00:00.123456+07:00", result.value(0)); + assert_eq!("2010-01-01T00:00:00.123456+07:00", result.value(1)); + assert!(result.is_null(2)); + } + + // Cast Timestamp(_, Some(timezone)) -> Timestamp(_, None) + #[test] + fn test_cast_timestamp_with_timezone_2() { + let string_array: Arc = Arc::new(StringArray::from(vec![ + Some("2000-01-01T07:00:00.123456789"), + Some("2010-01-01T07:00:00.123456789"), + None, + ])); + let to_type = DataType::Timestamp(TimeUnit::Millisecond, Some("+0700".into())); + let timestamp_array = cast(&string_array, &to_type).unwrap(); + + // Check intermediate representation is correct + let string_array = cast(×tamp_array, &DataType::Utf8).unwrap(); + let result = string_array.as_string::(); + assert_eq!("2000-01-01T07:00:00.123+07:00", result.value(0)); + assert_eq!("2010-01-01T07:00:00.123+07:00", result.value(1)); + assert!(result.is_null(2)); + + let to_type = DataType::Timestamp(TimeUnit::Nanosecond, None); + let timestamp_array = cast(×tamp_array, &to_type).unwrap(); + + let string_array = cast(×tamp_array, &DataType::Utf8).unwrap(); + let result = string_array.as_string::(); + assert_eq!("2000-01-01T00:00:00.123", result.value(0)); + assert_eq!("2010-01-01T00:00:00.123", result.value(1)); + assert!(result.is_null(2)); + } + + // Cast Timestamp(_, Some(timezone)) -> Timestamp(_, Some(timezone)) + #[test] + fn test_cast_timestamp_with_timezone_3() { + let string_array: Arc = Arc::new(StringArray::from(vec![ + Some("2000-01-01T07:00:00.123456789"), + Some("2010-01-01T07:00:00.123456789"), + None, + ])); + let to_type = DataType::Timestamp(TimeUnit::Microsecond, Some("+0700".into())); + let timestamp_array = cast(&string_array, &to_type).unwrap(); + + // Check intermediate representation is correct + let string_array = cast(×tamp_array, &DataType::Utf8).unwrap(); + let result = string_array.as_string::(); + assert_eq!("2000-01-01T07:00:00.123456+07:00", result.value(0)); + assert_eq!("2010-01-01T07:00:00.123456+07:00", result.value(1)); + assert!(result.is_null(2)); + + let to_type = DataType::Timestamp(TimeUnit::Second, Some("-08:00".into())); + let timestamp_array = cast(×tamp_array, &to_type).unwrap(); + + let string_array = cast(×tamp_array, &DataType::Utf8).unwrap(); + let result = string_array.as_string::(); + assert_eq!("1999-12-31T16:00:00-08:00", result.value(0)); + assert_eq!("2009-12-31T16:00:00-08:00", result.value(1)); + assert!(result.is_null(2)); + } + + /// The i64 value of a `Timestamp(Second, None)` holding the given wall clock + /// reading (a naive timestamp is stored as if it were UTC). + fn naive_seconds(y: i32, m: u32, d: u32, h: u32, min: u32) -> i64 { + NaiveDate::from_ymd_opt(y, m, d) + .unwrap() + .and_hms_opt(h, min, 0) + .unwrap() + .and_utc() + .timestamp() + } + + /// The instant, in seconds since the epoch, denoted by the given wall clock + /// reading at a fixed offset of `offset_hours`. + fn instant_seconds(y: i32, m: u32, d: u32, h: u32, min: u32, offset_hours: i32) -> i64 { + let offset = FixedOffset::east_opt(offset_hours * 3600).unwrap(); + NaiveDate::from_ymd_opt(y, m, d) + .unwrap() + .and_hms_opt(h, min, 0) + .unwrap() + .and_local_timezone(offset) + .unwrap() + .timestamp() + } + + // Cast Timestamp(_, None) -> Timestamp(_, Some(IANA timezone)) across DST + // transitions. See `resolve_local_offset`. + #[test] + fn test_cast_timestamp_to_named_timezone_dst() { + // Unambiguous, EDT (-04:00) is in effect. + let unambiguous = naive_seconds(2024, 11, 1, 0, 0); + // Ambiguous: the clocks go back at 2024-11-03T02:00 EDT, so 01:30 + // happens twice, first at -04:00 and then at -05:00. + let ambiguous = naive_seconds(2024, 11, 3, 1, 30); + // Nonexistent: the clocks go forward at 2024-03-10T02:00 EST, so 02:30 + // never happens. + let nonexistent = naive_seconds(2024, 3, 10, 2, 30); + assert_eq!( + [unambiguous, ambiguous, nonexistent], + [1_730_419_200, 1_730_597_400, 1_710_037_800] + ); + + let array = TimestampSecondArray::from(vec![ + Some(unambiguous), + Some(ambiguous), + Some(nonexistent), + None, + ]); + let to_type = DataType::Timestamp(TimeUnit::Second, Some("America/New_York".into())); + let options = CastOptions { + safe: false, + ..Default::default() + }; + + let b = cast_with_options(&array, &to_type, &options).unwrap(); + assert_eq!(b.data_type(), &to_type); + let c = b.as_primitive::(); + + assert_eq!(c.value(0), instant_seconds(2024, 11, 1, 0, 0, -4)); + assert_eq!(c.value(0), 1_730_433_600); + // The later of the two candidates, i.e. EST rather than EDT. + assert_eq!(c.value(1), instant_seconds(2024, 11, 3, 1, 30, -5)); + assert_eq!(c.value(1), 1_730_615_400); + // Shifted forward by the one hour gap: 02:30 EST is 03:30 EDT. + assert_eq!(c.value(2), instant_seconds(2024, 3, 10, 2, 30, -5)); + assert_eq!(c.value(2), instant_seconds(2024, 3, 10, 3, 30, -4)); + assert_eq!(c.value(2), 1_710_055_800); + assert!(c.is_null(3)); + } + + // The same values must not become null when `safe` casting is requested. + #[test] + fn test_cast_timestamp_to_named_timezone_dst_safe() { + let array = TimestampSecondArray::from(vec![ + Some(naive_seconds(2024, 11, 1, 0, 0)), + Some(naive_seconds(2024, 11, 3, 1, 30)), + Some(naive_seconds(2024, 3, 10, 2, 30)), + None, + ]); + let to_type = DataType::Timestamp(TimeUnit::Second, Some("America/New_York".into())); + let options = CastOptions { + safe: true, + ..Default::default() + }; + + let b = cast_with_options(&array, &to_type, &options).unwrap(); + assert_eq!(b.data_type(), &to_type); + let c = b.as_primitive::(); + assert_eq!(c.null_count(), 1); + assert_eq!(c.value(0), 1_730_433_600); + assert_eq!(c.value(1), 1_730_615_400); + assert_eq!(c.value(2), 1_710_055_800); + assert!(c.is_null(3)); + } + + // Southern hemisphere: the transitions run the other way around. + #[test] + fn test_cast_timestamp_to_named_timezone_dst_southern_hemisphere() { + // Ambiguous: clocks go back at 2024-04-07T03:00 AEDT (+11:00 -> +10:00). + let ambiguous = naive_seconds(2024, 4, 7, 2, 30); + // Nonexistent: clocks go forward at 2024-10-06T02:00 AEST (+10:00 -> +11:00). + let nonexistent = naive_seconds(2024, 10, 6, 2, 30); + + let array = TimestampSecondArray::from(vec![Some(ambiguous), Some(nonexistent)]); + let to_type = DataType::Timestamp(TimeUnit::Second, Some("Australia/Sydney".into())); + let options = CastOptions { + safe: false, + ..Default::default() + }; - let string_array = cast(×tamp_array, &DataType::Utf8).unwrap(); - let result = string_array.as_string::(); - assert_eq!("2000-01-01T00:00:00.123456+07:00", result.value(0)); - assert_eq!("2010-01-01T00:00:00.123456+07:00", result.value(1)); - assert!(result.is_null(2)); + let b = cast_with_options(&array, &to_type, &options).unwrap(); + assert_eq!(b.data_type(), &to_type); + let c = b.as_primitive::(); + + // The later of the two candidates, i.e. AEST (+10:00). + assert_eq!(c.value(0), instant_seconds(2024, 4, 7, 2, 30, 10)); + assert_eq!(c.value(0), 1_712_421_000); + // Shifted forward by the one hour gap: 02:30 AEST is 03:30 AEDT. + assert_eq!(c.value(1), instant_seconds(2024, 10, 6, 2, 30, 10)); + assert_eq!(c.value(1), instant_seconds(2024, 10, 6, 3, 30, 11)); + assert_eq!(c.value(1), 1_728_145_800); } - // Cast Timestamp(_, Some(timezone)) -> Timestamp(_, None) + // The resolution is independent of the time unit. #[test] - fn test_cast_timestamp_with_timezone_2() { - let string_array: Arc = Arc::new(StringArray::from(vec![ - Some("2000-01-01T07:00:00.123456789"), - Some("2010-01-01T07:00:00.123456789"), - None, - ])); - let to_type = DataType::Timestamp(TimeUnit::Millisecond, Some("+0700".into())); - let timestamp_array = cast(&string_array, &to_type).unwrap(); + fn test_cast_timestamp_to_named_timezone_dst_nanosecond() { + let ambiguous = naive_seconds(2024, 11, 3, 1, 30) * 1_000_000_000 + 123_456_789; + let nonexistent = naive_seconds(2024, 3, 10, 2, 30) * 1_000_000_000 + 123_456_789; + let array = TimestampNanosecondArray::from(vec![Some(ambiguous), Some(nonexistent)]); + let to_type = DataType::Timestamp(TimeUnit::Nanosecond, Some("America/New_York".into())); + let options = CastOptions { + safe: false, + ..Default::default() + }; - // Check intermediate representation is correct - let string_array = cast(×tamp_array, &DataType::Utf8).unwrap(); - let result = string_array.as_string::(); - assert_eq!("2000-01-01T07:00:00.123+07:00", result.value(0)); - assert_eq!("2010-01-01T07:00:00.123+07:00", result.value(1)); - assert!(result.is_null(2)); + let b = cast_with_options(&array, &to_type, &options).unwrap(); + assert_eq!(b.data_type(), &to_type); + let c = b.as_primitive::(); + assert_eq!(c.value(0), 1_730_615_400 * 1_000_000_000 + 123_456_789); + assert_eq!(c.value(1), 1_710_055_800 * 1_000_000_000 + 123_456_789); + } - let to_type = DataType::Timestamp(TimeUnit::Nanosecond, None); - let timestamp_array = cast(×tamp_array, &to_type).unwrap(); + // Unit conversion still composes with the timezone adjustment. + #[test] + fn test_cast_timestamp_to_named_timezone_dst_changing_unit() { + let array = TimestampSecondArray::from(vec![ + Some(naive_seconds(2024, 11, 3, 1, 30)), + Some(naive_seconds(2024, 3, 10, 2, 30)), + ]); + let to_type = DataType::Timestamp(TimeUnit::Millisecond, Some("America/New_York".into())); + let options = CastOptions { + safe: false, + ..Default::default() + }; - let string_array = cast(×tamp_array, &DataType::Utf8).unwrap(); - let result = string_array.as_string::(); - assert_eq!("2000-01-01T00:00:00.123", result.value(0)); - assert_eq!("2010-01-01T00:00:00.123", result.value(1)); - assert!(result.is_null(2)); + let b = cast_with_options(&array, &to_type, &options).unwrap(); + assert_eq!(b.data_type(), &to_type); + let c = b.as_primitive::(); + assert_eq!(c.value(0), 1_730_615_400_000); + assert_eq!(c.value(1), 1_710_055_800_000); } - // Cast Timestamp(_, Some(timezone)) -> Timestamp(_, Some(timezone)) + // A fixed offset has no transitions, so the same readings are unaffected. #[test] - fn test_cast_timestamp_with_timezone_3() { - let string_array: Arc = Arc::new(StringArray::from(vec![ - Some("2000-01-01T07:00:00.123456789"), - Some("2010-01-01T07:00:00.123456789"), + fn test_cast_timestamp_to_fixed_offset_timezone_unaffected() { + let array = TimestampSecondArray::from(vec![ + Some(naive_seconds(2024, 11, 1, 0, 0)), + Some(naive_seconds(2024, 11, 3, 1, 30)), + Some(naive_seconds(2024, 3, 10, 2, 30)), None, - ])); - let to_type = DataType::Timestamp(TimeUnit::Microsecond, Some("+0700".into())); - let timestamp_array = cast(&string_array, &to_type).unwrap(); - - // Check intermediate representation is correct - let string_array = cast(×tamp_array, &DataType::Utf8).unwrap(); - let result = string_array.as_string::(); - assert_eq!("2000-01-01T07:00:00.123456+07:00", result.value(0)); - assert_eq!("2010-01-01T07:00:00.123456+07:00", result.value(1)); - assert!(result.is_null(2)); - - let to_type = DataType::Timestamp(TimeUnit::Second, Some("-08:00".into())); - let timestamp_array = cast(×tamp_array, &to_type).unwrap(); + ]); + let to_type = DataType::Timestamp(TimeUnit::Second, Some("+08:00".into())); + let options = CastOptions { + safe: false, + ..Default::default() + }; - let string_array = cast(×tamp_array, &DataType::Utf8).unwrap(); - let result = string_array.as_string::(); - assert_eq!("1999-12-31T16:00:00-08:00", result.value(0)); - assert_eq!("2009-12-31T16:00:00-08:00", result.value(1)); - assert!(result.is_null(2)); + let b = cast_with_options(&array, &to_type, &options).unwrap(); + assert_eq!(b.data_type(), &to_type); + let c = b.as_primitive::(); + assert_eq!(c.value(0), instant_seconds(2024, 11, 1, 0, 0, 8)); + assert_eq!(c.value(1), instant_seconds(2024, 11, 3, 1, 30, 8)); + assert_eq!(c.value(2), instant_seconds(2024, 3, 10, 2, 30, 8)); + assert!(c.is_null(3)); } #[test] @@ -10012,6 +10437,120 @@ mod tests { test::(); } + /// Casts `input` to each of `targets` with `safe` true and false, and checks that the + /// result equals the cast of `expected` and has the same number of child values + fn check_cast_ignores_hidden_values( + input: &dyn Array, + expected: &dyn Array, + targets: &[DataType], + ) { + for target in targets { + for safe in [true, false] { + let options = CastOptions { + safe, + ..Default::default() + }; + let actual = cast_with_options(input, target, &options).unwrap(); + let expected = cast_with_options(expected, target, &options).unwrap(); + assert_eq!(actual.as_ref(), expected.as_ref()); + assert_eq!( + actual.to_data().child_data()[0].len(), + expected.to_data().child_data()[0].len() + ); + } + } + } + + #[test] + fn test_cast_sliced_list_ignores_hidden_values() { + fn test() { + // The first and last rows are outside the slice, and their values do not fit + // in Int16 + let input = GenericListArray::::from_iter_primitive::([ + Some(vec![Some(i32::MAX)]), + Some(vec![Some(1), Some(2)]), + Some(vec![]), + Some(vec![Some(3)]), + Some(vec![Some(i32::MAX); 2]), + ]); + let expected = GenericListArray::::from_iter_primitive::([ + Some(vec![Some(1), Some(2)]), + Some(vec![]), + Some(vec![Some(3)]), + ]); + // With a non-nullable child field, safe casts fail if they produce nulls + let field = Arc::new(Field::new_list_field(DataType::Int16, false)); + let targets = [ + DataType::List(field.clone()), + DataType::LargeList(field.clone()), + DataType::ListView(field.clone()), + DataType::LargeListView(field), + ]; + check_cast_ignores_hidden_values(&input.slice(1, 3), &expected, &targets); + } + test::(); + test::(); + } + + #[test] + fn test_cast_sliced_list_view_ignores_hidden_values() { + fn test() { + // Rows [[i32::MAX], [3], [], [1, 2], [i32::MAX]], out of order in the values. + // The empty row's offset is past the values of the other rows in the slice. + let input = GenericListViewArray::::new( + Arc::new(Field::new_list_field(DataType::Int32, true)), + ScalarBuffer::from_iter([0, 3, 5, 1, 4].map(O::usize_as)), + ScalarBuffer::from_iter([1, 1, 0, 2, 1].map(O::usize_as)), + Arc::new(Int32Array::from(vec![i32::MAX, 1, 2, 3, i32::MAX])), + None, + ); + let expected = GenericListViewArray::::from_iter_primitive::([ + Some(vec![Some(3)]), + Some(vec![]), + Some(vec![Some(1), Some(2)]), + ]); + let field = Arc::new(Field::new_list_field(DataType::Int16, false)); + let targets = [ + DataType::ListView(field.clone()), + DataType::LargeListView(field), + ]; + check_cast_ignores_hidden_values(&input.slice(1, 3), &expected, &targets); + } + test::(); + test::(); + } + + #[test] + fn test_cast_sliced_map_ignores_hidden_entries() { + // The first and last rows are outside the slice, and their keys do not fit in Int16 + let input = MapArray::from_vec_of_maps::( + vec![ + Some(vec![(i32::MAX, Some(0))]), + Some(vec![(1, Some(10)), (2, Some(20))]), + Some(vec![(3, Some(30))]), + Some(vec![(i32::MAX, Some(0))]), + ], + false, + ); + let expected = MapArray::from_vec_of_maps::( + vec![ + Some(vec![(1, Some(10)), (2, Some(20))]), + Some(vec![(3, Some(30))]), + ], + false, + ); + let entries = Field::new_struct( + "entries", + vec![ + Field::new("keys", DataType::Int16, false), + Field::new("values", DataType::Int32, true), + ], + false, + ); + let target = DataType::Map(Arc::new(entries), false); + check_cast_ignores_hidden_values(&input.slice(1, 2), &expected, &[target]); + } + #[test] fn test_cast_list_to_fsl() { // There four noteworthy cases we should handle: @@ -10565,6 +11104,63 @@ mod tests { ); } + /// Casts that cannot overflow, whether the source type or the data decides, + /// must match the checked conversion at the extremes, and a value beyond the + /// bound must still become null or an error. + #[test] + fn test_cast_integer_to_decimal_bounds() { + // Every i32 has at most ten digits, so the type decides. + let array = Int32Array::from(vec![Some(i32::MIN), None, Some(-1), Some(i32::MAX)]); + generate_cast_test_case!( + &array, + Decimal128Array, + &DataType::Decimal128(12, 2), + vec![ + Some(i32::MIN as i128 * 100), + None, + Some(-100), + Some(i32::MAX as i128 * 100) + ] + ); + // An i64 can exceed sixteen digits, so the data decides. + let bound = 9_999_999_999_999_999_i64; + let array = Int64Array::from(vec![Some(-bound), None, Some(bound)]); + generate_cast_test_case!( + &array, + Decimal128Array, + &DataType::Decimal128(18, 2), + vec![Some(-bound as i128 * 100), None, Some(bound as i128 * 100)] + ); + // A negative scale divides first, truncating toward zero. + let array = Int32Array::from(vec![Some(-1_234_567), Some(9_999_999)]); + generate_cast_test_case!( + &array, + Decimal32Array, + &DataType::Decimal32(5, -2), + vec![Some(-12_345), Some(99_999)] + ); + // One value beyond the bound sends the whole array down the checked path. + let array = Int64Array::from(vec![Some(bound + 1), Some(1)]); + let result = cast(&array, &DataType::Decimal128(18, 2)).unwrap(); + let result = result.as_primitive::(); + assert!(result.is_null(0)); + assert_eq!(result.value(1), 100); + let options = CastOptions { + safe: false, + ..Default::default() + }; + assert!(cast_with_options(&array, &DataType::Decimal128(18, 2), &options).is_err()); + // A precision beyond the target type is an error even when the values fit it. + let array = Int64Array::from(vec![5_000_000_000]); + for safe in [true, false] { + let options = CastOptions { + safe, + ..Default::default() + }; + assert!(cast_with_options(&array, &DataType::Decimal32(10, 0), &options).is_err()); + } + } + #[test] fn test_cast_numeric_to_decimal128_overflow() { let array = Int64Array::from(vec![i64::MAX]); @@ -10964,6 +11560,127 @@ mod tests { assert!(result.value(0) < 0.0); // Negative result } + #[test] + fn test_cast_decimal256_negative_scale_metadata() { + let input = Decimal256Array::from(vec![ + Some(i256::ZERO), + Some(i256::ONE), + Some(i256::MINUS_ONE), + None, + ]) + .with_precision_and_scale(76, -51) + .unwrap(); + let value = i256::from_i128(10).pow_wrapping(51); + let expected = + Decimal256Array::from(vec![Some(i256::ZERO), Some(value), Some(-value), None]) + .with_precision_and_scale(76, 0) + .unwrap(); + let options = CastOptions { + safe: false, + ..Default::default() + }; + let output_type = DataType::Decimal256(76, 0); + let result = cast(&input, &output_type).unwrap(); + assert_eq!(result.as_primitive::(), &expected); + let result = cast_with_options(&input, &output_type, &options).unwrap(); + assert_eq!(result.as_primitive::(), &expected); + + // The precision plus scale increase crosses i8::MAX at -52. + let input = input.with_precision_and_scale(76, -52).unwrap(); + let value = i256::from_i128(10).pow_wrapping(52); + let expected = + Decimal256Array::from(vec![Some(i256::ZERO), Some(value), Some(-value), None]) + .with_precision_and_scale(76, 0) + .unwrap(); + let result = cast(&input, &output_type).unwrap(); + assert_eq!(result.as_primitive::(), &expected); + let result = cast_with_options(&input, &output_type, &options).unwrap(); + assert_eq!(result.as_primitive::(), &expected); + } + + #[test] + fn test_cast_decimal256_negative_scale_precision_overflow() { + let input = Decimal256Array::from(vec![Some(i256::ZERO), Some(i256::ONE), None]) + .with_precision_and_scale(76, -76) + .unwrap(); + let output_type = DataType::Decimal256(76, 0); + let expected = Decimal256Array::from(vec![Some(i256::ZERO), None, None]) + .with_precision_and_scale(76, 0) + .unwrap(); + let result = cast(&input, &output_type).unwrap(); + assert_eq!(result.as_primitive::(), &expected); + + let options = CastOptions { + safe: false, + ..Default::default() + }; + let error = cast_with_options(&input, &output_type, &options).unwrap_err(); + let value = i256::from_i128(10).pow_wrapping(76); + assert_eq!( + error.to_string(), + format!( + "Invalid argument error: {value} is too large to store in a Decimal256 of precision 76. Max is {}", + value - i256::ONE + ) + ); + } + + #[test] + fn test_cast_decimal256_negative_scale_native_overflow() { + let input = Decimal256Array::from(vec![Some(i256::from_i128(6)), None]) + .with_precision_and_scale(76, -76) + .unwrap(); + let output_type = DataType::Decimal256(76, 0); + let expected = Decimal256Array::from(vec![None, None]) + .with_precision_and_scale(76, 0) + .unwrap(); + let result = cast(&input, &output_type).unwrap(); + assert_eq!(result.as_primitive::(), &expected); + + let options = CastOptions { + safe: false, + ..Default::default() + }; + let error = cast_with_options(&input, &output_type, &options).unwrap_err(); + assert_eq!( + error.to_string(), + "Cast error: Cannot cast to Decimal256(76, 0). Overflowing on 6" + ); + } + + #[test] + fn test_cast_decimal256_extreme_scale_difference() { + let input = Decimal256Array::from(vec![Some(i256::ONE), Some(i256::MINUS_ONE), None]) + .with_precision_and_scale(76, i8::MIN) + .unwrap(); + for safe in [true, false] { + let options = CastOptions { + safe, + ..Default::default() + }; + let error = + cast_with_options(&input, &DataType::Decimal256(76, 76), &options).unwrap_err(); + assert_eq!( + error.to_string(), + "Cast error: Cannot cast to Decimal256(76, 76). Value overflows for output scale" + ); + } + + let input = input.with_precision_and_scale(76, 76).unwrap(); + let expected = Decimal256Array::from(vec![Some(i256::ZERO), Some(i256::ZERO), None]) + .with_precision_and_scale(76, i8::MIN) + .unwrap(); + for safe in [true, false] { + let options = CastOptions { + safe, + ..Default::default() + }; + let result = + cast_with_options(&input, &DataType::Decimal256(76, i8::MIN), &options).unwrap(); + assert_eq!(result.as_primitive::(), &expected); + } + } + #[test] fn test_cast_decimal128_to_decimal128_negative_scale() { let input_type = DataType::Decimal128(20, 0); @@ -11861,6 +12578,34 @@ mod tests { assert_eq!(a.value(1), "\x00 Foo"); } + #[test] + fn test_cast_sliced_binary_with_hidden_invalid_utf8() { + // The invalid first value is outside the slice, so it must not fail the cast + let values: Vec<&[u8]> = vec![b"\xFF", b"a", b"b", b"c", b"d", b"e", b"f"]; + let arrays: [ArrayRef; 2] = [ + Arc::new(BinaryArray::from(values.clone()).slice(1, 6)), + Arc::new(LargeBinaryArray::from(values).slice(1, 6)), + ]; + let options = CastOptions { + safe: false, + format_options: FormatOptions::default(), + }; + let expected = StringArray::from(vec!["a", "b", "c", "d", "e", "f"]); + for array in &arrays { + for to_type in [DataType::Utf8, DataType::LargeUtf8, DataType::Utf8View] { + let result = cast_with_options(array, &to_type, &options).unwrap(); + assert_eq!(result.as_ref(), cast(&expected, &to_type).unwrap().as_ref()); + } + + // A dictionary with fewer keys than half its values takes a separate + // path to Utf8View + let keys = Int32Array::from(vec![5, 0]); + let dict = DictionaryArray::try_new(keys, Arc::clone(array)).unwrap(); + let result = cast_with_options(&dict, &DataType::Utf8View, &options).unwrap(); + assert_eq!(result.as_ref(), &StringViewArray::from(vec!["f", "a"])); + } + } + #[test] fn test_cast_utf8_to_timestamptz() { let valid = StringArray::from(vec!["2023-01-01"]); @@ -12558,11 +13303,14 @@ mod tests { }; for dt in data_types { - assert_eq!( - cast_with_options(&array, &dt, &cast_options) - .unwrap_err() - .to_string(), - "Parser error: Invalid timezone \"ZZTOP\": only offset based timezones supported without chrono-tz feature" + // The trailing detail of the message differs depending on whether + // `chrono-tz` is enabled, so only the common prefix is asserted. + let err = cast_with_options(&array, &dt, &cast_options) + .unwrap_err() + .to_string(); + assert!( + err.starts_with("Parser error: Invalid timezone \"ZZTOP\":"), + "{err}" ); } } @@ -13964,6 +14712,28 @@ mod tests { assert_eq!(got_list_view, &expected_list_view); } + #[test] + fn test_cast_list_to_list_view_null_in_non_nullable_field() { + // i32::MAX does not fit in Int16, so a safe cast turns it into a null + let list = ListArray::from_iter_primitive::([Some(vec![ + Some(1), + Some(i32::MAX), + ])]); + let field = Arc::new(Field::new_list_field(DataType::Int16, false)); + for (target_type, array_name) in [ + (DataType::ListView(field.clone()), "ListViewArray"), + (DataType::LargeListView(field), "LargeListViewArray"), + ] { + let err = cast(&list, &target_type).unwrap_err(); + assert_eq!( + err.to_string(), + format!( + "Invalid argument error: Non-nullable field of {array_name} \"item\" cannot contain nulls" + ) + ); + } + } + #[test] fn test_cast_large_list_view_to_large_list() { let list_view = diff --git a/arrow-cast/src/cast/string.rs b/arrow-cast/src/cast/string.rs index 8e4ff6a930c6..c193eacc6eea 100644 --- a/arrow-cast/src/cast/string.rs +++ b/arrow-cast/src/cast/string.rs @@ -18,6 +18,18 @@ use crate::cast::*; use arrow_buffer::NullBuffer; +pub(crate) fn cast_bool_to_string(array: &dyn Array) -> Result +where + A: Array + FromIterator> + 'static, +{ + let array: A = array + .as_boolean() + .iter() + .map(|v| v.map(|b| if b { "true" } else { "false" })) + .collect(); + Ok(Arc::new(array)) +} + pub(crate) fn value_to_string( array: &dyn Array, options: &CastOptions, @@ -357,9 +369,11 @@ pub(crate) fn cast_binary_to_string( Ok(a) => Ok(Arc::new(a)), Err(e) => match cast_options.safe { true => { - // Fallback to slow method to convert invalid sequences to nulls - let mut builder = - GenericStringBuilder::::with_capacity(array.len(), array.value_data().len()); + // Fallback to slow method to convert invalid sequences to nulls. + // Reserve only the bytes the offsets span, not the whole values buffer. + let offsets = array.offsets(); + let data_len = (offsets.last() - offsets.first()).as_usize(); + let mut builder = GenericStringBuilder::::with_capacity(array.len(), data_len); extend_valid_utf8(&mut builder, array.iter()); Ok(Arc::new(builder.finish())) diff --git a/arrow-cast/src/display.rs b/arrow-cast/src/display.rs index e27421f8c418..55f7d2661a01 100644 --- a/arrow-cast/src/display.rs +++ b/arrow-cast/src/display.rs @@ -674,7 +674,7 @@ impl<'a, F: DisplayIndexState<'a> + Array> DisplayIndex for ArrayFormat<'a, F> { impl DisplayIndex for &BooleanArray { fn write(&self, idx: usize, f: &mut dyn Write) -> FormatResult { - write!(f, "{}", self.value(idx))?; + f.write_str(if self.value(idx) { "true" } else { "false" })?; Ok(()) } } diff --git a/arrow-cast/src/parse.rs b/arrow-cast/src/parse.rs index ffc6738aa935..384eba200705 100644 --- a/arrow-cast/src/parse.rs +++ b/arrow-cast/src/parse.rs @@ -530,35 +530,65 @@ parser_primitive!(DurationSecondType); impl Parser for TimestampNanosecondType { fn parse(string: &str) -> Option { - string_to_timestamp_nanos(string).ok() + if let Ok(nanos) = string_to_timestamp_nanos(string) { + return Some(nanos); + } + + let trimmed = trim_pre_and_post_whitespace(string); + string_to_timestamp_nanos(trimmed).ok() } } impl Parser for TimestampMicrosecondType { fn parse(string: &str) -> Option { - let nanos = string_to_timestamp_nanos(string).ok(); - nanos.map(|x| x / 1000) + if let Ok(nanos) = string_to_timestamp_nanos(string) { + return Some(nanos / 1_000); + } + + let trimmed = trim_pre_and_post_whitespace(string); + string_to_timestamp_nanos(trimmed).ok().map(|x| x / 1_000) } } impl Parser for TimestampMillisecondType { fn parse(string: &str) -> Option { - let nanos = string_to_timestamp_nanos(string).ok(); - nanos.map(|x| x / 1_000_000) + if let Ok(nanos) = string_to_timestamp_nanos(string) { + return Some(nanos / 1_000_000); + } + + let trimmed = trim_pre_and_post_whitespace(string); + string_to_timestamp_nanos(trimmed) + .ok() + .map(|x| x / 1_000_000) } } impl Parser for TimestampSecondType { fn parse(string: &str) -> Option { - let nanos = string_to_timestamp_nanos(string).ok(); - nanos.map(|x| x / 1_000_000_000) + if let Ok(nanos) = string_to_timestamp_nanos(string) { + return Some(nanos / 1_000_000_000); + } + + let trimmed = trim_pre_and_post_whitespace(string); + string_to_timestamp_nanos(trimmed) + .ok() + .map(|x| x / 1_000_000_000) } } impl Parser for Time64NanosecondType { // Will truncate any fractions of a nanosecond fn parse(string: &str) -> Option { - string_to_time_nanoseconds(string) + let value = string_to_time_nanoseconds(string) + .ok() + .or_else(|| string.parse::().ok()); + + if value.is_some() { + return value; + } + + let trimmed = trim_pre_and_post_whitespace(string); + string_to_time_nanoseconds(trimmed) .ok() .or_else(|| string.parse::().ok()) } @@ -572,9 +602,19 @@ impl Parser for Time64NanosecondType { impl Parser for Time64MicrosecondType { // Will truncate any fractions of a microsecond fn parse(string: &str) -> Option { - string_to_time_nanoseconds(string) + let value = string_to_time_nanoseconds(string) .ok() .map(|nanos| nanos / 1_000) + .or_else(|| string.parse::().ok()); + + if value.is_some() { + return value; + } + + let trimmed = trim_pre_and_post_whitespace(string); + string_to_time_nanoseconds(trimmed) + .ok() + .map(|x| x / 1_000) .or_else(|| string.parse::().ok()) } @@ -587,9 +627,19 @@ impl Parser for Time64MicrosecondType { impl Parser for Time32MillisecondType { // Will truncate any fractions of a millisecond fn parse(string: &str) -> Option { - string_to_time_nanoseconds(string) + let value = string_to_time_nanoseconds(string) .ok() .map(|nanos| (nanos / 1_000_000) as i32) + .or_else(|| string.parse::().ok()); + + if value.is_some() { + return value; + } + + let trimmed = trim_pre_and_post_whitespace(string); + string_to_time_nanoseconds(trimmed) + .ok() + .map(|x| (x / 1_000_000) as i32) .or_else(|| string.parse::().ok()) } @@ -602,9 +652,19 @@ impl Parser for Time32MillisecondType { impl Parser for Time32SecondType { // Will truncate any fractions of a second fn parse(string: &str) -> Option { - string_to_time_nanoseconds(string) + let value = string_to_time_nanoseconds(string) .ok() .map(|nanos| (nanos / 1_000_000_000) as i32) + .or_else(|| string.parse::().ok()); + + if value.is_some() { + return value; + } + + let trimmed = trim_pre_and_post_whitespace(string); + string_to_time_nanoseconds(trimmed) + .ok() + .map(|x| (x / 1_000_000_000) as i32) .or_else(|| string.parse::().ok()) } @@ -749,7 +809,12 @@ fn parse_date_to_days(string: &str) -> Option { impl Parser for Date32Type { fn parse(string: &str) -> Option { - parse_date_to_days(string) + if let Some(days) = parse_date_to_days(string) { + return Some(days); + } + + let trimmed = trim_pre_and_post_whitespace(string); + parse_date_to_days(trimmed) } fn parse_formatted(string: &str, format: &str) -> Option { @@ -3623,4 +3688,93 @@ mod tests { assert_eq!(Int32Type::parse("3j"), None); assert_eq!(Int32Type::parse("3"), Some(3)); } + + #[test] + fn test_parse_temporal_with_surrounding_whitespace() { + let date = Date32Type::parse("2024-01-05"); + assert_eq!(Date32Type::parse(" 2024-01-05 "), date); + assert_eq!(Date32Type::parse("\t2024-01-05\n"), date); + + let timestamp = "2024-01-05T10:00:00"; + let padded_timestamp = " 2024-01-05T10:00:00 "; + + assert_eq!( + TimestampNanosecondType::parse(padded_timestamp), + TimestampNanosecondType::parse(timestamp) + ); + assert_eq!( + TimestampMicrosecondType::parse(padded_timestamp), + TimestampMicrosecondType::parse(timestamp) + ); + assert_eq!( + TimestampMillisecondType::parse(padded_timestamp), + TimestampMillisecondType::parse(timestamp) + ); + assert_eq!( + TimestampSecondType::parse(padded_timestamp), + TimestampSecondType::parse(timestamp) + ); + assert_eq!( + TimestampNanosecondType::parse("\t2024-01-05T10:00:00\n"), + TimestampNanosecondType::parse(timestamp) + ); + + // leet :) + let time = "13:37:00"; + let padded_time = " 13:37:00 "; + + assert_eq!( + Time64NanosecondType::parse(padded_time), + Time64NanosecondType::parse(time) + ); + assert_eq!( + Time64MicrosecondType::parse(padded_time), + Time64MicrosecondType::parse(time) + ); + assert_eq!( + Time64MicrosecondType::parse("\t13:37:00\n"), + Time64MicrosecondType::parse(time) + ); + } + + #[test] + fn test_parse_time_with_surrounding_ascii_whitespace() { + let time = "10:00:00"; + let padded = " 10:00:00 "; + let ascii_whitespace = "\t10:00:00\n"; + + assert_eq!( + Time64NanosecondType::parse(padded), + Time64NanosecondType::parse(time) + ); + assert_eq!( + Time64MicrosecondType::parse(padded), + Time64MicrosecondType::parse(time) + ); + assert_eq!( + Time32MillisecondType::parse(padded), + Time32MillisecondType::parse(time) + ); + assert_eq!( + Time32SecondType::parse(padded), + Time32SecondType::parse(time) + ); + + assert_eq!( + Time64NanosecondType::parse(ascii_whitespace), + Time64NanosecondType::parse(time) + ); + assert_eq!( + Time64MicrosecondType::parse(ascii_whitespace), + Time64MicrosecondType::parse(time) + ); + assert_eq!( + Time32MillisecondType::parse(ascii_whitespace), + Time32MillisecondType::parse(time) + ); + assert_eq!( + Time32SecondType::parse(ascii_whitespace), + Time32SecondType::parse(time) + ); + } } diff --git a/arrow-cmp/Cargo.toml b/arrow-cmp/Cargo.toml index 5a66ab341915..9109f2b9a1e9 100644 --- a/arrow-cmp/Cargo.toml +++ b/arrow-cmp/Cargo.toml @@ -19,7 +19,6 @@ name = "arrow-cmp" version = { workspace = true } description = "Basic comparator/ordering building blocks for Apache Arrow, shared by crates that need slot-wise comparison without pulling in the full arrow-ord kernels" -homepage = { workspace = true } repository = { workspace = true } authors = { workspace = true } license = { workspace = true } diff --git a/arrow-csv/Cargo.toml b/arrow-csv/Cargo.toml index aa1cbc21f60f..736c1777f46a 100644 --- a/arrow-csv/Cargo.toml +++ b/arrow-csv/Cargo.toml @@ -19,7 +19,6 @@ name = "arrow-csv" version = { workspace = true } description = "Support for parsing CSV format to and from the Arrow format" -homepage = { workspace = true } repository = { workspace = true } authors = { workspace = true } license = { workspace = true } diff --git a/arrow-data/Cargo.toml b/arrow-data/Cargo.toml index 3c51213474da..b546da6a2d4e 100644 --- a/arrow-data/Cargo.toml +++ b/arrow-data/Cargo.toml @@ -19,7 +19,6 @@ name = "arrow-data" version = { workspace = true } description = "Array data abstractions for Apache Arrow" -homepage = { workspace = true } repository = { workspace = true } authors = { workspace = true } license = { workspace = true } diff --git a/arrow-data/src/data.rs b/arrow-data/src/data.rs index 4128bffdd3d6..a8a69701256c 100644 --- a/arrow-data/src/data.rs +++ b/arrow-data/src/data.rs @@ -272,6 +272,19 @@ fn checked_len_plus_offset( }) } +/// Returns the first and last of `offsets`, and the bytes of `values` between +/// them as a `str`. Returns `None` if the offsets are not valid indexes into +/// `values` or those bytes are not valid UTF-8. +fn utf8_span<'a, T: ArrowNativeType>( + offsets: &[T], + values: &'a [u8], +) -> Option<(usize, usize, &'a str)> { + let first = offsets.first()?.to_usize()?; + let last = offsets.last()?.to_usize()?; + let values_str = std::str::from_utf8(values.get(first..last)?).ok()?; + Some((first, last, values_str)) +} + impl ArrayData { /// Create a new ArrayData instance; /// @@ -786,31 +799,46 @@ impl ArrayData { vec![ArrayData::new_empty(v.as_ref())], true, ), - DataType::Union(f, mode) => { - let (id, _) = f.iter().next().unwrap(); - let ids = Buffer::from_iter(std::iter::repeat_n(id, len)); - let buffers = match mode { - UnionMode::Sparse => vec![ids], - UnionMode::Dense => { - let end_offset = i32::from_usize(len).unwrap(); - vec![ids, Buffer::from_iter(0_i32..end_offset)] - } - }; - - let children = f - .iter() - .enumerate() - .map(|(idx, (_, f))| { - if idx == 0 || *mode == UnionMode::Sparse { - Self::new_null(f.data_type(), len) - } else { - Self::new_empty(f.data_type()) + DataType::Union(f, mode) => match f.iter().next() { + // Every slot carries a type id naming one of the children, + // so an empty union has nothing to put in a slot and can + // only be the empty array. + None => { + assert_eq!( + len, 0, + "cannot construct null data from an empty union of length {len}, a slot has no type id to carry" + ); + let buffers = match mode { + UnionMode::Sparse => vec![zeroed(0)], + UnionMode::Dense => vec![zeroed(0), zeroed(0)], + }; + (buffers, vec![], false) + } + Some((id, _)) => { + let ids = Buffer::from_iter(std::iter::repeat_n(id, len)); + let buffers = match mode { + UnionMode::Sparse => vec![ids], + UnionMode::Dense => { + let end_offset = i32::from_usize(len).unwrap(); + vec![ids, Buffer::from_iter(0_i32..end_offset)] } - }) - .collect(); + }; - (buffers, children, false) - } + let children = f + .iter() + .enumerate() + .map(|(idx, (_, f))| { + if idx == 0 || *mode == UnionMode::Sparse { + Self::new_null(f.data_type(), len) + } else { + Self::new_empty(f.data_type()) + } + }) + .collect(); + + (buffers, children, false) + } + }, DataType::RunEndEncoded(r, v) => { if len == 0 { // For empty arrays, create zero-length child arrays. @@ -1209,12 +1237,10 @@ impl ArrayData { DataType::List(field) => { let values_data = self.get_single_valid_child_data(field.data_type())?; self.validate_offsets::(values_data.len)?; - Ok(()) } DataType::LargeList(field) => { let values_data = self.get_single_valid_child_data(field.data_type())?; self.validate_offsets::(values_data.len)?; - Ok(()) } DataType::Map(field, _) => { let DataType::Struct(entries_fields) = field.data_type() else { @@ -1238,17 +1264,14 @@ impl ArrayData { } let values_data = self.get_single_valid_child_data(field.data_type())?; self.validate_offsets::(values_data.len)?; - Ok(()) } DataType::ListView(field) => { let values_data = self.get_single_valid_child_data(field.data_type())?; self.validate_offsets_and_sizes::(values_data.len)?; - Ok(()) } DataType::LargeListView(field) => { let values_data = self.get_single_valid_child_data(field.data_type())?; self.validate_offsets_and_sizes::(values_data.len)?; - Ok(()) } DataType::FixedSizeList(field, list_size) => { let values_data = self.get_single_valid_child_data(field.data_type())?; @@ -1260,18 +1283,18 @@ impl ArrayData { )) })?; - let expected_values_len = self.len + let len_plus_offset = + checked_len_plus_offset(&self.data_type, self.len, self.offset)?; + let expected_values_len = len_plus_offset .checked_mul(list_size) .expect("integer overflow computing expected number of expected values in FixedListSize"); if values_data.len < expected_values_len { return Err(ArrowError::InvalidArgumentError(format!( - "Values length {} is less than the length ({}) multiplied by the value size ({}) for {}", - values_data.len, self.len, list_size, self.data_type + "Values length {} is less than the length + offset ({}) multiplied by the value size ({}) for {}", + values_data.len, len_plus_offset, list_size, self.data_type ))); } - - Ok(()) } DataType::Struct(fields) => { self.validate_num_child_data(fields.len())?; @@ -1292,7 +1315,6 @@ impl ArrayData { ))); } } - Ok(()) } DataType::RunEndEncoded(run_ends_field, values_field) => { self.validate_num_child_data(2)?; @@ -1309,7 +1331,6 @@ impl ArrayData { "Found null values in run_ends array. The run_ends array should not have null values.".to_string(), )); } - Ok(()) } DataType::Union(fields, mode) => { self.validate_num_child_data(fields.len())?; @@ -1328,11 +1349,9 @@ impl ArrayData { } } } - Ok(()) } DataType::Dictionary(_key_type, value_type) => { self.get_single_valid_child_data(value_type)?; - Ok(()) } _ => { // other types do not have child data @@ -1343,9 +1362,9 @@ impl ArrayData { self.child_data.len() ))); } - Ok(()) } } + Ok(()) } /// Ensures that this array data has a single child_data with the @@ -1509,14 +1528,12 @@ impl ArrayData { DataType::FixedSizeList(field, len) => { let child = &self.child_data[0]; if !field.is_nullable() { - match &self.nulls { - Some(nulls) => { - let element_len = *len as usize; - let expanded = nulls.expand(element_len); - self.validate_non_nullable(Some(&expanded), child, child.nulls())?; - } - None => self.validate_non_nullable(None, child, child.nulls())?, - } + let element_len = *len as usize; + let child_nulls = child.nulls().map(|nulls| { + nulls.slice(self.offset * element_len, self.len * element_len) + }); + let expanded = self.nulls.as_ref().map(|nulls| nulls.expand(element_len)); + self.validate_non_nullable(expanded.as_ref(), child, child_nulls.as_ref())?; } } DataType::Struct(fields) => { @@ -1692,18 +1709,24 @@ impl ArrayData { } /// Ensures that all strings formed by the offsets in `buffers[0]` - /// into `buffers[1]` are valid utf8 sequences + /// into `buffers[1]` are valid utf8 sequences. Bytes of `buffers[1]` + /// outside the range of the offsets are not checked. fn validate_utf8(&self) -> Result<(), ArrowError> where T: ArrowNativeType + TryInto + num_traits::Num + std::fmt::Display, { let values_buffer = &self.buffer_at(1)?.as_slice(); - if let Ok(values_str) = std::str::from_utf8(values_buffer) { - // Validate Offsets are correct + let offsets = self.typed_offsets::()?; + if let Some((first, last, values_str)) = utf8_span(offsets, values_buffer) { self.validate_each_offset::(values_buffer.len(), |string_index, range| { - if !values_str.is_char_boundary(range.start) - || !values_str.is_char_boundary(range.end) - { + // `values_str` ends at the last offset. Checking each pair of offsets doesn't + // show that `range.end <= last`: a later offset can still be smaller, which + // `validate_each_offset` reports when it gets there. + if range.end > last { + return Ok(()); + } + let (start, end) = (range.start - first, range.end - first); + if !values_str.is_char_boundary(start) || !values_str.is_char_boundary(end) { return Err(ArrowError::InvalidArgumentError(format!( "incomplete utf-8 byte sequence from index {string_index}" ))); @@ -1711,7 +1734,7 @@ impl ArrayData { Ok(()) }) } else { - // find specific offset that failed utf8 validation + // Find specific offset that failed utf8 validation self.validate_each_offset::(values_buffer.len(), |string_index, range| { std::str::from_utf8(&values_buffer[range.clone()]).map_err(|e| { ArrowError::InvalidArgumentError(format!( @@ -2440,6 +2463,8 @@ pub(crate) fn get_fixed_size_binary_width(data_type: &DataType) -> usize { #[cfg(test)] mod tests { + use arrow_schema::UnionFields; + use super::*; use crate::ByteView; use crate::transform::MutableArrayData; @@ -2562,6 +2587,88 @@ mod tests { assert!(build(vec![true, true, false, true]).is_err()); } + #[test] + fn test_fixed_size_list_non_nullable_child_nulls_account_for_parent_offset() { + for child_offset in [0, 1] { + let child = ArrayData::builder(DataType::Int32) + .len(6) + .offset(child_offset) + .add_buffer(Buffer::from_slice_ref([0i32; 7])) + .nulls(Some(NullBuffer::from(vec![ + true, true, false, false, true, true, + ]))) + .build() + .unwrap(); + let build = |parent_nulls| { + ArrayData::builder(DataType::FixedSizeList( + Arc::new(Field::new("item", DataType::Int32, false)), + 2, + )) + .offset(1) + .len(2) + .nulls(Some(NullBuffer::from(parent_nulls))) + .add_child_data(child.clone()) + .build() + }; + + // The first visible list contains both child nulls. + assert!(build(vec![false, true]).is_ok()); + assert!(build(vec![true, false]).is_err()); + } + } + + #[test] + fn test_fixed_size_list_non_nullable_child_nulls_without_parent_nulls() { + let child = ArrayData::builder(DataType::Int32) + .len(6) + .add_buffer(Buffer::from_slice_ref([0i32; 6])) + .nulls(Some(NullBuffer::from(vec![ + false, false, true, true, false, false, + ]))) + .build() + .unwrap(); + let build = |offset, len, size| { + ArrayData::builder(DataType::FixedSizeList( + Arc::new(Field::new("item", DataType::Int32, false)), + size, + )) + .offset(offset) + .len(len) + .add_child_data(child.clone()) + .build() + }; + + // Nulls before and after the visible child range are irrelevant. + assert!(build(1, 1, 2).is_ok()); + assert!(build(0, 1, 2).is_err()); + assert!(build(2, 1, 2).is_err()); + assert!(build(3, 0, 2).is_ok()); + assert!(build(1, 2, 0).is_ok()); + } + + #[test] + fn test_fixed_size_list_validation_accounts_for_parent_offset() { + for nulls in [None, Some(NullBuffer::new_null(4))] { + let child = ArrayData::builder(DataType::Int32) + .len(4) + .add_buffer(Buffer::from_slice_ref([0i32; 4])) + .nulls(nulls) + .build() + .unwrap(); + let err = ArrayData::builder(DataType::FixedSizeList( + Arc::new(Field::new("item", DataType::Int32, false)), + 2, + )) + .offset(1) + .len(2) + .add_child_data(child) + .build() + .unwrap_err() + .to_string(); + assert!(err.contains("Values length 4 is less than")); + } + } + #[test] fn test_struct_equal_accounts_for_parent_offset() { let data_type = @@ -3677,4 +3784,23 @@ mod tests { [from_builder_res, from_try_new_res] } + + #[test] + fn test_new_null_empty_union() { + for mode in [UnionMode::Sparse, UnionMode::Dense] { + let data_type = DataType::Union(UnionFields::empty(), mode); + let data = ArrayData::new_null(&data_type, 0); + data.validate_full() + .expect("an empty union of length zero is valid"); + assert_eq!(data.len(), 0); + assert!(data.child_data().is_empty()); + } + } + + #[test] + #[should_panic(expected = "cannot construct null data from an empty union")] + fn test_new_null_empty_union_with_slots() { + let data_type = DataType::Union(UnionFields::empty(), UnionMode::Dense); + let _ = ArrayData::new_null(&data_type, 1); + } } diff --git a/arrow-data/src/equal/list_view.rs b/arrow-data/src/equal/list_view.rs index f431e84589df..285fbc8a9a04 100644 --- a/arrow-data/src/equal/list_view.rs +++ b/arrow-data/src/equal/list_view.rs @@ -17,7 +17,7 @@ use crate::ArrayData; use crate::data::count_nulls; -use crate::equal::equal_values; +use crate::equal::equal_range; use arrow_buffer::ArrowNativeType; use num_integer::Integer; @@ -66,8 +66,7 @@ pub(super) fn list_view_equal( let rhs_offset = rhs_offset.to_usize().unwrap(); let size = size.to_usize().unwrap(); - // Check if offsets are valid for the given range - if !equal_values(lhs_data, rhs_data, lhs_offset, rhs_offset, size) { + if !equal_range(lhs_data, rhs_data, lhs_offset, rhs_offset, size) { return false; } } @@ -77,10 +76,10 @@ pub(super) fn list_view_equal( let rhs_nulls = rhs.nulls().unwrap().slice(rhs_start, len); // Check values for equality, with null checking - for (index, ((&lhs_offset, &rhs_offset), &size)) in lhs_range_offsets + for (index, ((&lhs_offset, &rhs_offset), (&lhs_size, &rhs_size))) in lhs_range_offsets .iter() .zip(rhs_range_offsets) - .zip(lhs_range_sizes) + .zip(lhs_range_sizes.iter().zip(rhs_range_sizes)) .enumerate() { let lhs_is_null = lhs_nulls.is_null(index); @@ -90,12 +89,20 @@ pub(super) fn list_view_equal( return false; } + // Null rows are equal whatever their sizes and values + if lhs_is_null { + continue; + } + + if lhs_size != rhs_size { + return false; + } + let lhs_offset = lhs_offset.to_usize().unwrap(); let rhs_offset = rhs_offset.to_usize().unwrap(); - let size = size.to_usize().unwrap(); + let size = lhs_size.to_usize().unwrap(); - // Check if values match in the range - if !lhs_is_null && !equal_values(lhs_data, rhs_data, lhs_offset, rhs_offset, size) { + if !equal_range(lhs_data, rhs_data, lhs_offset, rhs_offset, size) { return false; } } diff --git a/arrow-data/src/equal/mod.rs b/arrow-data/src/equal/mod.rs index b241b8cf6641..011f591d1057 100644 --- a/arrow-data/src/equal/mod.rs +++ b/arrow-data/src/equal/mod.rs @@ -58,6 +58,9 @@ use variable_size::variable_sized_equal; /// Compares the values of two [ArrayData] starting at `lhs_start` and `rhs_start` respectively /// for `len` slots. +/// +/// Assumes the null masks of the two ranges are already known to be equal. To compare +/// ranges of child data, use [`equal_range`], which also compares the null masks. #[inline] fn equal_values( lhs: &ArrayData, diff --git a/arrow-flight/Cargo.toml b/arrow-flight/Cargo.toml index 4c47d4f3386d..dc5935d0cc0c 100644 --- a/arrow-flight/Cargo.toml +++ b/arrow-flight/Cargo.toml @@ -22,7 +22,6 @@ version = { workspace = true } edition = { workspace = true } rust-version = { workspace = true } authors = { workspace = true } -homepage = { workspace = true } repository = { workspace = true } license = { workspace = true } diff --git a/arrow-flight/gen/Cargo.toml b/arrow-flight/gen/Cargo.toml index da766e75e01e..ddd7259a27e7 100644 --- a/arrow-flight/gen/Cargo.toml +++ b/arrow-flight/gen/Cargo.toml @@ -22,7 +22,6 @@ version = "0.1.0" edition = { workspace = true } rust-version = { workspace = true } authors = { workspace = true } -homepage = { workspace = true } repository = { workspace = true } license = { workspace = true } publish = false diff --git a/arrow-flight/src/sql/server.rs b/arrow-flight/src/sql/server.rs index 6038f84efea4..138670e2e0a7 100644 --- a/arrow-flight/src/sql/server.rs +++ b/arrow-flight/src/sql/server.rs @@ -52,6 +52,8 @@ pub(crate) static END_SAVEPOINT: &str = "EndSavepoint"; pub(crate) static CANCEL_QUERY: &str = "CancelQuery"; /// Implements FlightSqlService to handle the flight sql protocol +// `async_trait` marks the boxed futures returned by provided methods `#[must_use]` +#[expect(clippy::double_must_use)] #[tonic::async_trait] pub trait FlightSqlService: Sync + Send + Sized + 'static { /// When impl FlightSqlService, you can always set FlightService to Self diff --git a/arrow-integration-test/Cargo.toml b/arrow-integration-test/Cargo.toml index 10fa57d744e1..ed306b888b47 100644 --- a/arrow-integration-test/Cargo.toml +++ b/arrow-integration-test/Cargo.toml @@ -19,7 +19,6 @@ name = "arrow-integration-test" version = { workspace = true } description = "Support for the Apache Arrow JSON test data format" -homepage = { workspace = true } repository = { workspace = true } authors = { workspace = true } license = { workspace = true } diff --git a/arrow-integration-testing/Cargo.toml b/arrow-integration-testing/Cargo.toml index ee6c344fa325..5e7b30efb759 100644 --- a/arrow-integration-testing/Cargo.toml +++ b/arrow-integration-testing/Cargo.toml @@ -19,7 +19,6 @@ name = "arrow-integration-testing" description = "Binaries used in the Arrow integration tests (NOT PUBLISHED TO crates.io)" version = { workspace = true } -homepage = { workspace = true } repository = { workspace = true } authors = { workspace = true } license = { workspace = true } diff --git a/arrow-ipc/Cargo.toml b/arrow-ipc/Cargo.toml index 71c4a4323ff1..3d0cd72e1f26 100644 --- a/arrow-ipc/Cargo.toml +++ b/arrow-ipc/Cargo.toml @@ -19,7 +19,6 @@ name = "arrow-ipc" version = { workspace = true } description = "Support for the Arrow IPC format" -homepage = { workspace = true } repository = { workspace = true } authors = { workspace = true } license = { workspace = true } diff --git a/arrow-ipc/src/writer.rs b/arrow-ipc/src/writer.rs index 30555a738a89..c435aa3147cf 100644 --- a/arrow-ipc/src/writer.rs +++ b/arrow-ipc/src/writer.rs @@ -1279,6 +1279,11 @@ pub(crate) fn unslice_run_array(arr: ArrayData) -> Result fn into_zero_offset_run_array( run_array: RunArray, ) -> Result, ArrowError> { + // Empty slices have no physical runs, regardless of their offset. + if run_array.is_empty() { + return Ok(ArrayData::new_empty(run_array.data_type()).into()); + } + let run_ends = run_array.run_ends(); if run_ends.offset() == 0 && run_ends.max_value() == run_ends.len() { return Ok(run_array); @@ -3764,6 +3769,56 @@ mod tests { assert_eq!(batch, batch2); } + fn assert_empty_run_array_roundtrip() { + let run_ends = PrimitiveArray::::from_iter_values( + [2, 5] + .into_iter() + .map(|v| R::Native::from_usize(v).unwrap()), + ); + let values = Int32Array::from(vec![10, 20]); + let array = RunArray::::try_new(&run_ends, &values).unwrap(); + let empty = RunArray::::from(ArrayData::new_empty(array.data_type())); + + for source in [&array, &empty] { + for offset in 0..=source.len() { + let sliced = source.slice(offset, 0); + let batch = RecordBatch::try_from_iter(vec![("run", Arc::new(sliced) as ArrayRef)]) + .unwrap(); + for decoded in [ + deserialize_stream(serialize_stream(&batch)), + deserialize_file(serialize_file(&batch)), + ] { + assert_eq!(decoded, batch); + let data = decoded.column(0).to_data(); + data.validate_full().unwrap(); + assert_eq!(data.offset(), 0); + assert!(data.child_data().iter().all(ArrayData::is_empty)); + } + + let normalized = into_zero_offset_run_array(source.slice(offset, 0)).unwrap(); + normalized.to_data().validate_full().unwrap(); + assert_eq!(normalized.offset(), 0); + assert!(normalized.run_ends().values().is_empty()); + assert!(normalized.values().is_empty()); + } + } + } + + #[test] + fn test_empty_run_array_roundtrip_i16() { + assert_empty_run_array_roundtrip::(); + } + + #[test] + fn test_empty_run_array_roundtrip_i32() { + assert_empty_run_array_roundtrip::(); + } + + #[test] + fn test_empty_run_array_roundtrip_i64() { + assert_empty_run_array_roundtrip::(); + } + #[test] fn test_run_array_unslice() { let total_len = 80; diff --git a/arrow-json/Cargo.toml b/arrow-json/Cargo.toml index c740450d61b0..7a377513a488 100644 --- a/arrow-json/Cargo.toml +++ b/arrow-json/Cargo.toml @@ -19,7 +19,6 @@ name = "arrow-json" version = { workspace = true } description = "Support for parsing JSON format to and from the Arrow format" -homepage = { workspace = true } repository = { workspace = true } authors = { workspace = true } license = { workspace = true } diff --git a/arrow-ord/Cargo.toml b/arrow-ord/Cargo.toml index d0e914e02abc..8d3f814f02ac 100644 --- a/arrow-ord/Cargo.toml +++ b/arrow-ord/Cargo.toml @@ -19,7 +19,6 @@ name = "arrow-ord" version = { workspace = true } description = "Ordering kernels for arrow arrays" -homepage = { workspace = true } repository = { workspace = true } authors = { workspace = true } license = { workspace = true } diff --git a/arrow-pyarrow-integration-testing/Cargo.toml b/arrow-pyarrow-integration-testing/Cargo.toml index b5762484eca0..65f903e51c7c 100644 --- a/arrow-pyarrow-integration-testing/Cargo.toml +++ b/arrow-pyarrow-integration-testing/Cargo.toml @@ -19,7 +19,6 @@ name = "arrow-pyarrow-integration-testing" description = "" version = "0.1.0" -homepage = "https://github.com/apache/arrow-rs" repository = "https://github.com/apache/arrow-rs" authors = ["Apache Arrow "] license = "Apache-2.0" diff --git a/arrow-pyarrow-testing/Cargo.toml b/arrow-pyarrow-testing/Cargo.toml index 688dfdd0ea6c..535f8ca20660 100644 --- a/arrow-pyarrow-testing/Cargo.toml +++ b/arrow-pyarrow-testing/Cargo.toml @@ -34,7 +34,6 @@ name = "arrow-pyarrow-testing" description = "Tests for arrow-pyarrow that require only a Python interpreter and pyarrow installed" version = "0.1.0" -homepage = "https://github.com/apache/arrow-rs" repository = "https://github.com/apache/arrow-rs" authors = ["Apache Arrow "] license = "Apache-2.0" diff --git a/arrow-pyarrow/Cargo.toml b/arrow-pyarrow/Cargo.toml index f7d6797c3f24..51e8092cf0e7 100644 --- a/arrow-pyarrow/Cargo.toml +++ b/arrow-pyarrow/Cargo.toml @@ -19,7 +19,6 @@ name = "arrow-pyarrow" version = { workspace = true } description = "Pyarrow bindings" -homepage = { workspace = true } repository = { workspace = true } authors = { workspace = true } license = { workspace = true } diff --git a/arrow-row/Cargo.toml b/arrow-row/Cargo.toml index a996300961d1..af4e3fd7eb28 100644 --- a/arrow-row/Cargo.toml +++ b/arrow-row/Cargo.toml @@ -19,7 +19,6 @@ name = "arrow-row" version = { workspace = true } description = "Arrow row format" -homepage = { workspace = true } repository = { workspace = true } authors = { workspace = true } license = { workspace = true } diff --git a/arrow-row/src/lib.rs b/arrow-row/src/lib.rs index d851e9dd5272..fa0d710e6c40 100644 --- a/arrow-row/src/lib.rs +++ b/arrow-row/src/lib.rs @@ -1502,7 +1502,7 @@ impl<'a> IntoIterator for &'a Rows { } /// An iterator over [`Rows`] -#[derive(Debug)] +#[derive(Debug, Clone)] pub struct RowsIter<'a> { rows: &'a Rows, start: usize, @@ -1527,6 +1527,48 @@ impl<'a> Iterator for RowsIter<'a> { let len = self.len(); (len, Some(len)) } + + fn count(self) -> usize + where + Self: Sized, + { + self.len() + } + + #[inline] + fn nth(&mut self, n: usize) -> Option { + // Check if we can advance to the desired offset. + // When n is 0 it means we want the next() value + // and when n is 1 we want the next().next() value + // so adding n to the current offset and not n - 1 + match self.start.checked_add(n) { + // Yes, and still within bounds + Some(new_offset) if new_offset < self.end => { + self.start = new_offset; + } + + // Either overflow or would exceed end + _ => { + self.start = self.end; + return None; + } + } + + self.next() + } + + fn last(mut self) -> Option { + // If already at the end, return None + if self.start == self.end { + return None; + } + + // Go to the one before the last bit + self.start = self.end - 1; + + // Return the last bit + self.next() + } } impl ExactSizeIterator for RowsIter<'_> { @@ -1548,6 +1590,27 @@ impl DoubleEndedIterator for RowsIter<'_> { let row = unsafe { self.rows.row_unchecked(self.end) }; Some(row) } + + fn nth_back(&mut self, n: usize) -> Option { + // Check if we can advance to the desired offset. + // When n is 0 it means we want the next_back() value + // and when n is 1 we want the next_back().next_back() value + // so subtracting n to the current offset and not n - 1 + match self.end.checked_sub(n) { + // Yes, and still within bounds + Some(new_offset) if self.start < new_offset => { + self.end = new_offset; + } + + // Either underflow or would exceed start + _ => { + self.start = self.end; + return None; + } + } + + self.next_back() + } } /// A comparable representation of a row. @@ -2230,11 +2293,11 @@ unsafe fn decode_column( DataType::Boolean => Arc::new(decode_bool(rows, options)), DataType::Binary => Arc::new(decode_binary::(rows, options)), DataType::LargeBinary => Arc::new(decode_binary::(rows, options)), - DataType::BinaryView => Arc::new(decode_binary_view(rows, options)), + DataType::BinaryView => Arc::new(decode_binary_view(rows, options)?), DataType::FixedSizeBinary(size) => Arc::new(decode_fixed_size_binary(rows, size, options)), DataType::Utf8 => Arc::new(unsafe{ decode_string::(rows, options, validate_utf8) }), DataType::LargeUtf8 => Arc::new(unsafe { decode_string::(rows, options, validate_utf8) }), - DataType::Utf8View => Arc::new(unsafe { decode_string_view(rows, options, validate_utf8) }), + DataType::Utf8View => Arc::new(unsafe { decode_string_view(rows, options, validate_utf8) }?), _ => return Err(ArrowError::NotYetImplemented(format!("unsupported data type: {data_type}" ))) } } @@ -2439,6 +2502,7 @@ unsafe fn decode_column( #[cfg(test)] mod tests { + use super::*; use arrow_array::builder::*; use arrow_array::types::*; use arrow_buffer::{Buffer, OffsetBuffer}; @@ -2449,8 +2513,9 @@ mod tests { use rand::distr::{Distribution, StandardUniform}; use rand::prelude::StdRng; use rand::{RngExt, SeedableRng}; - - use super::*; + use std::fmt::Debug; + use std::iter::Copied; + use std::slice::Iter; fn all_sort_options() -> [SortOptions; 4] { [ @@ -6778,4 +6843,429 @@ mod tests { assert_eq!(rows_iter.next_back(), None); assert_eq!(rows_iter.next(), None); } + + trait SharedBetweenRowsIteratorAndSliceIter<'a>: + ExactSizeIterator> + DoubleEndedIterator> + { + } + impl<'a, T: ?Sized + ExactSizeIterator> + DoubleEndedIterator>> + SharedBetweenRowsIteratorAndSliceIter<'a> for T + { + } + + fn get_rows_iterator_cases() -> impl Iterator { + let rows_converter = RowConverter::new(vec![SortField::new(DataType::Int32)]).unwrap(); + + [0, 1, 6, 8, 100, 164] + .map(|len| { + let source = (0..).take(len).collect::>(); + + let source = Int32Array::from(source); + + rows_converter.convert_columns(&[Arc::new(source)]).unwrap() + }) + .into_iter() + } + + fn setup_and_assert( + setup_iters: impl Fn(&mut dyn SharedBetweenRowsIteratorAndSliceIter), + assert_fn: impl Fn(RowsIter, Copied>), + ) { + for rows in get_rows_iterator_cases() { + let expected = (0..rows.num_rows()) + .map(|i| rows.row(i)) + .collect::>(); + let mut expected_iter = expected.iter().copied(); + + let mut actual = rows.iter(); + + setup_iters(&mut actual); + setup_iters(&mut expected_iter); + + assert_fn(actual, expected_iter); + } + } + + /// Trait representing an operation on a [`RowsIter`] + /// that can be compared against a slice iterator + trait RowsIteratorOp { + /// What the operation returns (e.g. Option for last, usize for count, etc) + type Output<'a>: PartialEq + Debug; + + /// The name of the operation, used for error messages + const NAME: &'static str; + + /// Get the value of the operation for the provided iterator + /// This will be either a [`RowsIter`] or a slice iterator to make sure they produce the same result + fn get_value<'a, T: SharedBetweenRowsIteratorAndSliceIter<'a>>(iter: T) + -> Self::Output<'a>; + } + + /// Helper function that will assert that the provided operation + /// produces the same result for both [`RowsIter`] and slice iterator + /// under various consumption patterns (e.g. some calls to next/next_back/consume_all/etc) + fn assert_rows_iterator_cases() { + setup_and_assert( + |_iter: &mut dyn SharedBetweenRowsIteratorAndSliceIter| {}, + |actual, expected| { + let current_iterator_values: Vec = expected.clone().collect(); + assert_eq!( + O::get_value(actual), + O::get_value(expected), + "Failed on op {} for new iter (left actual, right expected) ({current_iterator_values:?})", + O::NAME + ); + }, + ); + + setup_and_assert( + |iter: &mut dyn SharedBetweenRowsIteratorAndSliceIter| { + iter.next(); + }, + |actual, expected| { + let current_iterator_values: Vec = expected.clone().collect(); + + assert_eq!( + O::get_value(actual), + O::get_value(expected), + "Failed on op {} for new iter after consuming 1 element from the start (left actual, right expected) ({current_iterator_values:?})", + O::NAME + ); + }, + ); + + setup_and_assert( + |iter: &mut dyn SharedBetweenRowsIteratorAndSliceIter| { + iter.next_back(); + }, + |actual, expected| { + let current_iterator_values: Vec = expected.clone().collect(); + + assert_eq!( + O::get_value(actual), + O::get_value(expected), + "Failed on op {} for new iter after consuming 1 element from the end (left actual, right expected) ({current_iterator_values:?})", + O::NAME + ); + }, + ); + + setup_and_assert( + |iter: &mut dyn SharedBetweenRowsIteratorAndSliceIter| { + iter.next(); + iter.next_back(); + }, + |actual, expected| { + let current_iterator_values: Vec = expected.clone().collect(); + + assert_eq!( + O::get_value(actual), + O::get_value(expected), + "Failed on op {} for new iter after consuming 1 element from start and end (left actual, right expected) ({current_iterator_values:?})", + O::NAME + ); + }, + ); + + setup_and_assert( + |iter: &mut dyn SharedBetweenRowsIteratorAndSliceIter| { + while iter.len() > 1 { + iter.next(); + } + }, + |actual, expected| { + let current_iterator_values: Vec = expected.clone().collect(); + + assert_eq!( + O::get_value(actual), + O::get_value(expected), + "Failed on op {} for new iter after consuming all from the start but 1 (left actual, right expected) ({current_iterator_values:?})", + O::NAME + ); + }, + ); + + setup_and_assert( + |iter: &mut dyn SharedBetweenRowsIteratorAndSliceIter| { + while iter.len() > 1 { + iter.next_back(); + } + }, + |actual, expected| { + let current_iterator_values: Vec = expected.clone().collect(); + + assert_eq!( + O::get_value(actual), + O::get_value(expected), + "Failed on op {} for new iter after consuming all from the end but 1 (left actual, right expected) ({current_iterator_values:?})", + O::NAME + ); + }, + ); + + setup_and_assert( + |iter: &mut dyn SharedBetweenRowsIteratorAndSliceIter| { + while iter.next().is_some() {} + }, + |actual, expected| { + let current_iterator_values: Vec = expected.clone().collect(); + + assert_eq!( + O::get_value(actual), + O::get_value(expected), + "Failed on op {} for new iter after consuming all from the start (left actual, right expected) ({current_iterator_values:?})", + O::NAME + ); + }, + ); + + setup_and_assert( + |iter: &mut dyn SharedBetweenRowsIteratorAndSliceIter| { + while iter.next_back().is_some() {} + }, + |actual, expected| { + let current_iterator_values: Vec = expected.clone().collect(); + + assert_eq!( + O::get_value(actual), + O::get_value(expected), + "Failed on op {} for new iter after consuming all from the end (left actual, right expected) ({current_iterator_values:?})", + O::NAME + ); + }, + ); + } + + #[test] + fn assert_rows_iterator_count() { + struct CountOp; + + impl RowsIteratorOp for CountOp { + type Output<'a> = usize; + const NAME: &'static str = "count"; + + fn get_value<'a, T: SharedBetweenRowsIteratorAndSliceIter<'a>>( + iter: T, + ) -> Self::Output<'a> { + iter.count() + } + } + + assert_rows_iterator_cases::() + } + + #[test] + fn assert_rows_iterator_last() { + struct LastOp; + + impl RowsIteratorOp for LastOp { + type Output<'a> = Option>; + const NAME: &'static str = "last"; + + fn get_value<'a, T: SharedBetweenRowsIteratorAndSliceIter<'a>>( + iter: T, + ) -> Self::Output<'a> { + iter.last() + } + } + + assert_rows_iterator_cases::() + } + + #[test] + fn assert_rows_iterator_nth_0() { + struct NthOp; + + impl RowsIteratorOp for NthOp { + type Output<'a> = Option>; + const NAME: &'static str = if BACK { "nth_back(0)" } else { "nth(0)" }; + + fn get_value<'a, T: SharedBetweenRowsIteratorAndSliceIter<'a>>( + mut iter: T, + ) -> Self::Output<'a> { + if BACK { iter.nth_back(0) } else { iter.nth(0) } + } + } + + assert_rows_iterator_cases::>(); + assert_rows_iterator_cases::>(); + } + + #[test] + fn assert_rows_iterator_nth_1() { + struct NthOp; + + impl RowsIteratorOp for NthOp { + type Output<'a> = Option>; + const NAME: &'static str = if BACK { "nth_back(1)" } else { "nth(1)" }; + + fn get_value<'a, T: SharedBetweenRowsIteratorAndSliceIter<'a>>( + mut iter: T, + ) -> Self::Output<'a> { + if BACK { iter.nth_back(1) } else { iter.nth(1) } + } + } + + assert_rows_iterator_cases::>(); + assert_rows_iterator_cases::>(); + } + + #[test] + fn assert_rows_iterator_nth_after_end() { + struct NthOp; + + impl RowsIteratorOp for NthOp { + type Output<'a> = Option>; + const NAME: &'static str = if BACK { + "nth_back(iter.len() + 1)" + } else { + "nth(iter.len() + 1)" + }; + + fn get_value<'a, T: SharedBetweenRowsIteratorAndSliceIter<'a>>( + mut iter: T, + ) -> Self::Output<'a> { + if BACK { + iter.nth_back(iter.len() + 1) + } else { + iter.nth(iter.len() + 1) + } + } + } + + assert_rows_iterator_cases::>(); + assert_rows_iterator_cases::>(); + } + + #[test] + fn assert_rows_iterator_nth_len() { + struct NthOp; + + impl RowsIteratorOp for NthOp { + type Output<'a> = Option>; + const NAME: &'static str = if BACK { + "nth_back(iter.len())" + } else { + "nth(iter.len())" + }; + + fn get_value<'a, T: SharedBetweenRowsIteratorAndSliceIter<'a>>( + mut iter: T, + ) -> Self::Output<'a> { + if BACK { + iter.nth_back(iter.len()) + } else { + iter.nth(iter.len()) + } + } + } + + assert_rows_iterator_cases::>(); + assert_rows_iterator_cases::>(); + } + + #[test] + fn assert_rows_iterator_nth_last() { + struct NthOp; + + impl RowsIteratorOp for NthOp { + type Output<'a> = Option>; + const NAME: &'static str = if BACK { + "nth_back(iter.len().saturating_sub(1))" + } else { + "nth(iter.len().saturating_sub(1))" + }; + + fn get_value<'a, T: SharedBetweenRowsIteratorAndSliceIter<'a>>( + mut iter: T, + ) -> Self::Output<'a> { + if BACK { + iter.nth_back(iter.len().saturating_sub(1)) + } else { + iter.nth(iter.len().saturating_sub(1)) + } + } + } + + assert_rows_iterator_cases::>(); + assert_rows_iterator_cases::>(); + } + + #[test] + fn assert_rows_iterator_nth_and_reuse() { + setup_and_assert( + |_| {}, + |actual, expected| { + { + let mut actual = actual.clone(); + let mut expected = expected.clone(); + for _ in 0..expected.len() { + #[expect(clippy::iter_nth_zero)] + let actual_val = actual.nth(0); + #[expect(clippy::iter_nth_zero)] + let expected_val = expected.nth(0); + assert_eq!(actual_val, expected_val, "Failed on nth(0)"); + } + } + + { + let mut actual = actual.clone(); + let mut expected = expected.clone(); + for _ in 0..expected.len() { + let actual_val = actual.nth(1); + let expected_val = expected.nth(1); + assert_eq!(actual_val, expected_val, "Failed on nth(1)"); + } + } + + { + let mut actual = actual.clone(); + let mut expected = expected.clone(); + for _ in 0..expected.len() { + let actual_val = actual.nth(2); + let expected_val = expected.nth(2); + assert_eq!(actual_val, expected_val, "Failed on nth(2)"); + } + } + }, + ); + } + + #[test] + fn assert_rows_iterator_nth_back_and_reuse() { + setup_and_assert( + |_| {}, + |actual, expected| { + { + let mut actual = actual.clone(); + let mut expected = expected.clone(); + for _ in 0..expected.len() { + let actual_val = actual.nth_back(0); + let expected_val = expected.nth_back(0); + assert_eq!(actual_val, expected_val, "Failed on nth_back(0)"); + } + } + + { + let mut actual = actual.clone(); + let mut expected = expected.clone(); + for _ in 0..expected.len() { + let actual_val = actual.nth_back(1); + let expected_val = expected.nth_back(1); + assert_eq!(actual_val, expected_val, "Failed on nth_back(1)"); + } + } + + { + let mut actual = actual.clone(); + let mut expected = expected.clone(); + for _ in 0..expected.len() { + let actual_val = actual.nth_back(2); + let expected_val = expected.nth_back(2); + assert_eq!(actual_val, expected_val, "Failed on nth_back(2)"); + } + } + }, + ); + } } diff --git a/arrow-row/src/variable.rs b/arrow-row/src/variable.rs index 2f0e5f2c35c9..814d9c3b5d7b 100644 --- a/arrow-row/src/variable.rs +++ b/arrow-row/src/variable.rs @@ -23,7 +23,7 @@ use arrow_buffer::{ ArrowNativeType, BooleanBuffer, MutableBuffer, NullBuffer, OffsetBuffer, ScalarBuffer, }; use arrow_data::MAX_INLINE_VIEW_LEN; -use arrow_schema::SortOptions; +use arrow_schema::{ArrowError, SortOptions}; use builder::make_view; /// The block size of the variable length encoding @@ -204,8 +204,6 @@ fn encode_blocks(out: &mut [u8], val: &[u8]) -> usize { let to_write = &mut out[..end_offset]; let (chunks, remainder) = val.as_chunks::(); - #[expect(clippy::chunks_exact_to_as_chunks)] - // Requires using generic parameters in const operations: generic_const_exprs let to_write_chunks = to_write.chunks_exact_mut(SIZE + 1); for (input, output) in chunks.iter().zip(to_write_chunks) { let out_block: &mut [u8; SIZE] = (&mut output[..SIZE]).try_into().unwrap(); @@ -312,7 +310,7 @@ pub fn decode_binary( fn decode_binary_view_inner( rows: &mut [&[u8]], options: SortOptions, -) -> BinaryViewArray { +) -> Result { let len = rows.len(); let inline_str_max_len = MAX_INLINE_VIEW_LEN as usize; @@ -329,6 +327,15 @@ fn decode_binary_view_inner( inline_capacity += len; } } + // Every view offset is at most the total length of the non-inlined values, + // so checking it once here keeps the offsets below within the `i32` range + // the Arrow spec requires + let long_values_len = values_capacity - inline_str_max_len; + if long_values_len > i32::MAX as usize { + return Err(ArrowError::InvalidArgumentError(format!( + "{long_values_len} bytes of non-inlined values too long to decode into a view array with a single data buffer" + ))); + } let mut values = MutableBuffer::new(values_capacity); let mut view_utf8_validation_buffer = if VALIDATE_UTF8 { Vec::with_capacity(inline_capacity) @@ -375,11 +382,17 @@ fn decode_binary_view_inner( // SAFETY: // Valid by construction above - unsafe { BinaryViewArray::new_unchecked(views.into(), [values.into()].into(), nulls) } + Ok(unsafe { BinaryViewArray::new_unchecked(views.into(), [values.into()].into(), nulls) }) } /// Decodes a binary view array from `rows` with the provided `options` -pub fn decode_binary_view(rows: &mut [&[u8]], options: SortOptions) -> BinaryViewArray { +/// +/// Returns an error if the non-inlined values are longer than `i32::MAX` bytes, +/// the most a single data buffer can address +pub fn decode_binary_view( + rows: &mut [&[u8]], + options: SortOptions, +) -> Result { decode_binary_view_inner::(rows, options) } @@ -411,17 +424,20 @@ pub unsafe fn decode_string( /// # Safety /// /// The row must contain valid UTF-8 data +/// +/// Returns an error if the non-inlined values are longer than `i32::MAX` bytes, +/// the most a single data buffer can address pub unsafe fn decode_string_view( rows: &mut [&[u8]], options: SortOptions, validate_utf8: bool, -) -> StringViewArray { +) -> Result { let view = if validate_utf8 { - decode_binary_view_inner::(rows, options) + decode_binary_view_inner::(rows, options)? } else { - decode_binary_view_inner::(rows, options) + decode_binary_view_inner::(rows, options)? }; - unsafe { view.to_string_view_unchecked() } + Ok(unsafe { view.to_string_view_unchecked() }) } pub fn decode_null_value(rows: &mut [&[u8]], options: SortOptions) { diff --git a/arrow-schema/Cargo.toml b/arrow-schema/Cargo.toml index e85f8850d3e5..ae62aa4c42f9 100644 --- a/arrow-schema/Cargo.toml +++ b/arrow-schema/Cargo.toml @@ -19,7 +19,6 @@ name = "arrow-schema" version = { workspace = true } description = "Defines the logical types for arrow arrays" -homepage = { workspace = true } repository = { workspace = true } authors = { workspace = true } license = { workspace = true } diff --git a/arrow-schema/src/error.rs b/arrow-schema/src/error.rs index b9818b1b921d..35d480d54998 100644 --- a/arrow-schema/src/error.rs +++ b/arrow-schema/src/error.rs @@ -60,7 +60,7 @@ pub enum ArrowError { CDataInterface(String), /// Error when a dictionary key is bigger than the key type DictionaryKeyOverflowError, - /// Error when the run end index in a REE array is bigger than the array length + /// Error when a run-end index is bigger than the run-end type RunEndIndexOverflowError, /// Error when the offset overflows. OffsetOverflowError(usize), @@ -129,7 +129,7 @@ impl Display for ArrowError { write!(f, "Dictionary key bigger than the key type") } ArrowError::RunEndIndexOverflowError => { - write!(f, "Run end encoded array index overflow error") + write!(f, "Run end index bigger than the run-end type") } ArrowError::OffsetOverflowError(offset) => { write!(f, "Offset overflow error: {offset}") diff --git a/arrow-schema/src/extension/canonical/fixed_closedness_range.rs b/arrow-schema/src/extension/canonical/fixed_closedness_range.rs new file mode 100644 index 000000000000..b11264827f96 --- /dev/null +++ b/arrow-schema/src/extension/canonical/fixed_closedness_range.rs @@ -0,0 +1,597 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Fixed closedness range +//! +//! + +use serde_core::de::{IgnoredAny, MapAccess, Visitor}; +use serde_core::ser::SerializeStruct; +use serde_core::{Deserialize, Deserializer, Serialize, Serializer}; + +use crate::{ArrowError, DataType, Fields, extension::ExtensionType}; + +/// Which bound(s) of a [`FixedClosednessRange`] interval are inclusive. +/// +/// Uses the same vocabulary as pandas: "left", "right", "both", "neither". +/// A null (unbounded) bound is always treated as exclusive, regardless of +/// this value. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum RangeClosed { + /// The left (lower) endpoint is included; the right (upper) is excluded. + Left, + /// The left (lower) endpoint is excluded; the right (upper) is included. + Right, + /// Both endpoints are included (closed interval). + Both, + /// Neither endpoint is included (open interval). + Neither, +} + +impl Serialize for RangeClosed { + fn serialize(&self, serializer: S) -> Result + where + S: Serializer, + { + serializer.serialize_str(match self { + RangeClosed::Left => "left", + RangeClosed::Right => "right", + RangeClosed::Both => "both", + RangeClosed::Neither => "neither", + }) + } +} + +struct RangeClosedVisitor; + +impl Visitor<'_> for RangeClosedVisitor { + type Value = RangeClosed; + + fn expecting(&self, formatter: &mut std::fmt::Formatter) -> std::fmt::Result { + formatter.write_str("one of \"left\", \"right\", \"both\", \"neither\"") + } + + fn visit_str(self, value: &str) -> Result + where + E: serde_core::de::Error, + { + match value { + "left" => Ok(RangeClosed::Left), + "right" => Ok(RangeClosed::Right), + "both" => Ok(RangeClosed::Both), + "neither" => Ok(RangeClosed::Neither), + _ => Err(serde_core::de::Error::unknown_variant( + value, + &["left", "right", "both", "neither"], + )), + } + } +} + +impl<'de> Deserialize<'de> for RangeClosed { + fn deserialize(deserializer: D) -> Result + where + D: Deserializer<'de>, + { + deserializer.deserialize_str(RangeClosedVisitor) + } +} + +/// Extension type metadata for [`FixedClosednessRange`]. +#[derive(Debug, Clone, PartialEq)] +pub struct FixedClosednessRangeMetadata { + /// Whether the interval endpoints are included or excluded. + closed: RangeClosed, +} + +impl FixedClosednessRangeMetadata { + /// Returns a new `FixedClosednessRangeMetadata`. + pub fn new(closed: RangeClosed) -> Self { + FixedClosednessRangeMetadata { closed } + } + + /// Returns whether the interval endpoints are included or excluded. + pub fn closed(&self) -> RangeClosed { + self.closed + } +} + +impl Serialize for FixedClosednessRangeMetadata { + fn serialize(&self, serializer: S) -> Result + where + S: Serializer, + { + let mut state = serializer.serialize_struct("FixedClosednessRangeMetadata", 1)?; + state.serialize_field("closed", &self.closed)?; + state.end() + } +} + +#[derive(Debug)] +enum MetadataField { + Closed, + /// Any key other than `closed`. Unknown keys are ignored to allow + /// forward-compatible extensions, as required by the specification. + Other, +} + +struct MetadataFieldVisitor; + +impl Visitor<'_> for MetadataFieldVisitor { + type Value = MetadataField; + + fn expecting(&self, formatter: &mut std::fmt::Formatter) -> std::fmt::Result { + formatter.write_str("a metadata field name") + } + + fn visit_str(self, value: &str) -> Result + where + E: serde_core::de::Error, + { + match value { + "closed" => Ok(MetadataField::Closed), + _ => Ok(MetadataField::Other), + } + } +} + +impl<'de> Deserialize<'de> for MetadataField { + fn deserialize(deserializer: D) -> Result + where + D: Deserializer<'de>, + { + deserializer.deserialize_identifier(MetadataFieldVisitor) + } +} + +struct FixedClosednessRangeMetadataVisitor; + +impl<'de> Visitor<'de> for FixedClosednessRangeMetadataVisitor { + type Value = FixedClosednessRangeMetadata; + + fn expecting(&self, formatter: &mut std::fmt::Formatter) -> std::fmt::Result { + formatter.write_str("struct FixedClosednessRangeMetadata") + } + + fn visit_seq(self, mut seq: V) -> Result + where + V: serde_core::de::SeqAccess<'de>, + { + let closed = seq + .next_element()? + .ok_or_else(|| serde_core::de::Error::invalid_length(0, &self))?; + Ok(FixedClosednessRangeMetadata { closed }) + } + + fn visit_map(self, mut map: V) -> Result + where + V: MapAccess<'de>, + { + let mut closed = None; + + while let Some(key) = map.next_key()? { + match key { + MetadataField::Closed => { + if closed.is_some() { + return Err(serde_core::de::Error::duplicate_field("closed")); + } + closed = Some(map.next_value()?); + } + MetadataField::Other => { + // Unknown keys are ignored for forward compatibility. + map.next_value::()?; + } + } + } + + let closed = closed.ok_or_else(|| serde_core::de::Error::missing_field("closed"))?; + Ok(FixedClosednessRangeMetadata { closed }) + } +} + +impl<'de> Deserialize<'de> for FixedClosednessRangeMetadata { + fn deserialize(deserializer: D) -> Result + where + D: Deserializer<'de>, + { + deserializer.deserialize_struct( + "FixedClosednessRangeMetadata", + &["closed"], + FixedClosednessRangeMetadataVisitor, + ) + } +} + +/// The extension type for a bounded set (mathematical interval) whose +/// closedness is a type-level parameter shared by all values. +/// +/// Extension name: `arrow.fixed_closedness_range`. +/// +/// A range is a bounded set (mathematical interval) defined by a lower and +/// an upper bound over an orderable value type T. T may be any orderable +/// Arrow type, for example an integer, floating-point, decimal, date, time, +/// timestamp, duration, string or binary type. This specification defines +/// only the storage layout, not the order of any type. +/// +/// The storage type is a `Struct` with exactly two fields, in order: +/// - `lower` (type T): the lower bound. +/// - `upper` (type T): the upper bound. +/// +/// Both fields must share the same data type T. Each bound may +/// independently be nullable or non-nullable: a nullable bound can hold +/// null to represent an unbounded (infinite) endpoint on that side, while a +/// non-nullable bound is always finite. A null bound is always treated as +/// exclusive, regardless of the `closed` parameter: only a null bound means +/// the range is unbounded on that side. +/// +/// The `closed` parameter specifies which non-null bound(s) are inclusive; +/// see [`RangeClosed`]. It is required and is not defaulted on the wire: an +/// empty metadata string, or a JSON object without a `closed` key, is +/// invalid. +/// +/// Ranges whose closedness differs per value use the companion +/// `arrow.variable_closedness_range` extension type instead. +/// +/// +#[derive(Debug, Clone, PartialEq)] +pub struct FixedClosednessRange(FixedClosednessRangeMetadata); + +impl FixedClosednessRange { + /// Returns a new `FixedClosednessRange` extension type. + pub fn new(closed: RangeClosed) -> Self { + Self(FixedClosednessRangeMetadata::new(closed)) + } + + /// Returns whether the interval endpoints are included or excluded. + pub fn closed(&self) -> RangeClosed { + self.0.closed() + } +} + +impl From for FixedClosednessRange { + fn from(value: FixedClosednessRangeMetadata) -> Self { + Self(value) + } +} + +/// Validates that `data_type` is an acceptable storage type for +/// `arrow.fixed_closedness_range`. +/// +/// Checks that the data type is a struct with exactly two fields named +/// "lower" and "upper", both sharing the same data type. Each bound may be +/// nullable or non-nullable independently; nullability is not required. +fn validate_storage(data_type: &DataType) -> Result<(), ArrowError> { + let fields: &Fields = match data_type { + DataType::Struct(fields) => fields, + other => { + return Err(ArrowError::InvalidArgumentError(format!( + "FixedClosednessRange data type mismatch, expected Struct, found {other}" + ))); + } + }; + + if fields.len() != 2 { + return Err(ArrowError::InvalidArgumentError(format!( + "FixedClosednessRange data type mismatch, expected Struct with 2 fields, found {} field(s)", + fields.len() + ))); + } + + let lower = &fields[0]; + let upper = &fields[1]; + + if lower.name() != "lower" { + return Err(ArrowError::InvalidArgumentError(format!( + "FixedClosednessRange data type mismatch, expected first field named \"lower\", found \"{}\"", + lower.name() + ))); + } + + if upper.name() != "upper" { + return Err(ArrowError::InvalidArgumentError(format!( + "FixedClosednessRange data type mismatch, expected second field named \"upper\", found \"{}\"", + upper.name() + ))); + } + + if lower.data_type() != upper.data_type() { + return Err(ArrowError::InvalidArgumentError(format!( + "FixedClosednessRange data type mismatch, \"lower\" and \"upper\" fields must have the same data type, found \"{}\" and \"{}\"", + lower.data_type(), + upper.data_type() + ))); + } + + Ok(()) +} + +impl ExtensionType for FixedClosednessRange { + const NAME: &'static str = "arrow.fixed_closedness_range"; + + type Metadata = FixedClosednessRangeMetadata; + + fn metadata(&self) -> &Self::Metadata { + &self.0 + } + + fn serialize_metadata(&self) -> Option { + Some(serde_json::to_string(self.metadata()).expect("metadata serialization")) + } + + fn deserialize_metadata(metadata: Option<&str>) -> Result { + metadata.map_or_else( + || { + Err(ArrowError::InvalidArgumentError( + "FixedClosednessRange extension type requires metadata".to_owned(), + )) + }, + |value| { + serde_json::from_str(value).map_err(|e| { + ArrowError::InvalidArgumentError(format!( + "FixedClosednessRange metadata deserialization failed: {e}" + )) + }) + }, + ) + } + + fn supports_data_type(&self, data_type: &DataType) -> Result<(), ArrowError> { + validate_storage(data_type) + } + + fn try_new(data_type: &DataType, metadata: Self::Metadata) -> Result { + validate_storage(data_type)?; + Ok(Self::from(metadata)) + } + + fn validate(data_type: &DataType, _metadata: Self::Metadata) -> Result<(), ArrowError> { + validate_storage(data_type) + } +} + +#[cfg(test)] +mod tests { + #[cfg(feature = "canonical_extension_types")] + use crate::extension::CanonicalExtensionType; + use crate::{ + Field, + extension::{EXTENSION_TYPE_METADATA_KEY, EXTENSION_TYPE_NAME_KEY}, + }; + + use super::*; + + fn make_range_struct(value_type: DataType) -> DataType { + DataType::Struct( + [ + Field::new("lower", value_type.clone(), true), + Field::new("upper", value_type, true), + ] + .into_iter() + .collect(), + ) + } + + #[test] + fn valid() -> Result<(), ArrowError> { + let range = FixedClosednessRange::new(RangeClosed::Both); + let storage = make_range_struct(DataType::Int32); + let mut field = Field::new("", storage, false); + field.try_with_extension_type(range.clone())?; + assert_eq!(field.try_extension_type::()?, range); + #[cfg(feature = "canonical_extension_types")] + assert_eq!( + field.try_canonical_extension_type()?, + CanonicalExtensionType::FixedClosednessRange(range) + ); + Ok(()) + } + + #[test] + fn valid_string_value_type() -> Result<(), ArrowError> { + // T may be any orderable Arrow type, including a string type. + let range = FixedClosednessRange::new(RangeClosed::Left); + let storage = make_range_struct(DataType::Utf8); + let mut field = Field::new("", storage, false); + field.try_with_extension_type(range.clone())?; + assert_eq!(field.try_extension_type::()?, range); + Ok(()) + } + + #[test] + fn roundtrip_all_closed_values() -> Result<(), ArrowError> { + let storage = make_range_struct(DataType::Int32); + for closed in [ + RangeClosed::Left, + RangeClosed::Right, + RangeClosed::Both, + RangeClosed::Neither, + ] { + let range = FixedClosednessRange::new(closed); + let mut field = Field::new("", storage.clone(), false); + field.try_with_extension_type(range.clone())?; + let recovered = field.try_extension_type::()?; + assert_eq!(recovered.closed(), closed); + } + Ok(()) + } + + #[test] + #[should_panic(expected = "Extension type name missing")] + fn missing_name() { + let storage = make_range_struct(DataType::Int32); + let field = Field::new("", storage, false) + .with_metadata([(EXTENSION_TYPE_METADATA_KEY, r#"{"closed":"both"}"#)]); + field.extension_type::(); + } + + #[test] + #[should_panic(expected = "FixedClosednessRange extension type requires metadata")] + fn missing_metadata() { + let storage = make_range_struct(DataType::Int32); + let field = Field::new("", storage, false) + .with_metadata([(EXTENSION_TYPE_NAME_KEY, FixedClosednessRange::NAME)]); + field.extension_type::(); + } + + #[test] + #[should_panic(expected = "FixedClosednessRange metadata deserialization failed")] + fn invalid_metadata_bad_closed_string() { + let storage = make_range_struct(DataType::Int32); + let field = Field::new("", storage, false).with_metadata([ + (EXTENSION_TYPE_NAME_KEY, FixedClosednessRange::NAME), + (EXTENSION_TYPE_METADATA_KEY, r#"{"closed":"invalid"}"#), + ]); + field.extension_type::(); + } + + #[test] + #[should_panic(expected = "FixedClosednessRange metadata deserialization failed")] + fn invalid_metadata_missing_closed_key() { + let storage = make_range_struct(DataType::Int32); + let field = Field::new("", storage, false).with_metadata([ + (EXTENSION_TYPE_NAME_KEY, FixedClosednessRange::NAME), + (EXTENSION_TYPE_METADATA_KEY, r"{}"), + ]); + field.extension_type::(); + } + + #[test] + fn unknown_metadata_keys_are_ignored() -> Result<(), ArrowError> { + // Additional keys in the JSON object should be ignored to allow + // forward-compatible extensions. + let range = FixedClosednessRange::new(RangeClosed::Right); + let storage = make_range_struct(DataType::Int32); + let field = Field::new("", storage, false).with_metadata([ + (EXTENSION_TYPE_NAME_KEY, FixedClosednessRange::NAME), + ( + EXTENSION_TYPE_METADATA_KEY, + r#"{"closed":"right","extra":42}"#, + ), + ]); + assert_eq!(field.try_extension_type::()?, range); + Ok(()) + } + + #[test] + #[should_panic( + expected = "FixedClosednessRange data type mismatch, expected Struct, found Int32" + )] + fn invalid_storage_non_struct() { + let range = FixedClosednessRange::new(RangeClosed::Both); + let field = Field::new("", DataType::Int32, false); + field.with_extension_type(range); + } + + #[test] + #[should_panic( + expected = "FixedClosednessRange data type mismatch, expected Struct with 2 fields" + )] + fn invalid_storage_wrong_field_count() { + let range = FixedClosednessRange::new(RangeClosed::Both); + let storage = + DataType::Struct(std::iter::once(Field::new("lower", DataType::Int32, true)).collect()); + let field = Field::new("", storage, false); + field.with_extension_type(range); + } + + #[test] + #[should_panic( + expected = "FixedClosednessRange data type mismatch, expected first field named \"lower\"" + )] + fn invalid_storage_wrong_field_names() { + let range = FixedClosednessRange::new(RangeClosed::Both); + let storage = DataType::Struct( + [ + Field::new("start", DataType::Int32, true), + Field::new("end", DataType::Int32, true), + ] + .into_iter() + .collect(), + ); + let field = Field::new("", storage, false); + field.with_extension_type(range); + } + + #[test] + fn accepts_non_nullable_lower() -> Result<(), ArrowError> { + let range = FixedClosednessRange::new(RangeClosed::Both); + let storage = DataType::Struct( + [ + Field::new("lower", DataType::Int32, false), + Field::new("upper", DataType::Int32, true), + ] + .into_iter() + .collect(), + ); + let mut field = Field::new("", storage, false); + field.try_with_extension_type(range.clone())?; + assert_eq!(field.try_extension_type::()?, range); + Ok(()) + } + + #[test] + fn accepts_non_nullable_upper() -> Result<(), ArrowError> { + let range = FixedClosednessRange::new(RangeClosed::Both); + let storage = DataType::Struct( + [ + Field::new("lower", DataType::Int32, true), + Field::new("upper", DataType::Int32, false), + ] + .into_iter() + .collect(), + ); + let mut field = Field::new("", storage, false); + field.try_with_extension_type(range.clone())?; + assert_eq!(field.try_extension_type::()?, range); + Ok(()) + } + + #[test] + fn accepts_both_non_nullable() -> Result<(), ArrowError> { + let range = FixedClosednessRange::new(RangeClosed::Both); + let storage = DataType::Struct( + [ + Field::new("lower", DataType::Int32, false), + Field::new("upper", DataType::Int32, false), + ] + .into_iter() + .collect(), + ); + let mut field = Field::new("", storage, false); + field.try_with_extension_type(range.clone())?; + assert_eq!(field.try_extension_type::()?, range); + Ok(()) + } + + #[test] + #[should_panic( + expected = "FixedClosednessRange data type mismatch, \"lower\" and \"upper\" fields must have the same data type" + )] + fn invalid_storage_mismatched_types() { + let range = FixedClosednessRange::new(RangeClosed::Both); + let storage = DataType::Struct( + [ + Field::new("lower", DataType::Int32, true), + Field::new("upper", DataType::Int64, true), + ] + .into_iter() + .collect(), + ); + let field = Field::new("", storage, false); + field.with_extension_type(range); + } +} diff --git a/arrow-schema/src/extension/canonical/fixed_shape_tensor.rs b/arrow-schema/src/extension/canonical/fixed_shape_tensor.rs index d30df2608819..8e8f71b9229c 100644 --- a/arrow-schema/src/extension/canonical/fixed_shape_tensor.rs +++ b/arrow-schema/src/extension/canonical/fixed_shape_tensor.rs @@ -113,6 +113,23 @@ impl FixedShapeTensor { self.metadata.list_size() } + /// Returns the physical shape of the contained tensors. + /// + /// [`Self::dimensions`] and [`Self::list_size`] do not identify this + /// shape. `[2, 6]` and `[3, 4]` are both 2-dimensional and both have + /// 12 elements. + /// + /// ``` + /// # use arrow_schema::extension::FixedShapeTensor; + /// # use arrow_schema::DataType; + /// let tensor = FixedShapeTensor::try_new(DataType::Float32, [2, 6], None, None).unwrap(); + /// assert_eq!(tensor.shape(), &[2, 6]); + /// assert_eq!(tensor.list_size(), 12); + /// ``` + pub fn shape(&self) -> &[usize] { + self.metadata.shape() + } + /// Returns the number of dimensions in this fixed shape tensor. pub fn dimensions(&self) -> usize { self.metadata.dimensions() @@ -129,6 +146,30 @@ impl FixedShapeTensor { pub fn permutations(&self) -> Option<&[usize]> { self.metadata.permutations() } + + /// Returns the logical shape. + /// + /// Logical dimension `i` has the size of physical dimension + /// `permutations[i]`. Without a permutation, this is the physical shape. + /// + /// Physical shape `[100, 200, 500]` and permutation `[2, 0, 1]` have + /// logical shape `[500, 100, 200]`. + /// + /// ``` + /// # use arrow_schema::extension::FixedShapeTensor; + /// # use arrow_schema::DataType; + /// let tensor = FixedShapeTensor::try_new( + /// DataType::Float32, + /// [100, 200, 500], + /// None, + /// Some(vec![2, 0, 1]), + /// ) + /// .unwrap(); + /// assert_eq!(tensor.logical_shape(), vec![500, 100, 200]); + /// ``` + pub fn logical_shape(&self) -> Vec { + self.metadata.logical_shape() + } } /// Extension type metadata for [`FixedShapeTensor`]. @@ -334,6 +375,11 @@ impl FixedShapeTensorMetadata { }) } + /// Returns the physical shape of the contained tensors. + pub fn shape(&self) -> &[usize] { + &self.shape + } + /// Returns the product of all the elements in tensor shape. pub fn list_size(&self) -> usize { self.shape.iter().product() @@ -355,6 +401,16 @@ impl FixedShapeTensorMetadata { pub fn permutations(&self) -> Option<&[usize]> { self.permutations.as_ref().map(AsRef::as_ref) } + + /// Returns the logical shape. + /// + /// See [`FixedShapeTensor::logical_shape`]. + pub fn logical_shape(&self) -> Vec { + match &self.permutations { + Some(permutation) => permutation.iter().map(|&index| self.shape[index]).collect(), + None => self.shape.clone(), + } + } } impl ExtensionType for FixedShapeTensor { @@ -438,7 +494,7 @@ mod tests { use crate::extension::CanonicalExtensionType; use crate::{ Field, - extension::{EXTENSION_TYPE_METADATA_KEY, EXTENSION_TYPE_NAME_KEY}, + extension::{EXTENSION_TYPE_METADATA_KEY, EXTENSION_TYPE_NAME_KEY, ExtensionType}, }; use super::*; @@ -462,6 +518,8 @@ mod tests { field.try_extension_type::()?, fixed_shape_tensor ); + assert_eq!(fixed_shape_tensor.shape(), &[100, 200, 500]); + assert_eq!(fixed_shape_tensor.metadata().shape(), &[100, 200, 500]); #[cfg(feature = "canonical_extension_types")] assert_eq!( field.try_canonical_extension_type()?, @@ -470,6 +528,39 @@ mod tests { Ok(()) } + #[test] + fn logical_shape_follows_permutation() -> Result<(), ArrowError> { + let physical = FixedShapeTensor::try_new(DataType::Float32, [100, 200, 500], None, None)?; + assert_eq!(physical.logical_shape(), vec![100, 200, 500]); + assert_eq!(physical.metadata().logical_shape(), vec![100, 200, 500]); + + let permuted = FixedShapeTensor::try_new( + DataType::Float32, + [100, 200, 500], + Some(vec!["C".to_owned(), "H".to_owned(), "W".to_owned()]), + Some(vec![2, 0, 1]), + )?; + assert_eq!(permuted.logical_shape(), vec![500, 100, 200]); + assert_eq!(permuted.metadata().logical_shape(), vec![500, 100, 200]); + + let identity = + FixedShapeTensor::try_new(DataType::Float32, [2, 6], None, Some(vec![0, 1]))?; + assert_eq!(identity.logical_shape(), vec![2, 6]); + Ok(()) + } + + #[test] + fn shape_distinguishes_equal_sizes() -> Result<(), ArrowError> { + let wide = FixedShapeTensor::try_new(DataType::Float32, [2, 6], None, None)?; + let tall = FixedShapeTensor::try_new(DataType::Float32, [3, 4], None, None)?; + assert_eq!(wide.dimensions(), tall.dimensions()); + assert_eq!(wide.list_size(), tall.list_size()); + assert_eq!(wide.shape(), &[2, 6]); + assert_eq!(tall.shape(), &[3, 4]); + assert_ne!(wide.shape(), tall.shape()); + Ok(()) + } + #[test] #[should_panic(expected = "Extension type name missing")] fn missing_name() { diff --git a/arrow-schema/src/extension/canonical/mod.rs b/arrow-schema/src/extension/canonical/mod.rs index a68169c7015e..f62757315e37 100644 --- a/arrow-schema/src/extension/canonical/mod.rs +++ b/arrow-schema/src/extension/canonical/mod.rs @@ -27,6 +27,8 @@ mod bool8; pub use bool8::Bool8; +mod fixed_closedness_range; +pub use fixed_closedness_range::{FixedClosednessRange, FixedClosednessRangeMetadata, RangeClosed}; mod fixed_shape_tensor; pub use fixed_shape_tensor::{FixedShapeTensor, FixedShapeTensorMetadata}; mod json; @@ -37,6 +39,8 @@ mod timestamp_with_offset; pub use timestamp_with_offset::TimestampWithOffset; mod uuid; pub use uuid::Uuid; +mod variable_closedness_range; +pub use variable_closedness_range::VariableClosednessRange; mod variable_shape_tensor; pub use variable_shape_tensor::{VariableShapeTensor, VariableShapeTensorMetadata}; @@ -75,6 +79,16 @@ pub enum CanonicalExtensionType { /// Opaque(Opaque), + /// The extension type for `FixedClosednessRange`. + /// + /// + FixedClosednessRange(FixedClosednessRange), + + /// The extension type for `VariableClosednessRange`. + /// + /// + VariableClosednessRange(VariableClosednessRange), + /// The extension type for `Bool8`. /// /// @@ -103,6 +117,12 @@ impl TryFrom<&Field> for CanonicalExtensionType { Json::NAME => value.try_extension_type::().map(Into::into), Uuid::NAME => value.try_extension_type::().map(Into::into), Opaque::NAME => value.try_extension_type::().map(Into::into), + FixedClosednessRange::NAME => value + .try_extension_type::() + .map(Into::into), + VariableClosednessRange::NAME => value + .try_extension_type::() + .map(Into::into), Bool8::NAME => value.try_extension_type::().map(Into::into), TimestampWithOffset::NAME => value .try_extension_type::() @@ -153,6 +173,18 @@ impl From for CanonicalExtensionType { } } +impl From for CanonicalExtensionType { + fn from(value: FixedClosednessRange) -> Self { + CanonicalExtensionType::FixedClosednessRange(value) + } +} + +impl From for CanonicalExtensionType { + fn from(value: VariableClosednessRange) -> Self { + CanonicalExtensionType::VariableClosednessRange(value) + } +} + impl From for CanonicalExtensionType { fn from(value: Bool8) -> Self { CanonicalExtensionType::Bool8(value) diff --git a/arrow-schema/src/extension/canonical/variable_closedness_range.rs b/arrow-schema/src/extension/canonical/variable_closedness_range.rs new file mode 100644 index 000000000000..49a590aff87f --- /dev/null +++ b/arrow-schema/src/extension/canonical/variable_closedness_range.rs @@ -0,0 +1,472 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Variable closedness range +//! +//! + +use crate::{ArrowError, DataType, Fields, extension::ExtensionType}; + +/// The extension type for a bounded set (mathematical interval) whose bound +/// inclusivity is recorded per value rather than as a single type-level +/// parameter. +/// +/// Extension name: `arrow.variable_closedness_range`. +/// +/// A range is a bounded set (mathematical interval) defined by a lower and +/// an upper bound over an orderable value type T. T may be any orderable +/// Arrow type, for example an integer, floating-point, decimal, date, time, +/// timestamp, duration, string or binary type. This specification defines +/// only the storage layout, not the order of any type. +/// +/// The storage type is a `Struct` with exactly four fields, in order: +/// - `lower` (type T): the lower bound. +/// - `upper` (type T): the upper bound. +/// - `lower_inc` (non-nullable `Boolean`): whether the lower bound is +/// inclusive for that value. +/// - `upper_inc` (non-nullable `Boolean`): whether the upper bound is +/// inclusive for that value. +/// +/// `lower` and `upper` must share the same data type T. Each of them may +/// independently be nullable or non-nullable: a nullable bound can hold +/// null to represent an unbounded (infinite) endpoint on that side, while a +/// non-nullable bound is always finite. A null bound is always treated as +/// exclusive, regardless of its `lower_inc` / `upper_inc` flag: only a null +/// bound means the range is unbounded on that side. `lower_inc` and +/// `upper_inc` are always non-nullable. +/// +/// This type has no type-level parameters: inclusivity is carried per value +/// in the `lower_inc` and `upper_inc` fields rather than fixed by the type. +/// The extension metadata therefore serializes to the empty JSON object +/// `{}`. +/// +/// Ranges that canonicalize to a single closedness shared by all values use +/// the companion `arrow.fixed_closedness_range` extension type instead. +/// +/// +#[derive(Debug, Default, Clone, Copy, PartialEq)] +pub struct VariableClosednessRange; + +/// Validates that `data_type` is an acceptable storage type for +/// `arrow.variable_closedness_range`. +/// +/// Checks that the data type is a struct with exactly four fields named +/// "lower", "upper", "lower_inc" and "upper_inc", in that order. "lower" and +/// "upper" must share the same data type; each may be nullable or +/// non-nullable independently. "lower_inc" and "upper_inc" must be +/// non-nullable `Boolean`. +fn validate_storage(data_type: &DataType) -> Result<(), ArrowError> { + let fields: &Fields = match data_type { + DataType::Struct(fields) => fields, + other => { + return Err(ArrowError::InvalidArgumentError(format!( + "VariableClosednessRange data type mismatch, expected Struct, found {other}" + ))); + } + }; + + if fields.len() != 4 { + return Err(ArrowError::InvalidArgumentError(format!( + "VariableClosednessRange data type mismatch, expected Struct with 4 fields, found {} field(s)", + fields.len() + ))); + } + + let lower = &fields[0]; + let upper = &fields[1]; + let lower_inc = &fields[2]; + let upper_inc = &fields[3]; + + if lower.name() != "lower" { + return Err(ArrowError::InvalidArgumentError(format!( + "VariableClosednessRange data type mismatch, expected first field named \"lower\", found \"{}\"", + lower.name() + ))); + } + + if upper.name() != "upper" { + return Err(ArrowError::InvalidArgumentError(format!( + "VariableClosednessRange data type mismatch, expected second field named \"upper\", found \"{}\"", + upper.name() + ))); + } + + if lower_inc.name() != "lower_inc" { + return Err(ArrowError::InvalidArgumentError(format!( + "VariableClosednessRange data type mismatch, expected third field named \"lower_inc\", found \"{}\"", + lower_inc.name() + ))); + } + + if upper_inc.name() != "upper_inc" { + return Err(ArrowError::InvalidArgumentError(format!( + "VariableClosednessRange data type mismatch, expected fourth field named \"upper_inc\", found \"{}\"", + upper_inc.name() + ))); + } + + if lower.data_type() != upper.data_type() { + return Err(ArrowError::InvalidArgumentError(format!( + "VariableClosednessRange data type mismatch, \"lower\" and \"upper\" fields must have the same data type, found \"{}\" and \"{}\"", + lower.data_type(), + upper.data_type() + ))); + } + + if lower_inc.data_type() != &DataType::Boolean || upper_inc.data_type() != &DataType::Boolean { + return Err(ArrowError::InvalidArgumentError(format!( + "VariableClosednessRange data type mismatch, \"lower_inc\" and \"upper_inc\" fields must be Boolean, found \"{}\" and \"{}\"", + lower_inc.data_type(), + upper_inc.data_type() + ))); + } + + if lower_inc.is_nullable() || upper_inc.is_nullable() { + return Err(ArrowError::InvalidArgumentError( + "VariableClosednessRange data type mismatch, \"lower_inc\" and \"upper_inc\" fields must be non-nullable".to_owned(), + )); + } + + Ok(()) +} + +/// Validates that `metadata`, if present, is either empty or a valid JSON +/// object. +/// +/// This type has no type-level parameters, so the metadata carries no +/// information. A missing metadata key and an empty string are both +/// accepted, as is any JSON object (unknown keys are ignored for +/// forward-compatibility). Anything else, such as malformed JSON or a JSON +/// value that is not an object, is rejected. +fn validate_metadata(metadata: Option<&str>) -> Result<(), ArrowError> { + match metadata { + None => Ok(()), + Some("") => Ok(()), + Some(value) => { + let parsed: serde_json::Value = serde_json::from_str(value).map_err(|e| { + ArrowError::InvalidArgumentError(format!( + "VariableClosednessRange metadata deserialization failed: {e}" + )) + })?; + if parsed.is_object() { + Ok(()) + } else { + Err(ArrowError::InvalidArgumentError(format!( + "VariableClosednessRange metadata must be a JSON object, found \"{value}\"" + ))) + } + } + } +} + +impl ExtensionType for VariableClosednessRange { + const NAME: &'static str = "arrow.variable_closedness_range"; + + type Metadata = (); + + fn metadata(&self) -> &Self::Metadata { + &() + } + + fn serialize_metadata(&self) -> Option { + // There are no type-level parameters, but the metadata is still + // emitted explicitly as the empty JSON object for forward-compat. + Some("{}".to_owned()) + } + + fn deserialize_metadata(metadata: Option<&str>) -> Result { + validate_metadata(metadata) + } + + fn supports_data_type(&self, data_type: &DataType) -> Result<(), ArrowError> { + validate_storage(data_type) + } + + fn try_new(data_type: &DataType, _metadata: Self::Metadata) -> Result { + Self.supports_data_type(data_type).map(|()| Self) + } + + fn validate(data_type: &DataType, _metadata: Self::Metadata) -> Result<(), ArrowError> { + Self.supports_data_type(data_type) + } +} + +#[cfg(test)] +mod tests { + #[cfg(feature = "canonical_extension_types")] + use crate::extension::CanonicalExtensionType; + use crate::{ + Field, + extension::{EXTENSION_TYPE_METADATA_KEY, EXTENSION_TYPE_NAME_KEY}, + }; + + use super::*; + + fn make_range_struct(value_type: DataType) -> DataType { + DataType::Struct( + [ + Field::new("lower", value_type.clone(), true), + Field::new("upper", value_type, true), + Field::new("lower_inc", DataType::Boolean, false), + Field::new("upper_inc", DataType::Boolean, false), + ] + .into_iter() + .collect(), + ) + } + + #[test] + fn valid() -> Result<(), ArrowError> { + let storage = make_range_struct(DataType::Int32); + let mut field = Field::new("", storage, false); + field.try_with_extension_type(VariableClosednessRange)?; + field.try_extension_type::()?; + #[cfg(feature = "canonical_extension_types")] + assert_eq!( + field.try_canonical_extension_type()?, + CanonicalExtensionType::VariableClosednessRange(VariableClosednessRange) + ); + Ok(()) + } + + #[test] + fn valid_string_value_type() -> Result<(), ArrowError> { + // T may be any orderable Arrow type, including a string type. + let storage = make_range_struct(DataType::Utf8); + let mut field = Field::new("", storage, false); + field.try_with_extension_type(VariableClosednessRange)?; + field.try_extension_type::()?; + Ok(()) + } + + #[test] + fn metadata_serializes_to_empty_object() -> Result<(), ArrowError> { + let storage = make_range_struct(DataType::Int32); + let mut field = Field::new("", storage, false); + field.try_with_extension_type(VariableClosednessRange)?; + assert_eq!( + field.metadata().get(EXTENSION_TYPE_METADATA_KEY), + Some(&"{}".to_owned()) + ); + Ok(()) + } + + #[test] + #[should_panic(expected = "Extension type name missing")] + fn missing_name() { + let storage = make_range_struct(DataType::Int32); + let field = + Field::new("", storage, false).with_metadata([(EXTENSION_TYPE_METADATA_KEY, "{}")]); + field.extension_type::(); + } + + #[test] + fn missing_metadata_is_accepted() -> Result<(), ArrowError> { + // Unlike FixedClosednessRange, this type has no type-level + // parameters, so a missing metadata key is equivalent to `{}`. + let storage = make_range_struct(DataType::Int32); + let field = Field::new("", storage, false) + .with_metadata([(EXTENSION_TYPE_NAME_KEY, VariableClosednessRange::NAME)]); + field.try_extension_type::()?; + Ok(()) + } + + #[test] + fn empty_and_object_and_unknown_keys_metadata_are_accepted() -> Result<(), ArrowError> { + let storage = make_range_struct(DataType::Int32); + for serialized in ["", "{}", r#"{"extra":42}"#] { + let field = Field::new("", storage.clone(), false).with_metadata([ + (EXTENSION_TYPE_NAME_KEY, VariableClosednessRange::NAME), + (EXTENSION_TYPE_METADATA_KEY, serialized), + ]); + field.try_extension_type::()?; + } + Ok(()) + } + + #[test] + #[should_panic(expected = "VariableClosednessRange metadata deserialization failed")] + fn invalid_metadata_malformed_json() { + let storage = make_range_struct(DataType::Int32); + let field = Field::new("", storage, false).with_metadata([ + (EXTENSION_TYPE_NAME_KEY, VariableClosednessRange::NAME), + (EXTENSION_TYPE_METADATA_KEY, "{"), + ]); + field.extension_type::(); + } + + #[test] + #[should_panic(expected = "VariableClosednessRange metadata must be a JSON object")] + fn invalid_metadata_not_an_object() { + let storage = make_range_struct(DataType::Int32); + let field = Field::new("", storage, false).with_metadata([ + (EXTENSION_TYPE_NAME_KEY, VariableClosednessRange::NAME), + (EXTENSION_TYPE_METADATA_KEY, "[]"), + ]); + field.extension_type::(); + } + + #[test] + #[should_panic( + expected = "VariableClosednessRange data type mismatch, expected Struct, found Int32" + )] + fn invalid_storage_non_struct() { + Field::new("", DataType::Int32, false).with_extension_type(VariableClosednessRange); + } + + #[test] + #[should_panic( + expected = "VariableClosednessRange data type mismatch, expected Struct with 4 fields" + )] + fn invalid_storage_wrong_field_count() { + let storage = DataType::Struct( + [ + Field::new("lower", DataType::Int32, true), + Field::new("upper", DataType::Int32, true), + ] + .into_iter() + .collect(), + ); + Field::new("", storage, false).with_extension_type(VariableClosednessRange); + } + + #[test] + #[should_panic( + expected = "VariableClosednessRange data type mismatch, expected first field named \"lower\"" + )] + fn invalid_storage_wrong_field_names() { + let storage = DataType::Struct( + [ + Field::new("start", DataType::Int32, true), + Field::new("upper", DataType::Int32, true), + Field::new("lower_inc", DataType::Boolean, false), + Field::new("upper_inc", DataType::Boolean, false), + ] + .into_iter() + .collect(), + ); + Field::new("", storage, false).with_extension_type(VariableClosednessRange); + } + + #[test] + #[should_panic( + expected = "VariableClosednessRange data type mismatch, expected third field named \"lower_inc\"" + )] + fn invalid_storage_wrong_inc_field_order() { + let storage = DataType::Struct( + [ + Field::new("lower", DataType::Int32, true), + Field::new("upper", DataType::Int32, true), + Field::new("upper_inc", DataType::Boolean, false), + Field::new("lower_inc", DataType::Boolean, false), + ] + .into_iter() + .collect(), + ); + Field::new("", storage, false).with_extension_type(VariableClosednessRange); + } + + #[test] + #[should_panic( + expected = "VariableClosednessRange data type mismatch, \"lower\" and \"upper\" fields must have the same data type" + )] + fn invalid_storage_mismatched_bound_types() { + let storage = DataType::Struct( + [ + Field::new("lower", DataType::Int32, true), + Field::new("upper", DataType::Int64, true), + Field::new("lower_inc", DataType::Boolean, false), + Field::new("upper_inc", DataType::Boolean, false), + ] + .into_iter() + .collect(), + ); + Field::new("", storage, false).with_extension_type(VariableClosednessRange); + } + + #[test] + #[should_panic( + expected = "VariableClosednessRange data type mismatch, \"lower_inc\" and \"upper_inc\" fields must be Boolean" + )] + fn invalid_storage_non_boolean_flags() { + let storage = DataType::Struct( + [ + Field::new("lower", DataType::Int32, true), + Field::new("upper", DataType::Int32, true), + Field::new("lower_inc", DataType::Int8, false), + Field::new("upper_inc", DataType::Int8, false), + ] + .into_iter() + .collect(), + ); + Field::new("", storage, false).with_extension_type(VariableClosednessRange); + } + + #[test] + #[should_panic( + expected = "VariableClosednessRange data type mismatch, \"lower_inc\" and \"upper_inc\" fields must be non-nullable" + )] + fn invalid_storage_nullable_flags() { + let storage = DataType::Struct( + [ + Field::new("lower", DataType::Int32, true), + Field::new("upper", DataType::Int32, true), + Field::new("lower_inc", DataType::Boolean, true), + Field::new("upper_inc", DataType::Boolean, true), + ] + .into_iter() + .collect(), + ); + Field::new("", storage, false).with_extension_type(VariableClosednessRange); + } + + #[test] + fn accepts_non_nullable_bounds() -> Result<(), ArrowError> { + let storage = DataType::Struct( + [ + Field::new("lower", DataType::Int32, false), + Field::new("upper", DataType::Int32, false), + Field::new("lower_inc", DataType::Boolean, false), + Field::new("upper_inc", DataType::Boolean, false), + ] + .into_iter() + .collect(), + ); + let mut field = Field::new("", storage, false); + field.try_with_extension_type(VariableClosednessRange)?; + field.try_extension_type::()?; + Ok(()) + } + + #[test] + fn accepts_asymmetric_bound_nullability() -> Result<(), ArrowError> { + let storage = DataType::Struct( + [ + Field::new("lower", DataType::Int32, true), + Field::new("upper", DataType::Int32, false), + Field::new("lower_inc", DataType::Boolean, false), + Field::new("upper_inc", DataType::Boolean, false), + ] + .into_iter() + .collect(), + ); + let mut field = Field::new("", storage, false); + field.try_with_extension_type(VariableClosednessRange)?; + field.try_extension_type::()?; + Ok(()) + } +} diff --git a/arrow-select/Cargo.toml b/arrow-select/Cargo.toml index 52ee50075d06..133f2a61e6cf 100644 --- a/arrow-select/Cargo.toml +++ b/arrow-select/Cargo.toml @@ -19,7 +19,6 @@ name = "arrow-select" version = { workspace = true } description = "Selection kernels for arrow arrays" -homepage = { workspace = true } repository = { workspace = true } authors = { workspace = true } license = { workspace = true } diff --git a/arrow-select/benches/filter_bits.rs b/arrow-select/benches/filter_bits.rs index 90e70b1e7b75..ccce30027c3b 100644 --- a/arrow-select/benches/filter_bits.rs +++ b/arrow-select/benches/filter_bits.rs @@ -39,6 +39,31 @@ fn create_boolean_array(size: usize, true_density: f64, rng: &mut StdRng) -> Boo .collect() } +/// A mask whose set bits come in runs averaging `run` bits, with the unset +/// runs sized to give `true_density`: a clustered or sorted column under a +/// range predicate, rather than independent rows +fn create_clustered_array( + size: usize, + true_density: f64, + run: f64, + rng: &mut StdRng, +) -> BooleanArray { + let leave_set = 1.0 / run; + let leave_unset = true_density / (run * (1.0 - true_density)); + let mut set = rng.random_bool(true_density); + (0..size) + .map(|_| { + let bit = set; + set = if set { + !rng.random_bool(leave_set) + } else { + rng.random_bool(leave_unset) + }; + Some(bit) + }) + .collect() +} + fn bench_filter_bits(predicate: &FilterPredicate, array: &BooleanArray) { hint::black_box(predicate.filter(array).unwrap()); } @@ -86,5 +111,60 @@ fn add_benchmark(c: &mut Criterion) { } } -criterion_group!(benches, add_benchmark); +/// `filter_bits` over 512 batches of 8K rows, each with its own mask, as a +/// query filters a stream of record batches. A recent core learns the +/// per-word branches of masks repeated from one iteration to the next, up to +/// at least 16K words, which makes random masks look faster than a real +/// filter; 512 batches are 64K words, past that. Lazy strategies only, which +/// compress word by word +fn add_batched_benchmark(c: &mut Criterion) { + const BATCHES: usize = 512; + const ROWS: usize = 8192; + let mut rng = StdRng::seed_from_u64(43); + + let data: Vec = (0..BATCHES) + .map(|_| create_boolean_array(ROWS, 0.5, &mut rng)) + .collect(); + + let mut cases = vec![]; + for (label, true_density) in [ + ("1/1024", 1.0 / 1024.0), + ("1/256", 1.0 / 256.0), + ("1/64", 1.0 / 64.0), + ("1/16", 1.0 / 16.0), + ("1/4", 0.25), + ("1/2", 0.5), + ("3/4", 0.75), + ("15/16", 15.0 / 16.0), + ] { + let masks: Vec<_> = (0..BATCHES) + .map(|_| create_boolean_array(ROWS, true_density, &mut rng)) + .collect(); + cases.push((format!("random, kept {label}"), masks)); + } + for run in [64, 512] { + for (label, true_density) in [("1/8", 0.125), ("1/2", 0.5), ("7/8", 0.875)] { + let masks: Vec<_> = (0..BATCHES) + .map(|_| create_clustered_array(ROWS, true_density, run as f64, &mut rng)) + .collect(); + cases.push((format!("runs of {run}, kept {label}"), masks)); + } + } + + for (label, masks) in &cases { + let predicates: Vec<_> = masks + .iter() + .map(|m| FilterBuilder::new(m).build()) + .collect(); + c.bench_function(&format!("filter_bits batches ({label})"), |b| { + b.iter(|| { + for (predicate, array) in predicates.iter().zip(&data) { + bench_filter_bits(predicate, array); + } + }) + }); + } +} + +criterion_group!(benches, add_benchmark, add_batched_benchmark); criterion_main!(benches); diff --git a/arrow-select/src/coalesce.rs b/arrow-select/src/coalesce.rs index 929526ef4043..863d39182965 100644 --- a/arrow-select/src/coalesce.rs +++ b/arrow-select/src/coalesce.rs @@ -238,6 +238,10 @@ impl BatchCoalescer { /// This is semantically equivalent of calling [`Self::push_batch`] /// with the results from [`crate::filter::filter_record_batch`] /// + /// If the number of rows that `filter` selects is already known, + /// [`Self::push_batch_with_filter_builder`] with + /// [`FilterBuilder::with_count`] avoids counting them again. + /// /// # Example /// ``` /// # use arrow_array::{record_batch, BooleanArray}; @@ -262,7 +266,45 @@ impl BatchCoalescer { batch: RecordBatch, filter: &BooleanArray, ) -> Result<(), ArrowError> { - self.push_batch_with_filtered_columns(batch, filter) + self.push_batch_with_filter_builder(batch, FilterBuilder::new(filter)) + } + + /// Push a batch into the Coalescer after applying the filter described by + /// `filter_builder`. + /// + /// This is [`Self::push_batch_with_filter`] for a [`FilterBuilder`] the + /// caller has already created. For example, callers that already know how + /// many rows the filter selects can provide that number with + /// [`FilterBuilder::with_count`] so that it is not counted again. + /// + /// Callers do not need to call [`FilterBuilder::optimize`]: like + /// [`Self::push_batch_with_filter`], this optimizes the filter when `batch` + /// has more than one column, or one column for which + /// [`FilterBuilder::is_optimize_beneficial`] returns true. A filter the + /// caller already optimized stays optimized. + /// + /// # Example + /// ``` + /// # use arrow_array::{record_batch, BooleanArray}; + /// # use arrow_select::coalesce::BatchCoalescer; + /// # use arrow_select::filter::FilterBuilder; + /// let batch = record_batch!(("a", Int32, [1, 2, 3])).unwrap(); + /// let filter = BooleanArray::from(vec![true, false, true]); + /// // SAFETY: the filter selects two rows + /// let filter_builder = unsafe { FilterBuilder::new(&filter).with_count(2) }; + /// let mut coalescer = BatchCoalescer::new(batch.schema(), 1000); + /// coalescer.push_batch_with_filter_builder(batch, filter_builder).unwrap(); + /// coalescer.finish_buffered_batch().unwrap(); + /// let expected_batch = record_batch!(("a", Int32, [1, 3])).unwrap(); + /// assert_eq!(coalescer.next_completed_batch().unwrap(), expected_batch); + /// ``` + pub fn push_batch_with_filter_builder( + &mut self, + batch: RecordBatch, + filter_builder: FilterBuilder, + ) -> Result<(), ArrowError> { + let predicate = Self::filter_predicate_for_batch(&batch, filter_builder); + self.push_batch_with_filtered_columns(batch, &predicate) } /// Push a batch into the Coalescer after applying a set of indices @@ -606,10 +648,8 @@ impl BatchCoalescer { impl BatchCoalescer { fn filter_predicate_for_batch( batch: &RecordBatch, - filter: &BooleanArray, - selected_count: usize, + mut filter_builder: FilterBuilder, ) -> FilterPredicate { - let mut filter_builder = FilterBuilder::new_with_count(filter, selected_count); if batch.num_columns() > 1 || (batch.num_columns() > 0 && FilterBuilder::is_optimize_beneficial(batch.schema_ref().field(0).data_type())) @@ -622,9 +662,10 @@ impl BatchCoalescer { fn push_batch_with_filtered_columns( &mut self, batch: RecordBatch, - filter: &BooleanArray, + predicate: &FilterPredicate, ) -> Result<(), ArrowError> { - let filter_len = filter.len(); + let filter_len = predicate.filter_len(); + let selected_count = predicate.count(); let batch_num_rows = batch.num_rows(); let batch_num_columns = batch.num_columns(); @@ -634,7 +675,6 @@ impl BatchCoalescer { ))); } - let selected_count = filter.true_count(); if selected_count == 0 { return Ok(()); } @@ -662,16 +702,14 @@ impl BatchCoalescer { if should_materialize_filter { // Use materialized filtering when sparse per-column copying is unavailable. - let predicate = Self::filter_predicate_for_batch(&batch, filter, selected_count); let filtered_batch = predicate.filter_record_batch(&batch)?; return self.push_batch(filtered_batch); } - let predicate = Self::filter_predicate_for_batch(&batch, filter, selected_count); let (_schema, arrays, _num_rows) = batch.into_parts(); for (in_progress, array) in self.in_progress_arrays.iter_mut().zip(arrays) { - in_progress.copy_rows_by_filter_from(array, &predicate)?; + in_progress.copy_rows_by_filter_from(array, predicate)?; } self.buffered_rows += selected_count; @@ -788,7 +826,7 @@ mod tests { use arrow_array::types::Int32Type; use arrow_array::{ BinaryViewArray, Int32Array, Int64Array, RecordBatchOptions, StringArray, StringViewArray, - TimestampNanosecondArray, UInt32Array, UInt64Array, make_array, + TimestampNanosecondArray, UInt32Array, UInt64Array, make_array, record_batch, }; use arrow_buffer::BooleanBufferBuilder; use arrow_schema::{DataType, Field, Schema}; @@ -1025,6 +1063,29 @@ mod tests { .run(); } + #[test] + fn test_push_batch_with_filter_builder() { + let batch = record_batch!(("a", Int32, [1, 2, 3, 4])).unwrap(); + // The null is not selected, so the filter selects two rows. + let filter = BooleanArray::from(vec![Some(true), None, Some(false), Some(true)]); + let expected = record_batch!(("a", Int32, [1, 4])).unwrap(); + + // SAFETY: the count matches the filter. + let with_count = unsafe { FilterBuilder::new(&filter).with_count(2) }; + // The coalescer does not optimize a single primitive column itself, so + // this covers a predicate that only the caller optimized. + let optimized = FilterBuilder::new(&filter).optimize(); + for filter_builder in [with_count, optimized] { + let mut coalescer = BatchCoalescer::new(batch.schema(), 10); + coalescer + .push_batch_with_filter_builder(batch.clone(), filter_builder) + .unwrap(); + coalescer.finish_buffered_batch().unwrap(); + assert_eq!(coalescer.next_completed_batch().unwrap(), expected); + assert!(coalescer.next_completed_batch().is_none()); + } + } + #[test] fn test_coalesce_non_null() { Test::new("coalesce_non_null") diff --git a/arrow-select/src/concat.rs b/arrow-select/src/concat.rs index 72dd239bbb17..51b808285339 100644 --- a/arrow-select/src/concat.rs +++ b/arrow-select/src/concat.rs @@ -38,9 +38,9 @@ use arrow_array::cast::AsArray; use arrow_array::types::*; use arrow_array::*; use arrow_buffer::{ - ArrowNativeType, BooleanBufferBuilder, MutableBuffer, NullBuffer, OffsetBuffer, ScalarBuffer, + ArrowNativeType, BooleanBufferBuilder, MutableBuffer, NullBuffer, OffsetBuffer, RunEndBuffer, + ScalarBuffer, }; -use arrow_data::ArrayDataBuilder; use arrow_data::transform::{Capacities, MutableArrayData}; use arrow_schema::{ArrowError, DataType, FieldRef, Fields, SchemaRef}; use std::{collections::HashSet, ops::Add, sync::Arc}; @@ -430,6 +430,12 @@ where return Ok(new_empty_array(arrays[0].data_type())); } + // Reject lengths that do not fit in `R` before any `R::Native` arithmetic can wrap. + let total_len: usize = run_arrays.iter().map(|r| r.len()).sum(); + if R::Native::from_usize(total_len).is_none() { + return Err(ArrowError::RunEndIndexOverflowError); + } + // The run ends need to be adjusted by the sum of the lengths of the previous arrays. let needed_run_end_adjustments = std::iter::once(R::default_value()) .chain( @@ -442,9 +448,6 @@ where ) .collect::>(); - // This works out nicely to be the total (logical) length of the resulting array. - let total_len = needed_run_end_adjustments.last().unwrap().as_usize(); - let run_ends_array = PrimitiveArray::::from_iter_values(run_arrays.iter().enumerate().flat_map( move |(i, run_array)| { @@ -463,15 +466,15 @@ where let all_values = concat(&values_slices.iter().map(|x| x.as_ref()).collect::>())?; - let builder = ArrayDataBuilder::new(run_arrays[0].data_type().clone()) - .len(total_len) - .child_data(vec![run_ends_array.into_data(), all_values.into_data()]); + let data_type = run_arrays[0].data_type().clone(); + let (_, run_ends_values, _) = run_ends_array.into_parts(); - // `build_unchecked` is used to avoid recursive validation of child arrays. - let array_data = unsafe { builder.build_unchecked() }; - array_data.validate_data()?; - - Ok(Arc::>::new(array_data.into())) + // Safety: inputs are valid RunArrays; adjusted run ends are strictly increasing + // and end at `total_len`. Physical length matches `all_values`. + let run_ends = unsafe { RunEndBuffer::new_unchecked(run_ends_values, 0, total_len) }; + Ok(Arc::new(unsafe { + RunArray::::new_unchecked(data_type, run_ends, all_values) + })) } macro_rules! dict_helper { @@ -1889,6 +1892,24 @@ mod tests { assert_eq!(&[10, 20, 30, 40], values.values()); } + #[test] + fn test_concat_run_array_length_overflows_run_end_type() { + // 20_000 + 20_000 exceeds i16::MAX. + let array1 = RunArray::::try_new( + &Int16Array::from(vec![20_000]), + &Int16Array::from(vec![1]), + ) + .unwrap(); + let array2 = RunArray::::try_new( + &Int16Array::from(vec![20_000]), + &Int16Array::from(vec![2]), + ) + .unwrap(); + + let err = concat(&[&array1, &array2]).unwrap_err(); + assert!(matches!(err, ArrowError::RunEndIndexOverflowError)); + } + #[test] fn test_concat_sliced_run_array() { // Slicing away first run in both arrays diff --git a/arrow-select/src/filter.rs b/arrow-select/src/filter.rs index ed9b0da63748..9baccf745c4f 100644 --- a/arrow-select/src/filter.rs +++ b/arrow-select/src/filter.rs @@ -181,6 +181,9 @@ pub fn prep_null_mask_filter(filter: &BooleanArray) -> BooleanArray { /// In contrast to this function, it is then the responsibility of the caller /// to use [FilterBuilder::optimize] if appropriate. /// +/// If the number of rows that `predicate` selects is already known, +/// [`FilterBuilder::with_count`] avoids counting them again. +/// /// # See also /// * [`FilterBuilder`] for more control over the filtering process. /// * [`filter_record_batch`] to filter a [`RecordBatch`] @@ -223,6 +226,9 @@ pub fn filter(values: &dyn Array, predicate: &BooleanArray) -> Result, + optimize: bool, } impl FilterBuilder { /// Create a new [`FilterBuilder`] that can be used to construct a [`FilterPredicate`] pub fn new(filter: &BooleanArray) -> Self { - Self::new_with_count(filter, filter.true_count()) - } - - pub(crate) fn new_with_count(filter: &BooleanArray, count: usize) -> Self { let filter = match filter.null_count() { 0 => filter.clone(), _ => prep_null_mask_filter(filter), }; - let strategy = IterationStrategy::default_strategy(filter.len(), count); - Self { filter, - count, - strategy, + count: None, + optimize: false, } } - /// Compute an optimized representation of the provided `filter` mask that can be - /// applied to an array more quickly. + /// Set the number of rows that the filter selects, so that [`Self::build`] + /// does not have to count them. + /// + /// Callers that build a mask row by row, or derive it from a validity + /// buffer with a cached null count, often already hold this number. + /// + /// # Safety + /// + /// `count` must equal [`BooleanArray::true_count`] of the filter passed to + /// [`Self::new`]: the number of `true` values that are not null. + /// + /// # Example + /// ``` + /// # use arrow_array::{BooleanArray, Int32Array}; + /// # use arrow_select::filter::FilterBuilder; + /// let values = Int32Array::from(vec![1, 2, 3, 4]); + /// let mask = BooleanArray::from(vec![Some(true), None, Some(true), Some(false)]); + /// // The null is not selected, so the mask selects two rows. + /// // SAFETY: the count matches the mask. + /// let predicate = unsafe { FilterBuilder::new(&mask).with_count(2) }.build(); + /// assert_eq!(predicate.count(), 2); + /// let filtered = predicate.filter(&values).unwrap(); + /// assert_eq!(filtered.as_ref(), &Int32Array::from(vec![1, 3])); + /// ``` + pub unsafe fn with_count(mut self, count: usize) -> Self { + debug_assert_eq!( + count, + self.filter.true_count(), + "count must match the number of rows the filter selects" + ); + self.count = Some(count); + self + } + + /// Compute an optimized representation of the provided `filter` mask in + /// [`Self::build`], so that it can be applied to an array more quickly. /// /// When filtering multiple arrays (e.g. a [`RecordBatch`] or a /// [`StructArray`] with multiple fields), optimizing the filter can provide @@ -284,17 +319,7 @@ impl FilterBuilder { /// than the original mask, so it is often faster to filter a single array, /// without filter optimization. pub fn optimize(mut self) -> Self { - match self.strategy { - IterationStrategy::SlicesIterator => { - let slices = SlicesIterator::new(&self.filter).collect(); - self.strategy = IterationStrategy::Slices(slices) - } - IterationStrategy::IndexIterator => { - let indices = IndexIterator::new(&self.filter, self.count).collect(); - self.strategy = IterationStrategy::Indices(indices) - } - _ => {} - } + self.optimize = true; self } @@ -316,10 +341,26 @@ impl FilterBuilder { /// Construct the final `FilterPredicate` pub fn build(self) -> FilterPredicate { + let count = self.count.unwrap_or_else(|| self.filter.true_count()); + let mut strategy = IterationStrategy::default_strategy(self.filter.len(), count); + if self.optimize { + match strategy { + IterationStrategy::SlicesIterator => { + let slices = SlicesIterator::new(&self.filter).collect(); + strategy = IterationStrategy::Slices(slices) + } + IterationStrategy::IndexIterator => { + let indices = IndexIterator::new(&self.filter, count).collect(); + strategy = IterationStrategy::Indices(indices) + } + _ => {} + } + } + FilterPredicate { filter: self.filter, - count: self.count, - strategy: self.strategy, + count, + strategy, } } } @@ -483,6 +524,11 @@ impl FilterPredicate { self.count } + /// Length of the filter mask, including rows that are not selected + pub(crate) fn filter_len(&self) -> usize { + self.filter.len() + } + /// Return a [`FilterSelection`] for iterating over the rows selected by /// this [`FilterPredicate`]. pub(crate) fn selection(&self) -> FilterSelection<'_> { diff --git a/arrow-select/src/interleave.rs b/arrow-select/src/interleave.rs index 54f4412c5bac..74e29defe835 100644 --- a/arrow-select/src/interleave.rs +++ b/arrow-select/src/interleave.rs @@ -26,10 +26,13 @@ use arrow_array::types::*; use arrow_array::*; use arrow_buffer::bit_mask::set_bits; use arrow_buffer::bit_util; -use arrow_buffer::{ArrowNativeType, BooleanBuffer, MutableBuffer, NullBuffer, OffsetBuffer}; +use arrow_buffer::{ + ArrowNativeType, BooleanBuffer, Buffer, MutableBuffer, NullBuffer, OffsetBuffer, +}; use arrow_data::ByteView; use arrow_data::transform::MutableArrayData; use arrow_schema::{ArrowError, DataType, FieldRef, Fields}; +use std::collections::BTreeMap; use std::sync::Arc; macro_rules! primitive_helper { @@ -314,8 +317,17 @@ fn interleave_views( offsets.push(total_buffers); } + // Marks a buffer in `buffer_to_new_index` that has not yet been referenced. + // + // The view's buffer index is a signed 32-bit integer in the Arrow specification, + // so `u32::MAX` can never be a valid buffer index + const UNASSIGNED: u32 = u32::MAX; + // contains the mapping from old buffer index to new buffer index - let mut buffer_to_new_index = vec![None; total_buffers]; + let mut buffer_to_new_index = vec![UNASSIGNED; total_buffers]; + + // Contains the index in `buffers` of each buffer already emitted, keyed by identity + let mut seen_buffers = BTreeMap::new(); let views: Vec = indices .iter() @@ -329,12 +341,15 @@ fn interleave_views( // value is big enough to be in a variadic buffer let view = ByteView::from(*view); let buffer_to_new_idx = offsets[*array_idx] + view.buffer_index as usize; - let new_buffer_idx: u32 = - *buffer_to_new_index[buffer_to_new_idx].get_or_insert_with(|| { - buffers.push(array.data_buffers()[view.buffer_index as usize].clone()); - (buffers.len() - 1) as u32 - }); - view.with_buffer_index(new_buffer_idx).as_u128() + let new_buffer_idx = &mut buffer_to_new_index[buffer_to_new_idx]; + if *new_buffer_idx == UNASSIGNED { + *new_buffer_idx = push_buffer( + &mut seen_buffers, + &mut buffers, + &array.data_buffers()[view.buffer_index as usize], + ); + } + view.with_buffer_index(*new_buffer_idx).as_u128() }) .collect(); @@ -344,6 +359,24 @@ fn interleave_views( Ok(Arc::new(array)) } +/// Returns the index of `buffer` in `buffers`, adding it if not already present. +/// +/// Multiple input arrays, e.g. slices of the same array, may reference the same +/// buffer, so this checks by identity to ensure it is only emitted once +#[inline(never)] +fn push_buffer( + seen: &mut BTreeMap<(*const u8, usize), u32>, + buffers: &mut Vec, + buffer: &Buffer, +) -> u32 { + *seen + .entry((buffer.as_ptr(), buffer.len())) + .or_insert_with(|| { + buffers.push(buffer.clone()); + (buffers.len() - 1) as u32 + }) +} + fn interleave_struct( fields: &Fields, values: &[&dyn Array], @@ -1674,6 +1707,41 @@ mod tests { ); } + #[test] + fn test_interleave_views_shared_buffers() { + let long = |i: usize| format!("long_string_not_inlined_{i}"); + let base = StringViewArray::from_iter_values((0..10).map(long)); + assert_eq!(base.data_buffers().len(), 1); + + // Slices share the same buffers + let a = base.slice(0, 5); + let b = base.slice(5, 5); + // A distinct buffer list containing the same underlying buffer + let c = StringViewArray::try_new(base.views().clone(), base.data_buffers().to_vec(), None) + .unwrap(); + + let indices = &[(0, 1), (1, 2), (2, 3), (0, 4), (1, 0), (2, 9)]; + let values = interleave(&[&a, &b, &c], indices).unwrap(); + let result = values.as_string_view(); + assert_eq!(result.data_buffers().len(), 1); + + let expected: Vec<_> = [1, 7, 3, 4, 5, 9].into_iter().map(long).collect(); + let actual: Vec<_> = result.iter().map(|x| x.unwrap().to_string()).collect(); + assert_eq!(actual, expected); + + // Many slices, exceeding the linear scan threshold + let slices: Vec<_> = (0..40).map(|i| base.slice(i % 10, 1)).collect(); + let arrays: Vec<&dyn Array> = slices.iter().map(|x| x as &dyn Array).collect(); + let indices: Vec<_> = (0..40).rev().map(|i| (i, 0)).collect(); + let values = interleave(&arrays, &indices).unwrap(); + let result = values.as_string_view(); + assert_eq!(result.data_buffers().len(), 1); + + let expected: Vec<_> = (0..40).rev().map(|i| long(i % 10)).collect(); + let actual: Vec<_> = result.iter().map(|x| x.unwrap().to_string()).collect(); + assert_eq!(actual, expected); + } + #[test] fn test_interleave_views_multiple_buffers() { let str1 = "very_long_string_from_first_buffer".as_bytes(); diff --git a/arrow-select/src/nullif.rs b/arrow-select/src/nullif.rs index aa09a28631ee..cd6d0ec4a2f7 100644 --- a/arrow-select/src/nullif.rs +++ b/arrow-select/src/nullif.rs @@ -578,31 +578,13 @@ mod tests { /// Returns a new BooleanArray with no null buffer fn remove_null_buffer(array: &BooleanArray) -> BooleanArray { - make_array( - array - .into_data() - .into_builder() - .nulls(None) - .build() - .unwrap(), - ) - .as_boolean() - .clone() + BooleanArray::new(array.values().clone(), None) } /// Returns a new BooleanArray with a null buffer where all values are valid fn remove_null_values(array: &BooleanArray) -> BooleanArray { let len = array.len(); let new_nulls = NullBuffer::from_iter(std::iter::repeat_n(true, len)); - make_array( - array - .into_data() - .into_builder() - .nulls(Some(new_nulls)) - .build() - .unwrap(), - ) - .as_boolean() - .clone() + BooleanArray::new(array.values().clone(), Some(new_nulls)) } } diff --git a/arrow-select/src/take.rs b/arrow-select/src/take.rs index 8c9c4ba3da0a..38d7945700cc 100644 --- a/arrow-select/src/take.rs +++ b/arrow-select/src/take.rs @@ -163,6 +163,33 @@ pub fn take_arrays( .collect() } +/// For each [ArrayRef] in the [`Vec`], take elements by index and create a new +/// [`Vec`] from those indices, without bounds checking. +/// +/// # Safety +/// +/// Caller must ensure all values in `indices` are valid indices into each array +/// (i.e. `< array.len()`). Out-of-bounds indices will cause undefined behavior. +/// +/// # Errors +/// +/// Returns an error if an index value cannot be cast to `usize`. +pub unsafe fn take_arrays_unchecked( + arrays: &[ArrayRef], + indices: &dyn Array, +) -> Result, ArrowError> { + downcast_integer_array!( + indices => { + let indices = indices.to_indices(); + arrays + .iter() + .map(|array| take_impl::<_, false>(array.as_ref(), &indices)) + .collect() + }, + d => Err(ArrowError::InvalidArgumentError(format!("Take only supported for integers, got {d:?}"))) + ) +} + /// Verifies that the non-null values of `indices` are all `< len` fn check_bounds( len: usize, @@ -209,7 +236,7 @@ where } #[inline(never)] -fn take_impl( +fn take_impl( values: &dyn Array, indices: &PrimitiveArray, ) -> Result { @@ -224,38 +251,38 @@ fn take_impl( return Ok(new_empty_array(values.data_type())); } downcast_primitive_array! { - values => Ok(Arc::new(take_primitive::<_, _, CHECKED>(values, indices)?)), + values => Ok(Arc::new(take_primitive::<_, _, VALIDATE_INDICES>(values, indices)?)), DataType::Boolean => { let values = values.as_any().downcast_ref::().unwrap(); - Ok(Arc::new(take_boolean::<_, CHECKED>(values, indices))) + Ok(Arc::new(take_boolean::<_, VALIDATE_INDICES>(values, indices))) } DataType::Utf8 => { - Ok(Arc::new(take_bytes::<_, _, CHECKED>(values.as_string::(), indices)?)) + Ok(Arc::new(take_bytes::<_, _, VALIDATE_INDICES>(values.as_string::(), indices)?)) } DataType::LargeUtf8 => { - Ok(Arc::new(take_bytes::<_, _, CHECKED>(values.as_string::(), indices)?)) + Ok(Arc::new(take_bytes::<_, _, VALIDATE_INDICES>(values.as_string::(), indices)?)) } DataType::Utf8View => { - Ok(Arc::new(take_byte_view::<_, _, CHECKED>(values.as_string_view(), indices)?)) + Ok(Arc::new(take_byte_view::<_, _, VALIDATE_INDICES>(values.as_string_view(), indices)?)) } DataType::List(_) => { - Ok(Arc::new(take_list::<_, Int32Type, CHECKED>(values.as_list(), indices)?)) + Ok(Arc::new(take_list::<_, Int32Type, VALIDATE_INDICES>(values.as_list(), indices)?)) } DataType::LargeList(_) => { - Ok(Arc::new(take_list::<_, Int64Type, CHECKED>(values.as_list(), indices)?)) + Ok(Arc::new(take_list::<_, Int64Type, VALIDATE_INDICES>(values.as_list(), indices)?)) } DataType::ListView(_) => { - Ok(Arc::new(take_list_view::<_, Int32Type, CHECKED>(values.as_list_view(), indices)?)) + Ok(Arc::new(take_list_view::<_, Int32Type, VALIDATE_INDICES>(values.as_list_view(), indices)?)) } DataType::LargeListView(_) => { - Ok(Arc::new(take_list_view::<_, Int64Type, CHECKED>(values.as_list_view(), indices)?)) + Ok(Arc::new(take_list_view::<_, Int64Type, VALIDATE_INDICES>(values.as_list_view(), indices)?)) } DataType::FixedSizeList(_, length) => { let values = values .as_any() .downcast_ref::() .unwrap(); - Ok(Arc::new(take_fixed_size_list::<_, CHECKED>( + Ok(Arc::new(take_fixed_size_list::<_, VALIDATE_INDICES>( values, indices, *length as u32, @@ -263,7 +290,7 @@ fn take_impl( } DataType::Map(field, ordered) => { let list_arr = ListArray::from(values.as_map().clone()); - let list_data = take_list::<_, Int32Type, CHECKED>(&list_arr, indices)?; + let list_data = take_list::<_, Int32Type, VALIDATE_INDICES>(&list_arr, indices)?; let (_, offsets, entries, nulls) = list_data.into_parts(); let entries = entries.as_struct().clone(); Ok(Arc::new(MapArray::try_new( @@ -279,7 +306,7 @@ fn take_impl( let arrays = array .columns() .iter() - .map(|a| take_impl::<_, CHECKED>(a.as_ref(), indices)) + .map(|a| take_impl::<_, VALIDATE_INDICES>(a.as_ref(), indices)) .collect::, _>>()?; let fields: Vec<(FieldRef, ArrayRef)> = fields.iter().cloned().zip(arrays).collect(); @@ -304,7 +331,7 @@ fn take_impl( } } DataType::Dictionary(_, _) => downcast_dictionary_array! { - values => Ok(Arc::new(take_dict::<_, _, CHECKED>(values, indices)?)), + values => Ok(Arc::new(take_dict::<_, _, VALIDATE_INDICES>(values, indices)?)), t => unimplemented!("Take not supported for dictionary type {:?}", t) } DataType::RunEndEncoded(_, _) => downcast_run_array! { @@ -312,20 +339,20 @@ fn take_impl( t => unimplemented!("Take not supported for run type {:?}", t) } DataType::Binary => { - Ok(Arc::new(take_bytes::<_, _, CHECKED>(values.as_binary::(), indices)?)) + Ok(Arc::new(take_bytes::<_, _, VALIDATE_INDICES>(values.as_binary::(), indices)?)) } DataType::LargeBinary => { - Ok(Arc::new(take_bytes::<_, _, CHECKED>(values.as_binary::(), indices)?)) + Ok(Arc::new(take_bytes::<_, _, VALIDATE_INDICES>(values.as_binary::(), indices)?)) } DataType::BinaryView => { - Ok(Arc::new(take_byte_view::<_, _, CHECKED>(values.as_binary_view(), indices)?)) + Ok(Arc::new(take_byte_view::<_, _, VALIDATE_INDICES>(values.as_binary_view(), indices)?)) } DataType::FixedSizeBinary(size) => { let values = values .as_any() .downcast_ref::() .unwrap(); - Ok(Arc::new(take_fixed_size_binary::<_, CHECKED>(values, indices, *size)?)) + Ok(Arc::new(take_fixed_size_binary::<_, VALIDATE_INDICES>(values, indices, *size)?)) } DataType::Null => { // Take applied to a null array produces a null array. @@ -345,7 +372,7 @@ fn take_impl( let children = fields .iter() .map(|(type_id, _)| { - take_sparse_union_child::<_, CHECKED>( + take_sparse_union_child::<_, VALIDATE_INDICES>( values.child(type_id), null_type_id == Some(type_id), indices, @@ -363,7 +390,7 @@ fn take_impl( // Keep index nulls so `take` of each child writes a null instead of // reading child offset 0 (the default `take_native` fills in). let offsets = >::try_new( - take_native(values.offsets().unwrap(), indices), + take_native::<_, _, VALIDATE_INDICES>(values.offsets().unwrap(), indices), indices.nulls().cloned(), )?; @@ -375,7 +402,7 @@ fn take_impl( let values = values.child(field_type_id); - take_impl::<_, CHECKED>(values, indices.as_primitive::()) + take_impl::<_, VALIDATE_INDICES>(values, indices.as_primitive::()) }) .collect::>()?; @@ -418,11 +445,11 @@ fn take_union_type_ids( indices: &PrimitiveArray, ) -> Result<(ScalarBuffer, Option), ArrowError> { if indices.null_count() == 0 { - return Ok((take_native(type_ids, indices), None)); + return Ok((take_native::<_, _, true>(type_ids, indices), None)); } let null_type_id = union_null_type_id(fields)?; - let taken_type_ids = take_native(type_ids, indices); + let taken_type_ids = take_native::<_, _, true>(type_ids, indices); let type_ids = indices .iter() .zip(&taken_type_ids) @@ -474,13 +501,13 @@ fn field_can_represent_take_null(field: &FieldRef) -> bool { /// ([`take_union_type_ids`]). Values in the other children at those positions /// are unspecified, so they are taken with dummy indices instead of introducing /// nulls that would contradict field metadata. -fn take_sparse_union_child( +fn take_sparse_union_child( values: &dyn Array, represent_nulls: bool, indices: &PrimitiveArray, ) -> Result { if represent_nulls || indices.null_count() == 0 { - return take_impl::<_, CHECKED>(values, indices); + return take_impl::<_, VALIDATE_INDICES>(values, indices); } if values.is_empty() { @@ -496,7 +523,7 @@ fn take_sparse_union_child( )); } - take_impl::<_, CHECKED>(values, &indices_without_nulls(indices)) + take_impl::<_, VALIDATE_INDICES>(values, &indices_without_nulls(indices)) } /// Replaces null take indices with `0` and drops the null bitmap. @@ -526,7 +553,7 @@ pub struct TakeOptions { /// 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( +fn take_primitive( values: &PrimitiveArray, indices: &PrimitiveArray, ) -> Result, ArrowError> @@ -534,19 +561,19 @@ where T: ArrowPrimitiveType, I: ArrowPrimitiveType, { - let values_buf = take_native(values.values(), indices); - let nulls = take_nulls::<_, CHECKED>(values.nulls(), indices); + let values_buf = take_native::<_, _, VALIDATE_INDICES>(values.values(), indices); + let nulls = take_nulls::<_, VALIDATE_INDICES>(values.nulls(), indices); Ok(PrimitiveArray::try_new(values_buf, nulls)?.with_data_type(values.data_type().clone())) } #[inline(never)] -fn take_nulls( +fn take_nulls( values: Option<&NullBuffer>, indices: &PrimitiveArray, ) -> Option { match values.filter(|n| n.null_count() > 0) { Some(n) => NullBuffer::from_unsliced_buffer( - take_bits::<_, CHECKED>(n.inner(), indices).into_inner(), + take_bits::<_, VALIDATE_INDICES>(n.inner(), indices).into_inner(), indices.len(), ), None => indices.nulls().cloned(), @@ -554,7 +581,7 @@ fn take_nulls( } #[inline(never)] -fn take_native( +fn take_native( values: &[T], indices: &PrimitiveArray, ) -> ScalarBuffer { @@ -572,11 +599,24 @@ fn take_native( }, }) .collect(), - None => indices - .values() - .iter() - .map(|index| values[index.as_usize()]) - .collect(), + None => { + if VALIDATE_INDICES { + indices + .values() + .iter() + .map(|index| values[index.as_usize()]) + .collect() + } else { + indices + .values() + .iter() + .map(|index| { + // SAFETY: caller guarantees all indices are in-bounds. + unsafe { *values.get_unchecked(index.as_usize()) } + }) + .collect() + } + } } } @@ -617,7 +657,7 @@ unsafe fn pack_bit(src: *const u8, bit_idx: usize, out_pos: usize) -> u8 { } #[inline(never)] -fn take_bits( +fn take_bits( values: &BooleanBuffer, indices: &PrimitiveArray, ) -> BooleanBuffer { @@ -633,7 +673,7 @@ fn take_bits( index_nulls.valid_indices().for_each(|valid_idx| { // SAFETY: valid_idx < indices.len(), guaranteed by valid_indices(). let index_val = unsafe { indices.value_unchecked(valid_idx) }.as_usize(); - if CHECKED { + if VALIDATE_INDICES { if values.value(index_val) { // SAFETY: valid_idx < indices.len() = len, output buffer holds len bits. unsafe { bit_util::set_bit_raw(out_ptr, valid_idx) }; @@ -659,7 +699,7 @@ fn take_bits( // SAFETY: base + bit < full_bytes * 8 <= len, so base + bit is a valid // position in the indices array. let index_val = unsafe { indices.value_unchecked(base + bit) }.as_usize(); - if CHECKED { + if VALIDATE_INDICES { byte |= (values.value(index_val) as u8) << bit; } else { // SAFETY: caller guarantees index_val < values.len(). @@ -676,7 +716,7 @@ fn take_bits( // SAFETY: base + bit < len (remainder loop bound), so base + bit is a // valid position in the indices array. let index_val = unsafe { indices.value_unchecked(base + bit) }.as_usize(); - if CHECKED { + if VALIDATE_INDICES { byte |= (values.value(index_val) as u8) << bit; } else { // SAFETY: caller guarantees index_val < values.len(). @@ -693,7 +733,7 @@ fn take_bits( /// Gather value bits and validity bits from two boolean buffers in a single pass. /// Used when the values array itself has nulls, avoiding two separate `take_bits` calls. #[inline(never)] -fn take_bits_with_validity( +fn take_bits_with_validity( values: &BooleanBuffer, validity: &BooleanBuffer, indices: &PrimitiveArray, @@ -715,7 +755,7 @@ fn take_bits_with_validity( for out_pos in index_nulls.valid_indices() { // SAFETY: out_pos < indices.len(), guaranteed by valid_indices(). let src_idx = unsafe { indices.value_unchecked(out_pos) }.as_usize(); - if CHECKED { + if VALIDATE_INDICES { if values.value(src_idx) { // SAFETY: out_pos < indices.len() = len, output buffer holds len bits. unsafe { bit_util::set_bit_raw(value_out_ptr, out_pos) }; @@ -760,7 +800,7 @@ fn take_bits_with_validity( for bit_pos in 0..8usize { // SAFETY: bit_base + bit_pos < full_bytes * 8 <= len. let src_idx = unsafe { indices.value_unchecked(bit_base + bit_pos) }.as_usize(); - if CHECKED { + if VALIDATE_INDICES { packed_values |= (values.value(src_idx) as u8) << bit_pos; packed_validity |= (validity.value(src_idx) as u8) << bit_pos; } else { @@ -784,7 +824,7 @@ fn take_bits_with_validity( for bit_pos in 0..(len - bit_base) { // SAFETY: bit_base + bit_pos < len (remainder loop bound). let src_idx = unsafe { indices.value_unchecked(bit_base + bit_pos) }.as_usize(); - if CHECKED { + if VALIDATE_INDICES { packed_values |= (values.value(src_idx) as u8) << bit_pos; packed_validity |= (validity.value(src_idx) as u8) << bit_pos; } else { @@ -809,7 +849,7 @@ fn take_bits_with_validity( } /// `take` implementation for boolean arrays -fn take_boolean( +fn take_boolean( array: &BooleanArray, indices: &PrimitiveArray, ) -> BooleanArray { @@ -817,19 +857,19 @@ fn take_boolean( match array.nulls().filter(|n| n.null_count() > 0) { Some(array_nulls) => { let (val_buf, null_buf) = - take_bits_with_validity::<_, CHECKED>(bits, array_nulls.inner(), indices); + take_bits_with_validity::<_, VALIDATE_INDICES>(bits, array_nulls.inner(), indices); BooleanArray::new(val_buf, null_buf) } None => { - let val_buf = take_bits::<_, CHECKED>(bits, indices); - let null_buf = take_nulls::<_, CHECKED>(None, indices); + let val_buf = take_bits::<_, VALIDATE_INDICES>(bits, indices); + let null_buf = take_nulls::<_, VALIDATE_INDICES>(None, indices); BooleanArray::new(val_buf, null_buf) } } } /// `take` implementation for string arrays -fn take_bytes( +fn take_bytes( array: &GenericByteArray, indices: &PrimitiveArray, ) -> Result, ArrowError> { @@ -839,7 +879,7 @@ fn take_bytes(array.nulls(), indices); + let nulls = take_nulls::<_, VALIDATE_INDICES>(array.nulls(), indices); // Branch on output nulls — `None` means every output slot is valid. match nulls.as_ref().filter(|n| n.null_count() > 0) { @@ -960,12 +1000,12 @@ fn take_bytes( +fn take_byte_view( array: &GenericByteViewArray, indices: &PrimitiveArray, ) -> Result, ArrowError> { - let new_views = take_native(array.views(), indices); - let new_nulls = take_nulls::<_, CHECKED>(array.nulls(), indices); + let new_views = take_native::<_, _, VALIDATE_INDICES>(array.views(), indices); + let new_nulls = take_nulls::<_, VALIDATE_INDICES>(array.nulls(), indices); let buffers = Arc::clone(array.data_buffers()); // Safety: array.views was valid, and take_native copies only valid values, and verifies bounds Ok(unsafe { GenericByteViewArray::new_unchecked(new_views, buffers, new_nulls) }) @@ -975,7 +1015,7 @@ fn take_byte_view( +fn take_list( values: &GenericListArray, indices: &PrimitiveArray, ) -> Result, ArrowError> @@ -987,7 +1027,13 @@ where { let src_offsets = values.value_offsets(); let child_data = values.values().to_data(); - let nulls = take_nulls::<_, CHECKED>(values.nulls(), indices); + // The number of child values that the rows use, excluding values outside a slice + let visible_child_len = (values.offsets().last() - values.offsets().first()).as_usize(); + // An estimate of the output's child length, from the average length of the rows + let estimated_child_len = visible_child_len + .checked_div(values.len()) + .map_or(0, |avg| avg.saturating_mul(indices.len())); + let nulls = take_nulls::<_, VALIDATE_INDICES>(values.nulls(), indices); let mut dst_offsets = Vec::with_capacity(indices.len() + 1); dst_offsets.push(OffsetType::Native::zero()); @@ -1000,15 +1046,7 @@ where let values_buf = &child_data.buffers()[0]; let child_buf_offset = child_data.offset() * bytes_per_value; - let avg_row_len = child_data - .len() - .checked_div(values.len().max(1)) - .unwrap_or(0); - let mut dst_buf = MutableBuffer::new( - avg_row_len - .saturating_mul(indices.len()) - .saturating_mul(bytes_per_value), - ); + let mut dst_buf = MutableBuffer::new(estimated_child_len.saturating_mul(bytes_per_value)); let mut child_len = OffsetType::Native::zero(); @@ -1032,15 +1070,25 @@ where if prev < vidx { dst_offsets.extend(std::iter::repeat_n(child_len, vidx - prev)); } - let row = if CHECKED { - indices.value(vidx).as_usize() + // SAFETY: vidx < indices.len(), guaranteed by valid_indices(). + let row = unsafe { indices.value_unchecked(vidx) }.as_usize(); + let (start, end) = if VALIDATE_INDICES { + ( + child_buf_offset + src_offsets[row].as_usize() * bytes_per_value, + child_buf_offset + src_offsets[row + 1].as_usize() * bytes_per_value, + ) } else { - // SAFETY: !CHECKED means the caller guarantees all indices are valid; - // `vidx` is further bounded by the validity bitmap of `indices`. - unsafe { indices.value_unchecked(vidx) }.as_usize() + // SAFETY: caller guarantees row < values.len(), so row+1 is also in bounds. + unsafe { + ( + child_buf_offset + + src_offsets.get_unchecked(row).as_usize() * bytes_per_value, + child_buf_offset + + src_offsets.get_unchecked(row + 1).as_usize() + * bytes_per_value, + ) + } }; - let start = child_buf_offset + src_offsets[row].as_usize() * bytes_per_value; - let end = child_buf_offset + src_offsets[row + 1].as_usize() * bytes_per_value; dst_buf.extend_from_slice(&values_buf[start..end]); child_len = child_len .checked_add(&(src_offsets[row + 1] - src_offsets[row])) @@ -1071,13 +1119,11 @@ where return GenericListArray::::try_new(field, offsets, child, nulls); } - let capacity = child_data - .len() - .checked_div(values.len()) - .map(|avg| avg * indices.len()) - .unwrap_or_default(); - let mut mutable = - MutableArrayData::new(vec![&child_data], child_data.null_count() > 0, capacity); + let mut mutable = MutableArrayData::new( + vec![&child_data], + child_data.null_count() > 0, + estimated_child_len, + ); match nulls.as_ref().filter(|n| n.null_count() > 0) { None => { @@ -1102,18 +1148,20 @@ where if last < i { dst_offsets.extend(std::iter::repeat_n(current, i - last)); } - let row = if CHECKED { - indices.value(i).as_usize() + // SAFETY: i < indices.len(), guaranteed by valid_indices(). + let row = unsafe { indices.value_unchecked(i) }.as_usize(); + let (start, end) = if VALIDATE_INDICES { + (src_offsets[row].as_usize(), src_offsets[row + 1].as_usize()) } else { - // SAFETY: !CHECKED means the caller guarantees all indices are valid; - // `i` is further bounded by the validity bitmap of `indices`. - unsafe { indices.value_unchecked(i) }.as_usize() + // SAFETY: caller guarantees row < values.len(), so row+1 is also in bounds. + unsafe { + ( + src_offsets.get_unchecked(row).as_usize(), + src_offsets.get_unchecked(row + 1).as_usize(), + ) + } }; - mutable.try_extend( - 0, - src_offsets[row].as_usize(), - src_offsets[row + 1].as_usize(), - )?; + mutable.try_extend(0, start, end)?; dst_offsets.push( OffsetType::Native::from_usize(mutable.len()) .ok_or_else(|| ArrowError::OffsetOverflowError(mutable.len()))?, @@ -1135,7 +1183,7 @@ where GenericListArray::::try_new(field, offsets, child, nulls) } -fn take_list_view( +fn take_list_view( values: &GenericListViewArray, indices: &PrimitiveArray, ) -> Result, ArrowError> @@ -1144,9 +1192,9 @@ where OffsetType: ArrowPrimitiveType, OffsetType::Native: OffsetSizeTrait, { - let taken_offsets = take_native(values.offsets(), indices); - let taken_sizes = take_native(values.sizes(), indices); - let nulls = take_nulls::<_, CHECKED>(values.nulls(), indices); + let taken_offsets = take_native::<_, _, VALIDATE_INDICES>(values.offsets(), indices); + let taken_sizes = take_native::<_, _, VALIDATE_INDICES>(values.sizes(), indices); + let nulls = take_nulls::<_, VALIDATE_INDICES>(values.nulls(), indices); let field = match values.data_type() { DataType::ListView(field) | DataType::LargeListView(field) => field.clone(), @@ -1171,14 +1219,14 @@ where /// 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( +fn take_fixed_size_list( values: &FixedSizeListArray, indices: &PrimitiveArray, length: ::Native, ) -> Result { let field = values.value_field(); let child = values.values(); - let nulls = take_nulls::<_, CHECKED>(values.nulls(), indices); + let nulls = take_nulls::<_, VALIDATE_INDICES>(values.nulls(), indices); // Fast path: primitive child with no nulls copy row-sized byte blocks directly, let taken_child = if child.null_count() == 0 @@ -1193,7 +1241,7 @@ fn take_fixed_size_list( ) } else { let list_indices = take_value_indices_from_fixed_size_list(values, indices, length)?; - take_impl::(child.as_ref(), &list_indices)? + take_impl::(child.as_ref(), &list_indices)? }; FixedSizeListArray::try_new_with_length( @@ -1265,7 +1313,7 @@ fn take_fixed_size_list_primitive( /// The computation is done in two steps: /// - Compute the values buffer /// - Compute the null buffer -fn take_fixed_size_binary( +fn take_fixed_size_binary( values: &FixedSizeBinaryArray, indices: &PrimitiveArray, size: i32, @@ -1283,7 +1331,7 @@ fn take_fixed_size_binary( _ => take_fixed_size_binary_buffer_dynamic_length(values, indices, size_usize), }; - let value_nulls = take_nulls::<_, CHECKED>(values.nulls(), indices); + let value_nulls = take_nulls::<_, VALIDATE_INDICES>(values.nulls(), indices); let final_nulls = NullBuffer::union(value_nulls.as_ref(), indices.nulls()); return FixedSizeBinaryArray::try_new(size, result_buffer, final_nulls); @@ -1395,11 +1443,11 @@ fn take_fixed_size( /// /// 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( +fn take_dict( values: &DictionaryArray, indices: &PrimitiveArray, ) -> Result, ArrowError> { - let new_keys = take_primitive::<_, _, CHECKED>(values.keys(), indices)?; + let new_keys = take_primitive::<_, _, VALIDATE_INDICES>(values.keys(), indices)?; Ok(unsafe { DictionaryArray::new_unchecked(new_keys, values.values().clone()) }) } @@ -1661,12 +1709,22 @@ pub fn take_record_batch( record_batch: &RecordBatch, indices: &dyn Array, ) -> Result { - let columns = record_batch - .columns() - .iter() - .map(|c| take(c, indices, None)) - .collect::, _>>()?; - RecordBatch::try_new(record_batch.schema(), columns) + downcast_integer_array!( + indices => { + let indices = indices.to_indices(); + let cols = record_batch.columns(); + let mut columns = Vec::with_capacity(cols.len()); + if let Some(first) = cols.first() { + columns.push(take_impl::<_, true>(first.as_ref(), &indices)?); + for col in &cols[1..] { + columns.push(take_impl::<_, false>(col.as_ref(), &indices)?); + } + } + // Safety: indices were validated by the first take_impl call + Ok(unsafe { RecordBatch::new_unchecked(record_batch.schema(), columns, indices.len()) }) + }, + d => Err(ArrowError::InvalidArgumentError(format!("Take only supported for integers, got {d:?}"))) + ) } #[cfg(test)] @@ -1674,7 +1732,6 @@ mod tests { use super::*; use arrow_array::builder::*; use arrow_buffer::{IntervalDayTime, IntervalMonthDayNano}; - use arrow_data::ArrayData; use arrow_schema::{Field, Fields, TimeUnit, UnionFields}; use num_traits::ToPrimitive; @@ -2174,18 +2231,12 @@ mod tests { #[test] fn test_take_bool_nullable_index() { // indices where the masked invalid elements would be out of bounds - let index_data = ArrayData::try_new( - DataType::UInt32, - 6, - Some(Buffer::from_iter(vec![ + let index = UInt32Array::new( + ScalarBuffer::from(vec![99, 0, 999, 1, 9999, 2]), + Some(NullBuffer::from(vec![ false, true, false, true, false, true, ])), - 0, - vec![Buffer::from_iter(vec![99, 0, 999, 1, 9999, 2])], - vec![], - ) - .unwrap(); - let index = UInt32Array::from(index_data); + ); test_take_boolean_arrays( vec![Some(true), None, Some(false)], &index, @@ -2197,18 +2248,12 @@ mod tests { #[test] fn test_take_bool_nullable_index_nonnull_values() { // indices where the masked invalid elements would be out of bounds - let index_data = ArrayData::try_new( - DataType::UInt32, - 6, - Some(Buffer::from_iter(vec![ + let index = UInt32Array::new( + ScalarBuffer::from(vec![99, 0, 999, 1, 9999, 2]), + Some(NullBuffer::from(vec![ false, true, false, true, false, true, ])), - 0, - vec![Buffer::from_iter(vec![99, 0, 999, 1, 9999, 2])], - vec![], - ) - .unwrap(); - let index = UInt32Array::from(index_data); + ); test_take_boolean_arrays( vec![Some(true), Some(true), Some(false)], &index, @@ -2464,20 +2509,11 @@ mod tests { 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]).into_data(); - // 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(Arc::new(Field::new_list_field(DataType::Int32, false))); - let list_data = ArrayData::builder(list_data_type.clone()) - .len(4) - .add_buffer(value_offsets) - .add_child_data(value_data) - .build() - .unwrap(); - let list_array = $list_array_type::from(list_data); + let values = Arc::new(Int32Array::from(vec![0, 0, 0, -1, -2, -1, 2, 3])); + let value_offsets = + OffsetBuffer::<$offset_type>::new(ScalarBuffer::from(vec![0, 3, 6, 6, 8])); + let field = Arc::new(Field::new_list_field(DataType::Int32, false)); + let list_array = $list_array_type::new(Arc::clone(&field), value_offsets, values, None); // index returns: [[2,3], null, [-1,-2,-1], [], [0,0,0]] let index = UInt32Array::from(vec![Some(3), None, Some(1), Some(2), Some(0)]); @@ -2487,7 +2523,7 @@ mod tests { // construct a value array with expected results: // [[2,3], null, [-1,-2,-1], [], [0,0,0]] - let expected_data = Int32Array::from(vec![ + let expected_values = Arc::new(Int32Array::from(vec![ Some(2), Some(3), Some(-1), @@ -2496,21 +2532,16 @@ mod tests { Some(0), Some(0), Some(0), - ]) - .into_data(); - // construct offsets - let expected_offsets: [$offset_type; 6] = [0, 2, 2, 5, 5, 8]; - 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) + ])); + let expected_offsets = + OffsetBuffer::<$offset_type>::new(ScalarBuffer::from(vec![0, 2, 2, 5, 5, 8])); + let expected_list_array = $list_array_type::new( + field, + expected_offsets, + expected_values, // null buffer remains the same as only the indices have nulls - .nulls(index.nulls().cloned()) - .add_buffer(expected_offsets) - .add_child_data(expected_data) - .build() - .unwrap(); - let expected_list_array = $list_array_type::from(expected_list_data); + index.nulls().cloned(), + ); assert_eq!(a, &expected_list_array); }}; @@ -2519,7 +2550,7 @@ mod tests { 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![ + let values = Arc::new(Int32Array::from(vec![ Some(0), None, Some(0), @@ -2529,22 +2560,11 @@ mod tests { None, Some(5), None, - ]) - .into_data(); - // 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(Arc::new(Field::new_list_field(DataType::Int32, true))); - let list_data = ArrayData::builder(list_data_type.clone()) - .len(4) - .add_buffer(value_offsets) - .null_bit_buffer(Some(Buffer::from([0b11111111]))) - .add_child_data(value_data) - .build() - .unwrap(); - let list_array = $list_array_type::from(list_data); + ])); + let value_offsets = + OffsetBuffer::<$offset_type>::new(ScalarBuffer::from(vec![0, 3, 6, 7, 9])); + let field = Arc::new(Field::new_list_field(DataType::Int32, true)); + let list_array = $list_array_type::new(Arc::clone(&field), value_offsets, values, None); // 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)]); @@ -2554,7 +2574,7 @@ mod tests { // construct a value array with expected results: // [[null], null, [-1,-2,3], [5,null], [0,null,0]] - let expected_data = Int32Array::from(vec![ + let expected_values = Arc::new(Int32Array::from(vec![ None, Some(-1), Some(-2), @@ -2564,21 +2584,16 @@ mod tests { Some(0), None, Some(0), - ]) - .into_data(); - // 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) + ])); + let expected_offsets = + OffsetBuffer::<$offset_type>::new(ScalarBuffer::from(vec![0, 1, 1, 4, 6, 9])); + let expected_list_array = $list_array_type::new( + field, + expected_offsets, + expected_values, // null buffer remains the same as only the indices have nulls - .nulls(index.nulls().cloned()) - .add_buffer(expected_offsets) - .add_child_data(expected_data) - .build() - .unwrap(); - let expected_list_array = $list_array_type::from(expected_list_data); + index.nulls().cloned(), + ); assert_eq!(a, &expected_list_array); }}; @@ -2587,7 +2602,7 @@ mod tests { 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![ + let values = Arc::new(Int32Array::from(vec![ Some(0), None, Some(0), @@ -2596,22 +2611,14 @@ mod tests { Some(3), Some(5), None, - ]) - .into_data(); - // 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(Arc::new(Field::new_list_field(DataType::Int32, true))); - let list_data = ArrayData::builder(list_data_type.clone()) - .len(4) - .add_buffer(value_offsets) - .null_bit_buffer(Some(Buffer::from([0b11111011]))) - .add_child_data(value_data) - .build() - .unwrap(); - let list_array = $list_array_type::from(list_data); + ])); + let value_offsets = + OffsetBuffer::<$offset_type>::new(ScalarBuffer::from(vec![0, 3, 6, 6, 8])); + let field = Arc::new(Field::new_list_field(DataType::Int32, true)); + // the entry at index 2 is null + let list_nulls = NullBuffer::from(vec![true, true, false, true]); + let list_array = + $list_array_type::new(Arc::clone(&field), value_offsets, values, Some(list_nulls)); // 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)]); @@ -2621,7 +2628,7 @@ mod tests { // construct a value array with expected results: // [null, null, [-1,-2,3], [5,null], [0,null,0]] - let expected_data = Int32Array::from(vec![ + let expected_values = Arc::new(Int32Array::from(vec![ Some(-1), Some(-2), Some(3), @@ -2630,25 +2637,17 @@ mod tests { Some(0), None, Some(0), - ]) - .into_data(); - // 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(Some(Buffer::from(null_bits))) - .add_buffer(expected_offsets) - .add_child_data(expected_data) - .build() - .unwrap(); - let expected_list_array = $list_array_type::from(expected_list_data); + ])); + let expected_offsets = + OffsetBuffer::<$offset_type>::new(ScalarBuffer::from(vec![0, 0, 0, 3, 5, 8])); + // null buffer must be recalculated as both values and indices have nulls + let expected_nulls = NullBuffer::from(vec![false, false, true, true, true]); + let expected_list_array = $list_array_type::new( + field, + expected_offsets, + expected_values, + Some(expected_nulls), + ); assert_eq!(a, &expected_list_array); }}; @@ -2958,19 +2957,12 @@ mod tests { #[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]).into_data(); + let values = Arc::new(Int32Array::from(vec![0, 0, 0, -1, -2, -1, 2, 3])); // Construct offsets - let value_offsets = Buffer::from_slice_ref([0, 3, 6, 8]); + let value_offsets = OffsetBuffer::::new(ScalarBuffer::from(vec![0, 3, 6, 8])); // Construct a list array from the above two - let list_data_type = - DataType::List(Arc::new(Field::new_list_field(DataType::Int32, false))); - let list_data = ArrayData::builder(list_data_type) - .len(3) - .add_buffer(value_offsets) - .add_child_data(value_data) - .build() - .unwrap(); - let list_array = ListArray::from(list_data); + let field = Arc::new(Field::new_list_field(DataType::Int32, false)); + let list_array = ListArray::new(field, value_offsets, values, None); let index = UInt32Array::from(vec![1000]); @@ -3262,6 +3254,38 @@ mod tests { test_take_sliced_list_with_value_nulls_generic::(); } + #[test] + fn test_take_sliced_list_memory_size() { + // Takes the same rows from a two-row slice of a large array and from an unsliced + // copy of those two rows. The results should be equal and use the same memory. + fn check(parent: ArrayRef, unsliced: ArrayRef) { + let sliced = parent.slice(500, 2); + let indices = UInt32Array::from_iter_values((0..100).map(|i| i % 2)); + let actual = take(&sliced, &indices, None).unwrap(); + let expected = take(&unsliced, &indices, None).unwrap(); + assert_eq!(&actual, &expected); + assert_eq!( + actual.get_array_memory_size(), + expected.get_array_memory_size() + ); + } + + // Lists of 10 values per row + let list = |rows: usize, values: ArrayRef| -> ArrayRef { + let field = Arc::new(Field::new_list_field(values.data_type().clone(), false)); + let offsets = OffsetBuffer::from_repeated_length(10, rows); + Arc::new(ListArray::new(field, offsets, values, None)) + }; + check( + list(1000, Arc::new(Int64Array::from(vec![1; 10_000]))), + list(2, Arc::new(Int64Array::from(vec![1; 20]))), + ); + check( + list(1000, Arc::new(StringArray::from(vec!["a"; 10_000]))), + list(2, Arc::new(StringArray::from(vec!["a"; 20]))), + ); + } + #[test] fn test_take_runs() { let logical_array: Vec = vec![1_i32, 1, 2, 2, 1, 1, 1, 2, 2, 1, 1, 2, 2]; diff --git a/arrow-string/Cargo.toml b/arrow-string/Cargo.toml index 7d0b20edf52b..264c35b9d2d0 100644 --- a/arrow-string/Cargo.toml +++ b/arrow-string/Cargo.toml @@ -19,7 +19,6 @@ name = "arrow-string" version = { workspace = true } description = "String kernels for arrow arrays" -homepage = { workspace = true } repository = { workspace = true } authors = { workspace = true } license = { workspace = true } diff --git a/arrow-string/src/concat_elements.rs b/arrow-string/src/concat_elements.rs index 0df4216c66b9..fd83b45d223d 100644 --- a/arrow-string/src/concat_elements.rs +++ b/arrow-string/src/concat_elements.rs @@ -47,16 +47,15 @@ pub fn concat_elements_bytes( let nulls = NullBuffer::union(left.nulls(), right.nulls()); - let left_offsets = left.value_offsets(); - let right_offsets = right.value_offsets(); + let left_offsets = left.offsets(); + let right_offsets = right.offsets(); let left_values = left.value_data(); let right_values = right.value_data(); let mut output_values = Vec::with_capacity( - left_values.len() + right_values.len() - - left_offsets[0].as_usize() - - right_offsets[0].as_usize(), + (left_offsets.last() - left_offsets.first()).as_usize() + + (right_offsets.last() - right_offsets.first()).as_usize(), ); let mut output_offsets = Vec::with_capacity(left_offsets.len()); @@ -154,10 +153,9 @@ pub fn concat_elements_utf8_many( .collect::>(); let mut output_values = Vec::with_capacity( - data_values + arrays .iter() - .zip(offsets.iter_mut()) - .map(|(data, offset)| data.len() - offset.peek().unwrap().as_usize()) + .map(|array| (array.offsets().last() - array.offsets().first()).as_usize()) .sum(), ); @@ -279,23 +277,23 @@ where let null_buffer = NullBuffer::union(left.nulls(), right.nulls()); // Compute the required data buffer size, excluding any elements that are null - // or are small enough to be stored inline. - let data_size = match &null_buffer { + // or are small enough to be stored inline. Both lengths are `u32`, and a + // pair of them can sum past that, so the addition is done in `usize`. + let inline_max = MAX_INLINE_VIEW_LEN as usize; + let data_size: usize = match &null_buffer { None => left .lengths() .zip(right.lengths()) - .map(|(l, r)| l + r) - .filter(|len| *len > MAX_INLINE_VIEW_LEN) - .map(|len| len as usize) + .map(|(l, r)| l as usize + r as usize) + .filter(|len| *len > inline_max) .sum(), Some(nb) => left .lengths() .zip(right.lengths()) .zip(nb.iter()) .filter(|((_, _), not_null)| *not_null) - .map(|((l, r), _)| l + r) - .filter(|len| *len > MAX_INLINE_VIEW_LEN) - .map(|len| len as usize) + .map(|((l, r), _)| l as usize + r as usize) + .filter(|len| *len > inline_max) .sum(), }; @@ -608,6 +606,20 @@ mod tests { assert_eq!(output.as_string(), &expected); } + #[test] + fn test_concat_slice_capacity() { + let left = StringArray::from(vec!["hello"; 8]).slice(2, 3); + let right = StringArray::from(vec![" world"; 8]).slice(4, 3); + + let pair = concat_elements_utf8(&left, &right).unwrap(); + assert_eq!(pair, StringArray::from(vec!["hello world"; 3])); + assert_eq!(pair.values().capacity(), 33); + + let many = concat_elements_utf8_many(&[&left, &right, &left]).unwrap(); + assert_eq!(many, StringArray::from(vec!["hello worldhello"; 3])); + assert_eq!(many.values().capacity(), 48); + } + #[test] fn test_string_concat_error_empty() { assert_eq!( diff --git a/arrow-string/src/regexp.rs b/arrow-string/src/regexp.rs index 0474f2a7b8ea..015b8b477823 100644 --- a/arrow-string/src/regexp.rs +++ b/arrow-string/src/regexp.rs @@ -332,6 +332,8 @@ fn get_scalar_pattern_flag_utf8view<'a>( macro_rules! process_regexp_match { ($array:expr, $regex:expr, $list_builder:expr) => { + // Reuse capture location storage across matches + let mut capture_locations = None; $array .iter() .map(|value| { @@ -342,19 +344,25 @@ macro_rules! process_regexp_match { $list_builder.values().append_value(""); $list_builder.append(true); } - Some(value) => match $regex.captures(value) { - Some(caps) => { - let mut iter = caps.iter(); - if caps.len() > 1 { - iter.next(); - } - for m in iter.flatten() { - $list_builder.values().append_value(m.as_str()); + Some(value) => { + let locations = + capture_locations.get_or_insert_with(|| $regex.capture_locations()); + match $regex.captures_read(locations, value) { + Some(_) => { + let mut groups = 0..locations.len(); + if locations.len() > 1 { + // Skip group 0 (the whole match) when explicit capture groups exist. + groups.next(); + } + let iter = groups.map(|group| locations.get(group)); + for (start, end) in iter.flatten() { + $list_builder.values().append_value(&value[start..end]); + } + $list_builder.append(true); } - $list_builder.append(true); + None => $list_builder.append(false), } - None => $list_builder.append(false), - }, + } None => $list_builder.append(false), } Ok(()) @@ -519,6 +527,47 @@ pub fn regexp_match( mod tests { use super::*; + #[test] + fn match_scalar_pattern_capture_groups() { + macro_rules! check { + ($array_type:ty, $builder_type:ty) => {{ + // "β" follows "aβ" to check that an optional group matched in + // one row is not reported for the next row. + let array = <$array_type>::from(vec![Some("aβ"), Some("β"), Some("none"), None]); + let cases: [(&str, [Option<&[&str]>; 4]); 4] = [ + // Optional groups that did not match are omitted. + ("(a)?(β)", [Some(&["a", "β"]), Some(&["β"]), None, None]), + // Without explicit groups, the whole match is returned. + ("β", [Some(&["β"]), Some(&["β"]), None, None]), + ("()", [Some(&[""]), Some(&[""]), Some(&[""]), None]), + ("", [Some(&[""]), Some(&[""]), Some(&[""]), None]), + ]; + for (pattern, rows) in cases { + let mut expected = ListBuilder::new(<$builder_type>::new()); + for row in rows { + match row { + Some(values) => { + for value in values { + expected.values().append_value(value); + } + expected.append(true); + } + None => expected.append(false), + } + } + let expected = expected.finish(); + + let pattern_scalar = Scalar::new(<$array_type>::from(vec![pattern])); + let actual = regexp_match(&array, &pattern_scalar, None).unwrap(); + assert_eq!(actual.as_list::(), &expected, "{pattern}"); + } + }}; + } + check!(StringArray, GenericStringBuilder); + check!(LargeStringArray, GenericStringBuilder); + check!(StringViewArray, StringViewBuilder); + } + macro_rules! test_match_single_group { ($test_name:ident, $values:expr, $patterns:expr, $arr_type:ty, $builder_type:ty, $expected:expr) => { #[test] diff --git a/arrow-string/src/substring.rs b/arrow-string/src/substring.rs index 82171b5a1d8b..dfc10e099c06 100644 --- a/arrow-string/src/substring.rs +++ b/arrow-string/src/substring.rs @@ -350,14 +350,17 @@ where // When array is [Large]StringArray, we will check whether `offset` is at a valid char boundary. let check_char_boundary = { - |offset: T::Offset| { + let first = array.offsets().first().as_usize(); + let last = array.offsets().last().as_usize(); + move |offset: T::Offset| { if !matches!(T::DATA_TYPE, DataType::Utf8 | DataType::LargeUtf8) { return Ok(offset); } - // Safety: a StringArray must contain valid UTF8 data - let data_str = unsafe { std::str::from_utf8_unchecked(data) }; + // Safety: in a StringArray, the bytes between the first and last offsets + // are valid UTF-8. The bytes outside that range need not be. + let data_str = unsafe { std::str::from_utf8_unchecked(&data[first..last]) }; let offset_usize = offset.as_usize(); - if data_str.is_char_boundary(offset_usize) { + if data_str.is_char_boundary(offset_usize - first) { Ok(offset) } else { Err(ArrowError::ComputeError(format!( @@ -1129,6 +1132,15 @@ mod tests { assert!(err.contains("invalid utf-8 boundary")); } + #[test] + fn check_invalid_bytes_outside_offsets() { + // The values buffer holds invalid UTF-8 before and after the visible values + let values = Buffer::from_slice_ref(b"\xFFab\x80"); + let array = StringArray::new(OffsetBuffer::new(vec![1, 3].into()), values, None); + let result = substring(&array, 0, Some(5)).unwrap(); + assert_eq!(result.as_string::(), &StringArray::from(vec!["ab"])); + } + #[test] fn non_utf8_bytes() { // non-utf8 bytes diff --git a/arrow/Cargo.toml b/arrow/Cargo.toml index ba81249a63ff..cb8746d5851e 100644 --- a/arrow/Cargo.toml +++ b/arrow/Cargo.toml @@ -19,7 +19,6 @@ name = "arrow" version = { workspace = true } description = "Rust implementation of Apache Arrow" -homepage = { workspace = true } repository = { workspace = true } authors = { workspace = true } license = { workspace = true } diff --git a/arrow/benches/cast_kernels.rs b/arrow/benches/cast_kernels.rs index 5956bdb89ef1..811a46cd5818 100644 --- a/arrow/benches/cast_kernels.rs +++ b/arrow/benches/cast_kernels.rs @@ -259,7 +259,15 @@ fn cast_array(array: &ArrayRef, to_type: DataType) { fn add_benchmark(c: &mut Criterion) { let i32_array = build_array::(512); + let i32_array_8192 = build_array::(8192); + let bool_array: ArrayRef = Arc::new(create_boolean_array(8192, 0.1, 0.5)); + let bool_array_no_nulls: ArrayRef = Arc::new(create_boolean_array(8192, 0.0, 0.5)); let i64_array = build_array::(512); + let i64_within_1e6 = create_primitive_array_range::(512, 0.1, -1_000_000..1_000_000); + let mut values: Vec> = i64_within_1e6.iter().collect(); + *values.last_mut().unwrap() = Some(10_000_000_000_000_000); + let i64_late_overflow = Arc::new(Int64Array::from(values)) as ArrayRef; + let i64_within_1e6 = Arc::new(i64_within_1e6) as ArrayRef; let f32_array = build_array::(512); let f32_utf8_array = cast(&build_array::(512), &DataType::Utf8).unwrap(); let i32_utf8_array = cast(&build_array::(512), &DataType::Utf8).unwrap(); @@ -319,6 +327,24 @@ fn add_benchmark(c: &mut Criterion) { c.bench_function("cast int32 to int64 512", |b| { b.iter(|| cast_array(&i32_array, DataType::Int64)) }); + c.bench_function("cast int32 to bool 8192", |b| { + b.iter(|| cast_array(&i32_array_8192, DataType::Boolean)) + }); + c.bench_function("cast bool to int32 8192", |b| { + b.iter(|| cast_array(&bool_array, DataType::Int32)) + }); + c.bench_function("cast bool to string 8192", |b| { + b.iter(|| cast_array(&bool_array, DataType::Utf8)) + }); + c.bench_function("cast bool to string no nulls 8192", |b| { + b.iter(|| cast_array(&bool_array_no_nulls, DataType::Utf8)) + }); + c.bench_function("cast bool to string view 8192", |b| { + b.iter(|| cast_array(&bool_array, DataType::Utf8View)) + }); + c.bench_function("cast bool to string view no nulls 8192", |b| { + b.iter(|| cast_array(&bool_array_no_nulls, DataType::Utf8View)) + }); c.bench_function("cast float32 to int32 512", |b| { b.iter(|| cast_array(&f32_array, DataType::Int32)) }); @@ -337,6 +363,16 @@ fn add_benchmark(c: &mut Criterion) { c.bench_function("cast int64 to decimal32(9, -1) 512", |b| { b.iter(|| cast_array(&i64_array, DataType::Decimal32(9, -1))) }); + c.bench_function("cast int64 to decimal128(38, 10) 512", |b| { + b.iter(|| cast_array(&i64_array, DataType::Decimal128(38, 10))) + }); + c.bench_function("cast int64 within 1e6 to decimal128(18, 2) 512", |b| { + b.iter(|| cast_array(&i64_within_1e6, DataType::Decimal128(18, 2))) + }); + c.bench_function( + "cast int64 within 1e6 late overflow to decimal128(18, 2) 512", + |b| b.iter(|| cast_array(&i64_late_overflow, DataType::Decimal128(18, 2))), + ); c.bench_function("cast date64 to date32 512", |b| { b.iter(|| cast_array(&date64_array, DataType::Date32)) }); diff --git a/arrow/benches/concatenate_elements.rs b/arrow/benches/concatenate_elements.rs index fcacbe3284d1..33197cd11ed8 100644 --- a/arrow/benches/concatenate_elements.rs +++ b/arrow/benches/concatenate_elements.rs @@ -22,7 +22,7 @@ use criterion::Criterion; use arrow::array::*; use arrow::util::bench_util::*; -use arrow_string::concat_elements::concat_elements_dyn; +use arrow_string::concat_elements::{concat_elements_dyn, concat_elements_utf8_many}; use std::hint; fn bench_concat(v1: &dyn Array, v2: &dyn Array) { @@ -68,6 +68,21 @@ fn add_benchmark(c: &mut Criterion) { c.bench_function(&id, |b| b.iter(|| bench_concat(&array, &array))); } } + + // A slice whose values buffer extends far past its last row, compared with + // a compact copy of the same rows + let sliced = create_string_array_with_len::(131_072, 0.1, 32).slice(8_192, 8_192); + let compact = StringArray::from_iter(sliced.iter()); + for (name, array) in [("compact", compact), ("sliced", sliced)] { + c.bench_function(&format!("concat str {name} 8192"), |b| { + b.iter(|| bench_concat(&array, &array)) + }); + c.bench_function(&format!("concat str many {name} 8192"), |b| { + b.iter(|| { + hint::black_box(concat_elements_utf8_many(&[&array, &array, &array]).unwrap()) + }) + }); + } } criterion_group!(benches, add_benchmark); diff --git a/arrow/benches/concatenate_kernel.rs b/arrow/benches/concatenate_kernel.rs index 497c2e8f1ed1..6ae97c90a96d 100644 --- a/arrow/benches/concatenate_kernel.rs +++ b/arrow/benches/concatenate_kernel.rs @@ -235,6 +235,21 @@ fn add_benchmark(c: &mut Criterion) { |b| b.iter(|| bench_concat_arrays(&array_refs)), ); } + + // (name, logical_len, physical_len, num_arrays) + for (name, logical, physical, n) in [ + ("small logical=32 physical=4 x4", 32usize, 4usize, 4usize), + ("logical=1024 physical=128 x8", 1024, 128, 8), + ("logical=8192 physical=1024 x10", 8192, 1024, 10), + ] { + let arrays: Vec> = (0..n) + .map(|_| create_primitive_run_array::(logical, physical)) + .collect(); + let array_refs: Vec<&dyn Array> = arrays.iter().map(|a| a as &dyn Array).collect(); + c.bench_function(&format!("concat run_end i32 {name}"), |b| { + b.iter(|| bench_concat_arrays(&array_refs)) + }); + } } criterion_group!(benches, add_benchmark); diff --git a/arrow/benches/filter_kernels.rs b/arrow/benches/filter_kernels.rs index e8b24955492d..8d7a97360521 100644 --- a/arrow/benches/filter_kernels.rs +++ b/arrow/benches/filter_kernels.rs @@ -22,6 +22,7 @@ use arrow::util::bench_util::*; use arrow::array::*; use arrow::compute::filter; +use arrow::compute::{and, is_not_null}; use arrow::datatypes::{Field, Float32Type, Int32Type, Int64Type, Schema, UInt8Type}; use arrow_array::types::Decimal128Type; @@ -115,6 +116,40 @@ fn add_benchmark(c: &mut Criterion) { |b| b.iter(|| bench_built_filter(&data_array, &sparse_filter)), ); + let valid = is_not_null(&data_array).unwrap(); + let valid_filter = FilterBuilder::new(&and(&filter_array, &valid).unwrap()) + .optimize() + .build(); + let valid_dense_filter = FilterBuilder::new(&and(&dense_filter_array, &valid).unwrap()) + .optimize() + .build(); + let valid_sparse_filter = FilterBuilder::new(&and(&sparse_filter_array, &valid).unwrap()) + .optimize() + .build(); + c.bench_function("filter context i32 w NULLs, only valid (kept 1/4)", |b| { + b.iter(|| bench_built_filter(&data_array, &valid_filter)) + }); + c.bench_function( + "filter context i32 w NULLs, only valid high selectivity (kept 1023/2048)", + |b| b.iter(|| bench_built_filter(&data_array, &valid_dense_filter)), + ); + c.bench_function( + "filter context i32 w NULLs, only valid low selectivity (kept 1/2048)", + |b| b.iter(|| bench_built_filter(&data_array, &valid_sparse_filter)), + ); + + // Nulls only in the last word, so the filters below find a selected null late + let data_array: Int32Array = (0..size) + .map(|i| (i < size - 64).then_some(i as i32)) + .collect(); + c.bench_function("filter context i32 w NULLs at end (kept 1/2)", |b| { + b.iter(|| bench_built_filter(&data_array, &filter)) + }); + c.bench_function( + "filter context i32 w NULLs at end high selectivity (kept 1023/1024)", + |b| b.iter(|| bench_built_filter(&data_array, &dense_filter)), + ); + let data_array = create_primitive_array::(size, 0.5); c.bench_function("filter context u8 w NULLs (kept 1/2)", |b| { b.iter(|| bench_built_filter(&data_array, &filter)) diff --git a/arrow/benches/regexp_kernels.rs b/arrow/benches/regexp_kernels.rs index f2483f29f457..40e25cb75b9a 100644 --- a/arrow/benches/regexp_kernels.rs +++ b/arrow/benches/regexp_kernels.rs @@ -17,15 +17,52 @@ #[macro_use] extern crate criterion; -use criterion::Criterion; +use criterion::{BenchmarkId, Criterion, Throughput}; use arrow::array::*; use arrow::compute::kernels::regexp::*; use arrow::util::bench_util::*; use std::hint; -fn bench_regexp(arr: &GenericStringArray, regex_array: &dyn Datum) { - regexp_match(hint::black_box(arr), regex_array, None).unwrap(); +fn bench_regexp(arr: &dyn Array, regex_array: &dyn Datum) { + hint::black_box( + regexp_match(hint::black_box(arr), hint::black_box(regex_array), None).unwrap(), + ); +} + +fn bench_scalar_captures(c: &mut Criterion) { + let size = 8192; + let matches = StringArray::from_iter_values( + (0..size).map(|i| format!("customer-{i:08}-region-{}-purchase", i % 8)), + ); + let no_matches = StringArray::from_iter_values( + (0..size).map(|i| format!("visitor-{i:08}-region-{}-purchase", i % 8)), + ); + let mixed = StringArray::from_iter((0..size).map(|i| match i % 4 { + 0 => None, + 2 => Some(no_matches.value(i)), + _ => Some(matches.value(i)), + })); + let pattern = r"customer-([0-9]+)-region-([0-7])"; + let utf8_pattern = Scalar::new(StringArray::from(vec![pattern])); + let view_pattern = Scalar::new(StringViewArray::from(vec![pattern])); + + let mut group = c.benchmark_group("regexp_scalar_captures"); + group.throughput(Throughput::Elements(size as u64)); + for (name, array) in [ + ("matches", matches), + ("mixed", mixed), + ("no_matches", no_matches), + ] { + let view = StringViewArray::from(&array); + group.bench_function(BenchmarkId::new("utf8", name), |b| { + b.iter(|| bench_regexp(&array, &utf8_pattern)) + }); + group.bench_function(BenchmarkId::new("utf8view", name), |b| { + b.iter(|| bench_regexp(&view, &view_pattern)) + }); + } + group.finish(); } fn add_benchmark(c: &mut Criterion) { @@ -46,5 +83,5 @@ fn add_benchmark(c: &mut Criterion) { }); } -criterion_group!(benches, add_benchmark); +criterion_group!(benches, add_benchmark, bench_scalar_captures); criterion_main!(benches); diff --git a/arrow/tests/array_equal.rs b/arrow/tests/array_equal.rs index c108cc967022..cc45526a021e 100644 --- a/arrow/tests/array_equal.rs +++ b/arrow/tests/array_equal.rs @@ -822,6 +822,22 @@ fn test_test_list_view_array() { test_equal(&a, &b, true); } +fn test_list_view_sizes_and_child_nulls() { + // Same first elements but different sizes, with a null row + let a = create_list_view_array::([Some(vec![Some(1), Some(2)]), None]); + let b = create_list_view_array::([Some(vec![Some(1), Some(2), Some(3)]), None]); + test_equal(&a, &b, false); + + // Child values that differ only in their null masks, without and with a null row + let a = create_list_view_array::([Some(vec![Some(1), Some(0)])]); + let b = create_list_view_array::([Some(vec![Some(1), None])]); + test_equal(&a, &b, false); + + let a = create_list_view_array::([None, Some(vec![Some(1), Some(0)])]); + let b = create_list_view_array::([None, Some(vec![Some(1), None])]); + test_equal(&a, &b, false); +} + // Special test for List>. // This tests the equal_ranges kernel fn test_sliced_list_of_list_view() { @@ -871,6 +887,16 @@ fn test_large_list_view_array() { test_test_list_view_array::(); } +#[test] +fn test_list_view_array_sizes_and_child_nulls() { + test_list_view_sizes_and_child_nulls::(); +} + +#[test] +fn test_large_list_view_array_sizes_and_child_nulls() { + test_list_view_sizes_and_child_nulls::(); +} + #[test] fn test_nested_list_view_array() { test_sliced_list_of_list_view::(); diff --git a/arrow/tests/array_validation.rs b/arrow/tests/array_validation.rs index 6455a15b3eeb..70c1cf34fbd8 100644 --- a/arrow/tests/array_validation.rs +++ b/arrow/tests/array_validation.rs @@ -432,7 +432,7 @@ fn test_validate_large_list_view_negative_sizes() { #[test] #[should_panic( - expected = "Values length 4 is less than the length (3) multiplied by the value size (2) for FixedSizeList" + expected = "Values length 4 is less than the length + offset (3) multiplied by the value size (2) for FixedSizeList" )] fn test_validate_fixed_size_list() { // child has 4 elements, diff --git a/parquet-geospatial/Cargo.toml b/parquet-geospatial/Cargo.toml index 08d74ed7e963..ba28d6d28a03 100644 --- a/parquet-geospatial/Cargo.toml +++ b/parquet-geospatial/Cargo.toml @@ -20,7 +20,6 @@ name = "parquet-geospatial" version = { workspace = true } license = { workspace = true } description = "Apache Parquet Geometry and Geography implementation in Rust" -homepage = { workspace = true } repository = { workspace = true } authors = { workspace = true } keywords = ["arrow", "parquet", "geometry", "geography"] diff --git a/parquet-geospatial/src/interval.rs b/parquet-geospatial/src/interval.rs index 97a84ec00c7b..5ec05859d9cb 100644 --- a/parquet-geospatial/src/interval.rs +++ b/parquet-geospatial/src/interval.rs @@ -520,8 +520,6 @@ impl IntervalTrait for WraparoundInterval { #[cfg(test)] mod test { - use core::f64; - use super::*; fn test_empty(empty: T) { diff --git a/parquet-variant-compute/Cargo.toml b/parquet-variant-compute/Cargo.toml index 7505d3f45c3d..d803b5f47bbb 100644 --- a/parquet-variant-compute/Cargo.toml +++ b/parquet-variant-compute/Cargo.toml @@ -20,7 +20,6 @@ name = "parquet-variant-compute" version = { workspace = true } license = { workspace = true } description = "Apache Parquet Variant Batch Processing" -homepage = { workspace = true } repository = { workspace = true } authors = { workspace = true } keywords = ["arrow", "parquet", "variant"] diff --git a/parquet-variant-compute/benches/variant_kernels.rs b/parquet-variant-compute/benches/variant_kernels.rs index 5a8862f1a963..0d35d689018f 100644 --- a/parquet-variant-compute/benches/variant_kernels.rs +++ b/parquet-variant-compute/benches/variant_kernels.rs @@ -15,14 +15,16 @@ // specific language governing permissions and limitations // under the License. -use arrow::array::{Array, ArrayRef, BinaryViewArray, BinaryViewBuilder, StringArray, StructArray}; +use arrow::array::{ + Array, ArrayRef, BinaryViewArray, BinaryViewBuilder, Int32Array, StringArray, StructArray, +}; use arrow::buffer::Buffer; use arrow_schema::{DataType, Field, FieldRef, Fields}; use chrono::{DateTime, NaiveDate, NaiveDateTime, NaiveTime}; use criterion::{BatchSize, Criterion, criterion_group, criterion_main}; use parquet_variant::{ EMPTY_VARIANT_METADATA_BYTES, Variant, VariantBuilder, VariantBuilderExt, VariantDecimal8, - VariantPath, + VariantPath, VariantPathElement, }; use parquet_variant_compute::{ GetOptions, VariantArray, VariantArrayBuilder, json_to_variant, shred_variant, variant_get, @@ -178,6 +180,55 @@ pub fn variant_get_bench(c: &mut Criterion) { }); } +pub fn variant_get_list_index_bench(c: &mut Criterion) { + let mut builder = VariantArrayBuilder::new(8192); + for i in 0..8192 { + let mut list = builder.new_list(); + // Alternate empty lists with two-element lists to exercise missing and present indexes. + if i % 2 != 0 { + list.append_value(Variant::Int64(i)); + list.append_value(Variant::Int64(i + 1)); + } + list.finish(); + } + let variant_array = builder.build(); + let list_type = DataType::List(Arc::new(Field::new("item", DataType::Int64, true))); + let shredded = shred_variant(&variant_array, &list_type).unwrap(); + let options = GetOptions::new_with_path(VariantPath::from(0)); + + let input = ArrayRef::from(variant_array); + c.bench_function("variant_get_list_index_unshredded", |b| { + b.iter(|| variant_get(&input, options.clone()).unwrap()) + }); + + let input = ArrayRef::from(shredded); + c.bench_function("variant_get_list_index_shredded", |b| { + b.iter(|| variant_get(&input, options.clone()).unwrap()) + }); +} + +pub fn variant_get_list_index_all_oob_int64_bench(c: &mut Criterion) { + let mut builder = VariantArrayBuilder::new(64); + for _ in 0..64 { + let mut list = builder.new_list(); + list.append_value(Variant::Int64(1)); + list.finish(); + } + let variant_array = builder.build(); + let list_type = DataType::List(Arc::new(Field::new("item", DataType::Int64, true))); + let input = ArrayRef::from(shred_variant(&variant_array, &list_type).unwrap()); + let options = GetOptions::new_with_path(VariantPath::from(9)) + .with_as_type(Some(Arc::new(Field::new("value", DataType::Int64, true)))); + let result = variant_get(&input, options.clone()).unwrap(); + assert_eq!(result.data_type(), &DataType::Int64); + assert_eq!(result.len(), 64); + assert_eq!(result.null_count(), 64); + + c.bench_function("variant_get_list_index_all_oob_int64_64_rows", |b| { + b.iter(|| variant_get(&input, options.clone()).unwrap()) + }); +} + pub fn variant_get_shredded_utf8_bench(c: &mut Criterion) { let variant_array = create_shredded_utf8_variant_array(8192); let input = ArrayRef::from(variant_array); @@ -198,8 +249,13 @@ pub fn variant_get_unshredded_object_path_bench(c: &mut Criterion) { let variant_array = create_unshredded_object_variant_array(VARIANT_GET_UNSHREDDED_OBJECT_ROWS); let input = ArrayRef::from(variant_array); let field: FieldRef = Arc::new(Field::new("typed_value", DataType::Int32, true)); - let options = GetOptions::new_with_path(VariantPath::try_from("attr.140").unwrap()) - .with_as_type(Some(field)); + // The dot is part of the field name, not a separator between nested fields. + let path = VariantPath::from(vec![VariantPathElement::field("attr.140")]); + let options = GetOptions::new_with_path(path).with_as_type(Some(field)); + + let result = variant_get(&input, options.clone()).unwrap(); + let result = result.as_any().downcast_ref::().unwrap(); + assert!(result.iter().all(|value| value == Some(140))); c.bench_function("variant_get_unshredded_object_path_262k_rows", |b| { b.iter(|| variant_get(&input, options.clone()).unwrap()) @@ -467,6 +523,8 @@ pub fn variant_get_binary_from_string_bench(c: &mut Criterion) { criterion_group!( benches, variant_get_bench, + variant_get_list_index_bench, + variant_get_list_index_all_oob_int64_bench, variant_get_shredded_utf8_bench, variant_get_unshredded_object_path_bench, shred_variant_partial_object_bench, diff --git a/parquet-variant-compute/src/shred_variant.rs b/parquet-variant-compute/src/shred_variant.rs index 2ad0342ed21b..2d78bc9014e2 100644 --- a/parquet-variant-compute/src/shred_variant.rs +++ b/parquet-variant-compute/src/shred_variant.rs @@ -17,6 +17,7 @@ //! Module for shredding VariantArray with a given schema. +use crate::type_conversion::VariantCastMode; use crate::variant_array::{ShreddedVariantFieldArray, StructArrayBuilder}; use crate::variant_to_arrow::{ ArrayVariantToArrowRowBuilder, PrimitiveVariantToArrowRowBuilder, @@ -87,7 +88,7 @@ pub(crate) fn shred_variant_with_options( cast_options, array.len(), NullValue::TopLevelVariant, - true, + VariantCastMode::Shred, )?; for i in 0..array.len() { if array.is_null(i) { @@ -141,7 +142,7 @@ pub(crate) fn make_variant_to_shredded_variant_arrow_row_builder<'a>( cast_options: &'a CastOptions, capacity: usize, null_value: NullValue, - shred: bool, + cast_mode: VariantCastMode, ) -> Result> { let builder = match data_type { DataType::Struct(fields) => { @@ -150,7 +151,7 @@ pub(crate) fn make_variant_to_shredded_variant_arrow_row_builder<'a>( cast_options, capacity, null_value, - shred, + cast_mode, )?; VariantToShreddedVariantRowBuilder::Object(typed_value_builder) } @@ -190,7 +191,7 @@ pub(crate) fn make_variant_to_shredded_variant_arrow_row_builder<'a>( | DataType::FixedSizeBinary(16) // UUID => { let builder = - make_primitive_variant_to_arrow_row_builder(data_type, cast_options, capacity, shred)?; + make_primitive_variant_to_arrow_row_builder(data_type, cast_options, capacity, cast_mode)?; let typed_value_builder = VariantToShreddedPrimitiveVariantRowBuilder::new(builder, capacity, null_value); VariantToShreddedVariantRowBuilder::Primitive(typed_value_builder) @@ -368,7 +369,7 @@ impl<'a> VariantToShreddedObjectVariantRowBuilder<'a> { cast_options: &'a CastOptions, capacity: usize, null_value: NullValue, - shred: bool, + cast_mode: VariantCastMode, ) -> Result { let typed_value_builders = fields.iter().map(|field| { let builder = make_variant_to_shredded_variant_arrow_row_builder( @@ -376,7 +377,7 @@ impl<'a> VariantToShreddedObjectVariantRowBuilder<'a> { cast_options, capacity, NullValue::ObjectField, - shred, + cast_mode, )?; Ok((field.name().as_str(), builder)) }); @@ -1082,7 +1083,7 @@ mod tests { &cast_options, 1, mode, - true, + VariantCastMode::Shred, ) .unwrap(); primitive_builder.append_null().unwrap(); @@ -1113,7 +1114,7 @@ mod tests { &cast_options, 1, mode, - true, + VariantCastMode::Shred, ) .unwrap(); array_builder.append_null().unwrap(); @@ -1142,7 +1143,7 @@ mod tests { &cast_options, 1, mode, - true, + VariantCastMode::Shred, ) .unwrap(); object_builder.append_null().unwrap(); @@ -1172,7 +1173,7 @@ mod tests { fn test_already_shredded_input_error() { // Create a VariantArray that already has typed_value_field // First create a valid VariantArray, then extract its parts to construct a shredded one - let temp_array = VariantArray::from_iter(vec![Some(Variant::from("test"))]); + let temp_array = VariantArray::from_iter([Some("test")]); let metadata = temp_array.metadata_column().clone(); let value = temp_array.value_column().clone(); let typed_value = Arc::new(Int64Array::from(vec![42])) as ArrayRef; @@ -1206,7 +1207,7 @@ mod tests { fn test_invalid_fixed_size_binary_shredding() { let mock_uuid_1 = Uuid::new_v4(); - let input = VariantArray::from_iter([Some(Variant::from(mock_uuid_1)), None]); + let input = VariantArray::from_iter([Some(mock_uuid_1), None]); // shred_variant only supports FixedSizeBinary(16). Any other length will err. let err = shred_variant(&input, &DataType::FixedSizeBinary(17)).unwrap_err(); @@ -1498,7 +1499,7 @@ mod tests { #[test] fn test_invalid_shredded_types_rejected() { - let input = VariantArray::from_iter([Variant::from(42)]); + let input = VariantArray::from_iter([Some(42)]); let invalid_types = vec![ DataType::UInt8, diff --git a/parquet-variant-compute/src/type_conversion.rs b/parquet-variant-compute/src/type_conversion.rs index b3d3b7094f0c..2882b70bbe67 100644 --- a/parquet-variant-compute/src/type_conversion.rs +++ b/parquet-variant-compute/src/type_conversion.rs @@ -34,16 +34,25 @@ use half::f16; use num_traits::NumCast; use parquet_variant::{Variant, VariantDecimal4, VariantDecimal8, VariantDecimal16}; +/// Controls the conversions allowed when extracting a typed value from a variant. +#[derive(Debug, Clone, Copy)] +pub(crate) enum VariantCastMode { + /// Use the conversions permitted by the variant shredding specification. + Shred, + /// Allow casts when extracting a value with `variant_get`. + Get, +} + /// Extension trait for Arrow primitive types that can extract their native value from a Variant pub(crate) trait PrimitiveFromVariant: ArrowPrimitiveType { - fn from_variant(variant: &Variant<'_, '_>, shred: bool) -> Option; + fn from_variant(variant: &Variant<'_, '_>, cast_mode: VariantCastMode) -> Option; } /// Extension trait for Arrow timestamp types that can extract their native value from a Variant /// We can't use [`PrimitiveFromVariant`] directly because we need _two_ implementations for each /// timestamp type -- the `NTZ` param here. pub(crate) trait TimestampFromVariant: ArrowTimestampType { - fn from_variant(variant: &Variant<'_, '_>, shred: bool) -> Option; + fn from_variant(variant: &Variant<'_, '_>, cast_mode: VariantCastMode) -> Option; } /// Cast a single `Variant` value with safe/strict semantics. @@ -71,10 +80,10 @@ pub(crate) fn variant_cast_with_options<'a, 'm, 'v, T>( macro_rules! impl_primitive_from_variant { ($arrow_type:ty, $shred_fun:expr, $get_method:ident $(, $cast_fn:expr)?) => { impl PrimitiveFromVariant for $arrow_type { - fn from_variant(variant: &Variant<'_, '_>, shred: bool) -> Option { - let value = match shred { - true => $shred_fun(variant), - false => $get_method(variant), + fn from_variant(variant: &Variant<'_, '_>, cast_mode: VariantCastMode) -> Option { + let value = match cast_mode { + VariantCastMode::Shred => $shred_fun(variant), + VariantCastMode::Get => $get_method(variant), }; $( let value = value.and_then($cast_fn); )? value @@ -86,10 +95,13 @@ macro_rules! impl_primitive_from_variant { macro_rules! impl_timestamp_from_variant { ($timestamp_type:ty, $shred_fun:expr, $variant_method:expr, ntz=$ntz:ident, $cast_fn:expr $(,)?) => { impl TimestampFromVariant<{ $ntz }> for $timestamp_type { - fn from_variant(variant: &Variant<'_, '_>, shred: bool) -> Option { - let value = match shred { - true => ($shred_fun)(variant), - false => $variant_method(variant), + fn from_variant( + variant: &Variant<'_, '_>, + cast_mode: VariantCastMode, + ) -> Option { + let value = match cast_mode { + VariantCastMode::Shred => ($shred_fun)(variant), + VariantCastMode::Get => $variant_method(variant), }; value.and_then($cast_fn) @@ -685,8 +697,11 @@ impl ShredDecimalVariant for Decimal256Type { } } -pub(crate) fn variant_to_boolean(variant: &Variant<'_, '_>, shred: bool) -> Option { - if shred { +pub(crate) fn variant_to_boolean( + variant: &Variant<'_, '_>, + cast_mode: VariantCastMode, +) -> Option { + if matches!(cast_mode, VariantCastMode::Shred) { return variant.as_boolean(); } diff --git a/parquet-variant-compute/src/unshred_variant.rs b/parquet-variant-compute/src/unshred_variant.rs index 40a22a3f5491..cc414dadc39d 100644 --- a/parquet-variant-compute/src/unshred_variant.rs +++ b/parquet-variant-compute/src/unshred_variant.rs @@ -857,9 +857,7 @@ mod tests { #[test] fn test_shred_unshred_round_trip_annotates_value_non_nullable() { - let mut builder = VariantArrayBuilder::new(1); - builder.append_variant(Variant::from(42i64)); - let original = builder.build(); + let original = VariantArray::from_iter([Some(42_i64)]); assert!(!value_field_is_nullable(&original)); let shredded = shred_variant(&original, &DataType::Int64).unwrap(); @@ -892,10 +890,7 @@ mod tests { #[test] fn test_unshred_already_unshredded_reannotates_nullable_value() { - let mut builder = VariantArrayBuilder::new(2); - builder.append_variant(Variant::from(42i64)); - builder.append_null(); - let original = builder.build(); + let original = VariantArray::from_iter([Some(42_i64), None]); // same data, but with the out-of-spec nullable `value` annotation let nullable_input = VariantArray::from_parts( @@ -965,18 +960,14 @@ mod tests { /// shape has its own expansion of `handle_unshredded_case`. #[test] fn test_unshred_missing_row_for_decimal_timestamp_object_list() { - let mut builder = VariantArrayBuilder::new(2); - builder.append_variant(Variant::from(VariantDecimal8::try_new(1234, 2).unwrap())); - builder.append_null(); - assert_missing_row_unshreds_to_variant_null(&builder.build(), &DataType::Decimal64(18, 2)); + let decimals = + VariantArray::from_iter([Some(VariantDecimal8::try_new(1234, 2).unwrap()), None]); + assert_missing_row_unshreds_to_variant_null(&decimals, &DataType::Decimal64(18, 2)); - let mut builder = VariantArrayBuilder::new(2); - builder.append_variant(Variant::from( - chrono::DateTime::from_timestamp(1, 0).unwrap(), - )); - builder.append_null(); + let timestamps = + VariantArray::from_iter([Some(chrono::DateTime::from_timestamp(1, 0).unwrap()), None]); assert_missing_row_unshreds_to_variant_null( - &builder.build(), + ×tamps, &DataType::Timestamp(TimeUnit::Microsecond, Some("UTC".into())), ); @@ -1010,9 +1001,7 @@ mod tests { #[test] fn test_unshred_value_only_with_unmasked_nulls_materializes_variant_null() { - let mut builder = VariantArrayBuilder::new(1); - builder.append_variant(Variant::from(42i64)); - let single = builder.build(); + let single = VariantArray::from_iter([Some(42_i64)]); let metadata_bytes = single.metadata_column().as_binary_view().value(0); let value_bytes = single.value_column().as_binary_view().value(0); diff --git a/parquet-variant-compute/src/variant_array.rs b/parquet-variant-compute/src/variant_array.rs index da051fe549d3..cdf3436f3192 100644 --- a/parquet-variant-compute/src/variant_array.rs +++ b/parquet-variant-compute/src/variant_array.rs @@ -179,6 +179,31 @@ impl ExtensionType for VariantType { /// See the examples below from converting between `VariantArray` and /// `StructArray`. /// +/// # Example: collecting optional values +/// +/// `VariantArray` can collect optional values that convert into [`Variant`]. +/// `None` creates a null row, while `Some(Variant::Null)` creates a valid row +/// whose value is `Variant::Null`. +/// +/// ``` +/// # use parquet_variant::Variant; +/// # use parquet_variant_compute::VariantArray; +/// let values = [Some(42_i64), None, Some(-1_i64)]; +/// let array: VariantArray = values.into_iter().collect(); +/// +/// assert_eq!(array.value(0), Variant::Int64(42)); +/// assert!(array.is_null(1)); +/// ``` +/// +/// For an all-null array, specify the type when no `Some` value can infer it: +/// +/// ``` +/// # use parquet_variant::Variant; +/// # use parquet_variant_compute::VariantArray; +/// let null_rows = VariantArray::from_iter(vec![None::; 3]); +/// assert!(null_rows.is_null(0)); +/// ``` +/// /// [`VariantArrayBuilder`]: crate::VariantArrayBuilder /// /// # Documentation @@ -200,12 +225,9 @@ impl ExtensionType for VariantType { /// ``` /// # use arrow::array::StructArray; /// # use arrow_schema::{Schema, Field, DataType}; -/// # use parquet_variant::Variant; -/// # use parquet_variant_compute::{VariantArrayBuilder, VariantArray, VariantType}; +/// # use parquet_variant_compute::{VariantArray, VariantType}; /// # fn get_variant_array() -> VariantArray { -/// # let mut builder = VariantArrayBuilder::new(10); -/// # builder.append_variant(Variant::from("such wow")); -/// # builder.build() +/// # VariantArray::from_iter([Some("such wow")]) /// # } /// # fn get_schema() -> Schema { /// # Schema::new(vec![ @@ -228,12 +250,9 @@ impl ExtensionType for VariantType { /// /// ``` /// # use arrow_schema::{Schema, Field, DataType}; -/// # use parquet_variant::Variant; -/// # use parquet_variant_compute::{VariantArrayBuilder, VariantArray, VariantType}; +/// # use parquet_variant_compute::{VariantArray, VariantType}; /// # fn get_variant_array() -> VariantArray { -/// # let mut builder = VariantArrayBuilder::new(10); -/// # builder.append_variant(Variant::from("such wow")); -/// # builder.build() +/// # VariantArray::from_iter([Some("such wow")]) /// # } /// let variant_array = get_variant_array(); /// // First field is an integer id, second field is a variant @@ -248,12 +267,9 @@ impl ExtensionType for VariantType { /// /// ``` /// # use arrow_schema::{Schema, Field, DataType}; -/// # use parquet_variant::Variant; -/// # use parquet_variant_compute::{VariantArrayBuilder, VariantArray, VariantType}; +/// # use parquet_variant_compute::{VariantArray, VariantType}; /// # fn get_variant_array() -> VariantArray { -/// # let mut builder = VariantArrayBuilder::new(10); -/// # builder.append_variant(Variant::from("such wow")); -/// # builder.build() +/// # VariantArray::from_iter([Some("such wow")]) /// # } /// # let variant_array = get_variant_array(); /// // The DataType of a VariantArray varies depending on how it is shredded @@ -271,12 +287,9 @@ impl ExtensionType for VariantType { /// /// ``` /// # use arrow::array::StructArray; -/// # use parquet_variant::Variant; -/// # use parquet_variant_compute::VariantArrayBuilder; +/// # use parquet_variant_compute::VariantArray; /// // Create Variant Array -/// let mut builder = VariantArrayBuilder::new(10); -/// builder.append_variant(Variant::from("such wow")); -/// let variant_array = builder.build(); +/// let variant_array = VariantArray::from_iter([Some("such wow")]); /// // convert to StructArray /// let struct_array: StructArray = variant_array.into(); /// ``` @@ -286,11 +299,9 @@ impl ExtensionType for VariantType { /// ``` /// # use arrow::array::StructArray; /// # use parquet_variant::Variant; -/// # use parquet_variant_compute::{VariantArrayBuilder, VariantArray}; +/// # use parquet_variant_compute::VariantArray; /// # fn get_struct_array() -> StructArray { -/// # let mut builder = VariantArrayBuilder::new(10); -/// # builder.append_variant(Variant::from("such wow")); -/// # builder.build().into() +/// # VariantArray::from_iter([Some("such wow")]).into() /// # } /// let struct_array: StructArray = get_struct_array(); /// // try and create a VariantArray from it @@ -615,9 +626,9 @@ impl From for ArrayRef { } } -impl<'m, 'v> FromIterator>> for VariantArray { - fn from_iter>>>(iter: T) -> Self { - let iter = iter.into_iter(); +impl<'m, 'v, V: Into>> FromIterator> for VariantArray { + fn from_iter>>(iter: T) -> Self { + let iter = iter.into_iter().map(|value| value.map(Into::into)); let mut b = VariantArrayBuilder::new(iter.size_hint().0); b.extend(iter); @@ -1801,15 +1812,7 @@ mod test { #[test] fn test_variant_array_iter_double_ended() { - let mut b = VariantArrayBuilder::new(5); - - b.append_variant(Variant::from(0_i32)); - b.append_null(); - b.append_variant(Variant::from(2_i32)); - b.append_null(); - b.append_variant(Variant::from(4_i32)); - - let array = b.build(); + let array = VariantArray::from_iter([Some(0_i32), None, Some(2_i32), None, Some(4_i32)]); let mut iter = array.iter(); assert_eq!(iter.next(), Some(Some(Variant::from(0_i32)))); @@ -1825,15 +1828,7 @@ mod test { #[test] fn test_variant_array_iter_reverse() { - let mut b = VariantArrayBuilder::new(5); - - b.append_variant(Variant::from("a")); - b.append_null(); - b.append_variant(Variant::from("aaa")); - b.append_null(); - b.append_variant(Variant::from("aaaaa")); - - let array = b.build(); + let array = VariantArray::from_iter([Some("a"), None, Some("aaa"), None, Some("aaaaa")]); let result: Vec<_> = array.iter().rev().collect(); assert_eq!( @@ -1875,6 +1870,56 @@ mod test { assert!(variant_array.is_null(3)); } + #[test] + fn test_from_option_into_variants_into_variant_array() { + // Items that convert cleanly to `Variant` can be collected directly, + // without wrapping them in `Some(Variant::from(..))` first + let v = vec![Some(42_i64), None, Some(-1_i64)]; + + let variant_array = VariantArray::from_iter(v); + + assert_eq!(variant_array.len(), 3); + + assert!(!variant_array.is_null(0)); + assert_eq!(variant_array.value(0), Variant::Int64(42)); + + assert!(variant_array.is_null(1)); + + assert!(!variant_array.is_null(2)); + assert_eq!(variant_array.value(2), Variant::Int64(-1)); + } + + #[test] + fn test_from_option_str_into_variant_array() { + let v = vec![ + Some("hello"), + None, + Some( + "hello world this is much longer than the maximum short string length of sixty three bytes", + ), + ]; + + let variant_array: VariantArray = v.into_iter().collect(); + + assert_eq!(variant_array.len(), 3); + + assert!(!variant_array.is_null(0)); + assert_eq!( + variant_array.value(0), + Variant::ShortString(ShortString::try_new("hello").unwrap()) + ); + + assert!(variant_array.is_null(1)); + + assert!(!variant_array.is_null(2)); + assert_eq!( + variant_array.value(2), + Variant::String( + "hello world this is much longer than the maximum short string length of sixty three bytes" + ) + ); + } + #[test] fn test_from_variants_into_variant_array() { let v = vec![ diff --git a/parquet-variant-compute/src/variant_get.rs b/parquet-variant-compute/src/variant_get.rs index 63af0dad08ee..96d634d5161e 100644 --- a/parquet-variant-compute/src/variant_get.rs +++ b/parquet-variant-compute/src/variant_get.rs @@ -3023,9 +3023,7 @@ mod test { #[test] fn test_error_message_boolean_type_display() { - let mut builder = VariantArrayBuilder::new(1); - builder.append_variant(Variant::from("abcd")); - let variant_array: ArrayRef = ArrayRef::from(builder.build()); + let variant_array: ArrayRef = ArrayRef::from(VariantArray::from_iter([Some("abcd")])); // Request Boolean with strict casting to force an error let options = GetOptions { @@ -3044,9 +3042,7 @@ mod test { #[test] fn test_error_message_numeric_type_display() { - let mut builder = VariantArrayBuilder::new(1); - builder.append_variant(Variant::from("abcd")); - let variant_array: ArrayRef = ArrayRef::from(builder.build()); + let variant_array: ArrayRef = ArrayRef::from(VariantArray::from_iter([Some("abcd")])); // Request Float32 with strict casting to force an error let options = GetOptions { @@ -3065,9 +3061,7 @@ mod test { #[test] fn test_error_message_temporal_type_display() { - let mut builder = VariantArrayBuilder::new(1); - builder.append_variant(Variant::BooleanFalse); - let variant_array: ArrayRef = ArrayRef::from(builder.build()); + let variant_array: ArrayRef = ArrayRef::from(VariantArray::from_iter([Some(false)])); // Request Timestamp with strict casting to force an error let options = GetOptions { @@ -3939,15 +3933,15 @@ mod test { #[test] fn get_decimal32_scale_down_rounding() { - let mut builder = crate::VariantArrayBuilder::new(7); - builder.append_variant(VariantDecimal4::try_new(1235, 0).unwrap().into()); - builder.append_variant(VariantDecimal4::try_new(1245, 0).unwrap().into()); - builder.append_variant(VariantDecimal4::try_new(-1235, 0).unwrap().into()); - builder.append_variant(VariantDecimal4::try_new(-1245, 0).unwrap().into()); - builder.append_variant(VariantDecimal4::try_new(1235, 2).unwrap().into()); // 12.35 rounded down to 10 for scale -1 - builder.append_variant(VariantDecimal4::try_new(1235, 3).unwrap().into()); // 1.235 rounded down to 0 for scale -1 - builder.append_variant(VariantDecimal4::try_new(5235, 3).unwrap().into()); // 5.235 rounded up to 10 for scale -1 - let variant_array: ArrayRef = ArrayRef::from(builder.build()); + let variant_array: ArrayRef = ArrayRef::from(VariantArray::from_iter([ + Some(VariantDecimal4::try_new(1235, 0).unwrap()), + Some(VariantDecimal4::try_new(1245, 0).unwrap()), + Some(VariantDecimal4::try_new(-1235, 0).unwrap()), + Some(VariantDecimal4::try_new(-1245, 0).unwrap()), + Some(VariantDecimal4::try_new(1235, 2).unwrap()), // 12.35 rounded down to 10 for scale -1 + Some(VariantDecimal4::try_new(1235, 3).unwrap()), // 1.235 rounded down to 0 for scale -1 + Some(VariantDecimal4::try_new(5235, 3).unwrap()), // 5.235 rounded up to 10 for scale -1 + ])); let field = Field::new("result", DataType::Decimal32(9, -1), true); let options = GetOptions::new().with_as_type(Some(FieldRef::from(field))); @@ -3968,18 +3962,10 @@ mod test { #[test] fn get_decimal32_large_scale_reduction() { - let mut builder = crate::VariantArrayBuilder::new(2); - builder.append_variant( - VariantDecimal4::try_new(-VariantDecimal4::MAX_UNSCALED_VALUE, 0) - .unwrap() - .into(), - ); - builder.append_variant( - VariantDecimal4::try_new(VariantDecimal4::MAX_UNSCALED_VALUE, 0) - .unwrap() - .into(), - ); - let variant_array: ArrayRef = ArrayRef::from(builder.build()); + let variant_array: ArrayRef = ArrayRef::from(VariantArray::from_iter([ + Some(VariantDecimal4::try_new(-VariantDecimal4::MAX_UNSCALED_VALUE, 0).unwrap()), + Some(VariantDecimal4::try_new(VariantDecimal4::MAX_UNSCALED_VALUE, 0).unwrap()), + ])); let field = Field::new("result", DataType::Decimal32(9, -9), true); let options = GetOptions::new().with_as_type(Some(FieldRef::from(field))); @@ -4007,18 +3993,11 @@ mod test { #[test] fn get_decimal32_precision_overflow_safe() { // Exceed Decimal32 after scaling and rounding - let mut builder = crate::VariantArrayBuilder::new(2); - builder.append_variant( - VariantDecimal4::try_new(VariantDecimal4::MAX_UNSCALED_VALUE, 0) - .unwrap() - .into(), - ); - builder.append_variant( - VariantDecimal4::try_new(VariantDecimal4::MAX_UNSCALED_VALUE, 9) - .unwrap() - .into(), - ); // integer value round up overflows - let variant_array: ArrayRef = ArrayRef::from(builder.build()); + let variant_array: ArrayRef = ArrayRef::from(VariantArray::from_iter([ + Some(VariantDecimal4::try_new(VariantDecimal4::MAX_UNSCALED_VALUE, 0).unwrap()), + // integer value round up overflows + Some(VariantDecimal4::try_new(VariantDecimal4::MAX_UNSCALED_VALUE, 9).unwrap()), + ])); let field = Field::new("result", DataType::Decimal32(2, 2), true); let options = GetOptions::new().with_as_type(Some(FieldRef::from(field))); @@ -4031,13 +4010,9 @@ mod test { #[test] fn get_decimal32_precision_overflow_unsafe_errors() { - let mut builder = crate::VariantArrayBuilder::new(1); - builder.append_variant( - VariantDecimal4::try_new(VariantDecimal4::MAX_UNSCALED_VALUE, 0) - .unwrap() - .into(), - ); - let variant_array: ArrayRef = ArrayRef::from(builder.build()); + let variant_array: ArrayRef = ArrayRef::from(VariantArray::from_iter([Some( + VariantDecimal4::try_new(VariantDecimal4::MAX_UNSCALED_VALUE, 0).unwrap(), + )])); let field = Field::new("result", DataType::Decimal32(9, 2), true); let cast_options = CastOptions { @@ -4089,15 +4064,15 @@ mod test { #[test] fn get_decimal64_scale_down_rounding() { - let mut builder = crate::VariantArrayBuilder::new(7); - builder.append_variant(VariantDecimal8::try_new(1235, 0).unwrap().into()); - builder.append_variant(VariantDecimal8::try_new(1245, 0).unwrap().into()); - builder.append_variant(VariantDecimal8::try_new(-1235, 0).unwrap().into()); - builder.append_variant(VariantDecimal8::try_new(-1245, 0).unwrap().into()); - builder.append_variant(VariantDecimal8::try_new(1235, 2).unwrap().into()); // 12.35 rounded down to 10 for scale -1 - builder.append_variant(VariantDecimal8::try_new(1235, 3).unwrap().into()); // 1.235 rounded down to 0 for scale -1 - builder.append_variant(VariantDecimal8::try_new(5235, 3).unwrap().into()); // 5.235 rounded up to 10 for scale -1 - let variant_array: ArrayRef = ArrayRef::from(builder.build()); + let variant_array: ArrayRef = ArrayRef::from(VariantArray::from_iter([ + Some(VariantDecimal8::try_new(1235, 0).unwrap()), + Some(VariantDecimal8::try_new(1245, 0).unwrap()), + Some(VariantDecimal8::try_new(-1235, 0).unwrap()), + Some(VariantDecimal8::try_new(-1245, 0).unwrap()), + Some(VariantDecimal8::try_new(1235, 2).unwrap()), // 12.35 rounded down to 10 for scale -1 + Some(VariantDecimal8::try_new(1235, 3).unwrap()), // 1.235 rounded down to 0 for scale -1 + Some(VariantDecimal8::try_new(5235, 3).unwrap()), // 5.235 rounded up to 10 for scale -1 + ])); let field = Field::new("result", DataType::Decimal64(18, -1), true); let options = GetOptions::new().with_as_type(Some(FieldRef::from(field))); @@ -4118,18 +4093,10 @@ mod test { #[test] fn get_decimal64_large_scale_reduction() { - let mut builder = crate::VariantArrayBuilder::new(2); - builder.append_variant( - VariantDecimal8::try_new(-VariantDecimal8::MAX_UNSCALED_VALUE, 0) - .unwrap() - .into(), - ); - builder.append_variant( - VariantDecimal8::try_new(VariantDecimal8::MAX_UNSCALED_VALUE, 0) - .unwrap() - .into(), - ); - let variant_array: ArrayRef = ArrayRef::from(builder.build()); + let variant_array: ArrayRef = ArrayRef::from(VariantArray::from_iter([ + Some(VariantDecimal8::try_new(-VariantDecimal8::MAX_UNSCALED_VALUE, 0).unwrap()), + Some(VariantDecimal8::try_new(VariantDecimal8::MAX_UNSCALED_VALUE, 0).unwrap()), + ])); let field = Field::new("result", DataType::Decimal64(18, -18), true); let options = GetOptions::new().with_as_type(Some(FieldRef::from(field))); @@ -4157,18 +4124,11 @@ mod test { #[test] fn get_decimal64_precision_overflow_safe() { // Exceed Decimal64 after scaling and rounding - let mut builder = crate::VariantArrayBuilder::new(2); - builder.append_variant( - VariantDecimal8::try_new(VariantDecimal8::MAX_UNSCALED_VALUE, 0) - .unwrap() - .into(), - ); - builder.append_variant( - VariantDecimal8::try_new(VariantDecimal8::MAX_UNSCALED_VALUE, 18) - .unwrap() - .into(), - ); // integer value round up overflows - let variant_array: ArrayRef = ArrayRef::from(builder.build()); + let variant_array: ArrayRef = ArrayRef::from(VariantArray::from_iter([ + Some(VariantDecimal8::try_new(VariantDecimal8::MAX_UNSCALED_VALUE, 0).unwrap()), + // integer value round up overflows + Some(VariantDecimal8::try_new(VariantDecimal8::MAX_UNSCALED_VALUE, 18).unwrap()), + ])); let field = Field::new("result", DataType::Decimal64(2, 2), true); let options = GetOptions::new().with_as_type(Some(FieldRef::from(field))); @@ -4181,13 +4141,9 @@ mod test { #[test] fn get_decimal64_precision_overflow_unsafe_errors() { - let mut builder = crate::VariantArrayBuilder::new(1); - builder.append_variant( - VariantDecimal8::try_new(VariantDecimal8::MAX_UNSCALED_VALUE, 0) - .unwrap() - .into(), - ); - let variant_array: ArrayRef = ArrayRef::from(builder.build()); + let variant_array: ArrayRef = ArrayRef::from(VariantArray::from_iter([Some( + VariantDecimal8::try_new(VariantDecimal8::MAX_UNSCALED_VALUE, 0).unwrap(), + )])); let field = Field::new("result", DataType::Decimal64(18, 2), true); let cast_options = CastOptions { @@ -4208,12 +4164,12 @@ mod test { #[test] fn get_decimal128_rescaled_to_scale2() { - let mut builder = crate::VariantArrayBuilder::new(4); - builder.append_variant(VariantDecimal16::try_new(1234, 2).unwrap().into()); - builder.append_variant(VariantDecimal16::try_new(1234, 3).unwrap().into()); - builder.append_variant(VariantDecimal16::try_new(1234, 0).unwrap().into()); - builder.append_null(); - let variant_array: ArrayRef = ArrayRef::from(builder.build()); + let variant_array: ArrayRef = ArrayRef::from(VariantArray::from_iter([ + Some(VariantDecimal16::try_new(1234, 2).unwrap()), + Some(VariantDecimal16::try_new(1234, 3).unwrap()), + Some(VariantDecimal16::try_new(1234, 0).unwrap()), + None, + ])); let field = Field::new("result", DataType::Decimal128(38, 2), true); let options = GetOptions::new().with_as_type(Some(FieldRef::from(field))); @@ -4230,15 +4186,15 @@ mod test { #[test] fn get_decimal128_scale_down_rounding() { - let mut builder = crate::VariantArrayBuilder::new(7); - builder.append_variant(VariantDecimal16::try_new(1235, 0).unwrap().into()); - builder.append_variant(VariantDecimal16::try_new(1245, 0).unwrap().into()); - builder.append_variant(VariantDecimal16::try_new(-1235, 0).unwrap().into()); - builder.append_variant(VariantDecimal16::try_new(-1245, 0).unwrap().into()); - builder.append_variant(VariantDecimal16::try_new(1235, 2).unwrap().into()); // 12.35 rounded down to 10 for scale -1 - builder.append_variant(VariantDecimal16::try_new(1235, 3).unwrap().into()); // 1.235 rounded down to 0 for scale -1 - builder.append_variant(VariantDecimal16::try_new(5235, 3).unwrap().into()); // 5.235 rounded up to 10 for scale -1 - let variant_array: ArrayRef = ArrayRef::from(builder.build()); + let variant_array: ArrayRef = ArrayRef::from(VariantArray::from_iter([ + Some(VariantDecimal16::try_new(1235, 0).unwrap()), + Some(VariantDecimal16::try_new(1245, 0).unwrap()), + Some(VariantDecimal16::try_new(-1235, 0).unwrap()), + Some(VariantDecimal16::try_new(-1245, 0).unwrap()), + Some(VariantDecimal16::try_new(1235, 2).unwrap()), // 12.35 rounded down to 10 for scale -1 + Some(VariantDecimal16::try_new(1235, 3).unwrap()), // 1.235 rounded down to 0 for scale -1 + Some(VariantDecimal16::try_new(5235, 3).unwrap()), // 5.235 rounded up to 10 for scale -1 + ])); let field = Field::new("result", DataType::Decimal128(38, -1), true); let options = GetOptions::new().with_as_type(Some(FieldRef::from(field))); @@ -4260,18 +4216,11 @@ mod test { #[test] fn get_decimal128_precision_overflow_safe() { // Exceed Decimal128 after scaling and rounding - let mut builder = crate::VariantArrayBuilder::new(2); - builder.append_variant( - VariantDecimal16::try_new(VariantDecimal16::MAX_UNSCALED_VALUE, 0) - .unwrap() - .into(), - ); - builder.append_variant( - VariantDecimal16::try_new(VariantDecimal16::MAX_UNSCALED_VALUE, 38) - .unwrap() - .into(), - ); // integer value round up overflows - let variant_array: ArrayRef = ArrayRef::from(builder.build()); + let variant_array: ArrayRef = ArrayRef::from(VariantArray::from_iter([ + Some(VariantDecimal16::try_new(VariantDecimal16::MAX_UNSCALED_VALUE, 0).unwrap()), + // integer value round up overflows + Some(VariantDecimal16::try_new(VariantDecimal16::MAX_UNSCALED_VALUE, 38).unwrap()), + ])); let field = Field::new("result", DataType::Decimal128(2, 2), true); let options = GetOptions::new().with_as_type(Some(FieldRef::from(field))); @@ -4284,13 +4233,9 @@ mod test { #[test] fn get_decimal128_precision_overflow_unsafe_errors() { - let mut builder = crate::VariantArrayBuilder::new(1); - builder.append_variant( - VariantDecimal16::try_new(VariantDecimal16::MAX_UNSCALED_VALUE, 0) - .unwrap() - .into(), - ); - let variant_array: ArrayRef = ArrayRef::from(builder.build()); + let variant_array: ArrayRef = ArrayRef::from(VariantArray::from_iter([Some( + VariantDecimal16::try_new(VariantDecimal16::MAX_UNSCALED_VALUE, 0).unwrap(), + )])); let field = Field::new("result", DataType::Decimal128(38, 2), true); let cast_options = CastOptions { @@ -4310,12 +4255,12 @@ mod test { #[test] fn get_decimal256_rescaled_to_scale2() { // Build unshredded variant values with different scales using Decimal16 source - let mut builder = crate::VariantArrayBuilder::new(4); - builder.append_variant(VariantDecimal16::try_new(1234, 2).unwrap().into()); // 12.34 - builder.append_variant(VariantDecimal16::try_new(1234, 3).unwrap().into()); // 1.234 - builder.append_variant(VariantDecimal16::try_new(1234, 0).unwrap().into()); // 1234 - builder.append_null(); - let variant_array: ArrayRef = ArrayRef::from(builder.build()); + let variant_array: ArrayRef = ArrayRef::from(VariantArray::from_iter([ + Some(VariantDecimal16::try_new(1234, 2).unwrap()), // 12.34 + Some(VariantDecimal16::try_new(1234, 3).unwrap()), // 1.234 + Some(VariantDecimal16::try_new(1234, 0).unwrap()), // 1234 + None, + ])); let field = Field::new("result", DataType::Decimal256(76, 2), true); let options = GetOptions::new().with_as_type(Some(FieldRef::from(field))); @@ -4332,15 +4277,15 @@ mod test { #[test] fn get_decimal256_scale_down_rounding() { - let mut builder = crate::VariantArrayBuilder::new(7); - builder.append_variant(VariantDecimal16::try_new(1235, 0).unwrap().into()); - builder.append_variant(VariantDecimal16::try_new(1245, 0).unwrap().into()); - builder.append_variant(VariantDecimal16::try_new(-1235, 0).unwrap().into()); - builder.append_variant(VariantDecimal16::try_new(-1245, 0).unwrap().into()); - builder.append_variant(VariantDecimal16::try_new(1235, 2).unwrap().into()); // 12.35 rounded down to 10 for scale -1 - builder.append_variant(VariantDecimal16::try_new(1235, 3).unwrap().into()); // 1.235 rounded down to 0 for scale -1 - builder.append_variant(VariantDecimal16::try_new(5235, 3).unwrap().into()); // 5.235 rounded up to 10 for scale -1 - let variant_array: ArrayRef = ArrayRef::from(builder.build()); + let variant_array: ArrayRef = ArrayRef::from(VariantArray::from_iter([ + Some(VariantDecimal16::try_new(1235, 0).unwrap()), + Some(VariantDecimal16::try_new(1245, 0).unwrap()), + Some(VariantDecimal16::try_new(-1235, 0).unwrap()), + Some(VariantDecimal16::try_new(-1245, 0).unwrap()), + Some(VariantDecimal16::try_new(1235, 2).unwrap()), // 12.35 rounded down to 10 for scale -1 + Some(VariantDecimal16::try_new(1235, 3).unwrap()), // 1.235 rounded down to 0 for scale -1 + Some(VariantDecimal16::try_new(5235, 3).unwrap()), // 5.235 rounded up to 10 for scale -1 + ])); let field = Field::new("result", DataType::Decimal256(76, -1), true); let options = GetOptions::new().with_as_type(Some(FieldRef::from(field))); @@ -4362,18 +4307,10 @@ mod test { #[test] fn get_decimal256_precision_overflow_safe() { // Exceed Decimal128 max precision (38) after scaling - let mut builder = crate::VariantArrayBuilder::new(2); - builder.append_variant( - VariantDecimal16::try_new(VariantDecimal16::MAX_UNSCALED_VALUE, 1) - .unwrap() - .into(), - ); - builder.append_variant( - VariantDecimal16::try_new(VariantDecimal16::MAX_UNSCALED_VALUE, 0) - .unwrap() - .into(), - ); - let variant_array: ArrayRef = ArrayRef::from(builder.build()); + let variant_array: ArrayRef = ArrayRef::from(VariantArray::from_iter([ + Some(VariantDecimal16::try_new(VariantDecimal16::MAX_UNSCALED_VALUE, 1).unwrap()), + Some(VariantDecimal16::try_new(VariantDecimal16::MAX_UNSCALED_VALUE, 0).unwrap()), + ])); let field = Field::new("result", DataType::Decimal256(76, 39), true); let options = GetOptions::new().with_as_type(Some(FieldRef::from(field))); @@ -4394,18 +4331,10 @@ mod test { #[test] fn get_decimal256_precision_overflow_unsafe_errors() { // Exceed Decimal128 max precision (38) after scaling - let mut builder = crate::VariantArrayBuilder::new(2); - builder.append_variant( - VariantDecimal16::try_new(VariantDecimal16::MAX_UNSCALED_VALUE, 1) - .unwrap() - .into(), - ); - builder.append_variant( - VariantDecimal16::try_new(VariantDecimal16::MAX_UNSCALED_VALUE, 0) - .unwrap() - .into(), - ); - let variant_array: ArrayRef = ArrayRef::from(builder.build()); + let variant_array: ArrayRef = ArrayRef::from(VariantArray::from_iter([ + Some(VariantDecimal16::try_new(VariantDecimal16::MAX_UNSCALED_VALUE, 1).unwrap()), + Some(VariantDecimal16::try_new(VariantDecimal16::MAX_UNSCALED_VALUE, 0).unwrap()), + ])); let field = Field::new("result", DataType::Decimal256(76, 39), true); let cast_options = CastOptions { @@ -4452,11 +4381,11 @@ mod test { #[test] fn get_variant_as_dictionary() { - let variant_array: ArrayRef = ArrayRef::from(VariantArray::from_iter(vec![ - Some(Variant::from("apple")), - Some(Variant::from("banana")), + let variant_array: ArrayRef = ArrayRef::from(VariantArray::from_iter([ + Some("apple"), + Some("banana"), None, - Some(Variant::from("apple")), + Some("apple"), ])); let data_type = DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)); let options = GetOptions::new().with_as_type(Some(FieldRef::from(Field::new( @@ -4475,12 +4404,8 @@ mod test { #[test] fn get_variant_as_numeric_dictionary() { - let variant_array: ArrayRef = ArrayRef::from(VariantArray::from_iter(vec![ - Some(Variant::from(42)), - Some(Variant::from(7)), - None, - Some(Variant::from(42)), - ])); + let variant_array: ArrayRef = + ArrayRef::from(VariantArray::from_iter([Some(42), Some(7), None, Some(42)])); let data_type = DataType::Dictionary(Box::new(DataType::Int16), Box::new(DataType::Int32)); let options = GetOptions::new().with_as_type(Some(FieldRef::from(Field::new( "dict", @@ -4498,12 +4423,12 @@ mod test { #[test] fn get_variant_as_run_end_encoded() { - let variant_array: ArrayRef = ArrayRef::from(VariantArray::from_iter(vec![ - Some(Variant::from("apple")), - Some(Variant::from("apple")), + let variant_array: ArrayRef = ArrayRef::from(VariantArray::from_iter([ + Some("apple"), + Some("apple"), None, - Some(Variant::from("banana")), - Some(Variant::from("banana")), + Some("banana"), + Some("banana"), ])); let run_ends = Arc::new(Field::new( Field::REE_RUN_ENDS_FIELD_DEFAULT_NAME, @@ -4695,12 +4620,8 @@ mod test { } fn invalid_time_variant_array() -> ArrayRef { - let mut builder = VariantArrayBuilder::new(3); // 86401000000 is invalid for Time64Microsecond (max is 86400000000) - builder.append_variant(Variant::Int64(86401000000)); - builder.append_variant(Variant::Int64(86401000000)); - builder.append_variant(Variant::Int64(86401000000)); - Arc::new(builder.build().into_inner()) + ArrayRef::from(VariantArray::from_iter([Some(86401000000_i64); 3])) } #[test] @@ -5459,10 +5380,7 @@ mod test { vec![Field::new("encoded", data_type.clone(), true)], ) .unwrap(); - let mut builder = VariantArrayBuilder::new(2); - builder.append_variant(Variant::from("apple")); - builder.append_variant(Variant::from("banana")); - let array = ArrayRef::from(builder.build()); + let array = ArrayRef::from(VariantArray::from_iter([Some("apple"), Some("banana")])); let options = union_get_options(&fields, UnionMode::Dense).with_cast_options(CastOptions { safe: false, @@ -5524,9 +5442,9 @@ mod test { ], ) .unwrap(); - let mut builder = VariantArrayBuilder::new(1); - builder.append_variant(VariantDecimal4::try_new(12_345, 2).unwrap().into()); - let array = ArrayRef::from(builder.build()); + let array = ArrayRef::from(VariantArray::from_iter([Some( + VariantDecimal4::try_new(12_345, 2).unwrap(), + )])); for safe in [true, false] { let options = @@ -5725,9 +5643,7 @@ mod test { #[test] fn get_variant_as_union_empty_fields_errors() { - let mut builder = VariantArrayBuilder::new(1); - builder.append_variant(Variant::Int8(1)); - let array = ArrayRef::from(builder.build()); + let array = ArrayRef::from(VariantArray::from_iter([Some(1_i8)])); let err = variant_get( &array, diff --git a/parquet-variant-compute/src/variant_to_arrow.rs b/parquet-variant-compute/src/variant_to_arrow.rs index 2c811b8b7ba1..7febecaab575 100644 --- a/parquet-variant-compute/src/variant_to_arrow.rs +++ b/parquet-variant-compute/src/variant_to_arrow.rs @@ -20,7 +20,7 @@ use crate::shred_variant::{ make_variant_to_shredded_variant_arrow_row_builder, }; use crate::type_conversion::{ - PrimitiveFromVariant, ShredDecimalVariant, TimestampFromVariant, + PrimitiveFromVariant, ShredDecimalVariant, TimestampFromVariant, VariantCastMode, shred_variant_to_unscaled_decimal, variant_cast_with_options, variant_to_boolean, variant_to_unscaled_decimal, }; @@ -106,7 +106,7 @@ fn make_typed_variant_to_arrow_row_builder<'a>( data_type: &'a DataType, cast_options: &'a CastOptions, capacity: usize, - shred: bool, + cast_mode: VariantCastMode, ) -> Result> { use VariantToArrowRowBuilder::*; @@ -139,7 +139,7 @@ fn make_typed_variant_to_arrow_row_builder<'a>( *ordered, cast_options, capacity, - shred, + cast_mode, )?; Ok(Map(builder)) } @@ -166,7 +166,7 @@ fn make_typed_variant_to_arrow_row_builder<'a>( data_type, cast_options, capacity, - shred, + cast_mode, )?; Ok(Primitive(builder)) } @@ -188,9 +188,12 @@ pub(crate) fn make_variant_to_arrow_row_builder<'a>( metadata.clone(), capacity, )), - Some(data_type) => { - make_typed_variant_to_arrow_row_builder(data_type, cast_options, capacity, false)? - } + Some(data_type) => make_typed_variant_to_arrow_row_builder( + data_type, + cast_options, + capacity, + VariantCastMode::Get, + )?, }; // Wrap with path extraction if needed @@ -403,7 +406,7 @@ impl<'a> EncodedVariantToArrowRowBuilder<'a> { value_type, cast_options, capacity, - false, + VariantCastMode::Get, )?); Ok(Self { data_type, @@ -431,7 +434,7 @@ pub(crate) fn make_primitive_variant_to_arrow_row_builder<'a>( data_type: &'a DataType, cast_options: &'a CastOptions, capacity: usize, - shred: bool, + cast_mode: VariantCastMode, ) -> Result> { use PrimitiveVariantToArrowRowBuilder::*; @@ -440,108 +443,108 @@ pub(crate) fn make_primitive_variant_to_arrow_row_builder<'a>( DataType::Boolean => Boolean(VariantToBooleanArrowRowBuilder::new( cast_options, capacity, - shred, + cast_mode, )), DataType::Int8 => Int8(VariantToPrimitiveArrowRowBuilder::new( cast_options, capacity, - shred, + cast_mode, )), DataType::Int16 => Int16(VariantToPrimitiveArrowRowBuilder::new( cast_options, capacity, - shred, + cast_mode, )), DataType::Int32 => Int32(VariantToPrimitiveArrowRowBuilder::new( cast_options, capacity, - shred, + cast_mode, )), DataType::Int64 => Int64(VariantToPrimitiveArrowRowBuilder::new( cast_options, capacity, - shred, + cast_mode, )), DataType::UInt8 => UInt8(VariantToPrimitiveArrowRowBuilder::new( cast_options, capacity, - shred, + cast_mode, )), DataType::UInt16 => UInt16(VariantToPrimitiveArrowRowBuilder::new( cast_options, capacity, - shred, + cast_mode, )), DataType::UInt32 => UInt32(VariantToPrimitiveArrowRowBuilder::new( cast_options, capacity, - shred, + cast_mode, )), DataType::UInt64 => UInt64(VariantToPrimitiveArrowRowBuilder::new( cast_options, capacity, - shred, + cast_mode, )), DataType::Float16 => Float16(VariantToPrimitiveArrowRowBuilder::new( cast_options, capacity, - shred, + cast_mode, )), DataType::Float32 => Float32(VariantToPrimitiveArrowRowBuilder::new( cast_options, capacity, - shred, + cast_mode, )), DataType::Float64 => Float64(VariantToPrimitiveArrowRowBuilder::new( cast_options, capacity, - shred, + cast_mode, )), DataType::Decimal32(precision, scale) => Decimal32(VariantToDecimalArrowRowBuilder::new( cast_options, capacity, *precision, *scale, - shred, + cast_mode, )?), DataType::Decimal64(precision, scale) => Decimal64(VariantToDecimalArrowRowBuilder::new( cast_options, capacity, *precision, *scale, - shred, + cast_mode, )?), DataType::Decimal128(precision, scale) => Decimal128(VariantToDecimalArrowRowBuilder::new( cast_options, capacity, *precision, *scale, - shred, + cast_mode, )?), DataType::Decimal256(precision, scale) => Decimal256(VariantToDecimalArrowRowBuilder::new( cast_options, capacity, *precision, *scale, - shred, + cast_mode, )?), DataType::Date32 => Date32(VariantToPrimitiveArrowRowBuilder::new( cast_options, capacity, - shred, + cast_mode, )), DataType::Date64 => Date64(VariantToPrimitiveArrowRowBuilder::new( cast_options, capacity, - shred, + cast_mode, )), DataType::Time32(TimeUnit::Second) => Time32Second(VariantToPrimitiveArrowRowBuilder::new( cast_options, capacity, - shred, + cast_mode, )), DataType::Time32(TimeUnit::Millisecond) => Time32Milli( - VariantToPrimitiveArrowRowBuilder::new(cast_options, capacity, shred), + VariantToPrimitiveArrowRowBuilder::new(cast_options, capacity, cast_mode), ), DataType::Time32(t) => { return Err(ArrowError::InvalidArgumentError(format!( @@ -549,10 +552,10 @@ pub(crate) fn make_primitive_variant_to_arrow_row_builder<'a>( ))); } DataType::Time64(TimeUnit::Microsecond) => Time64Micro( - VariantToPrimitiveArrowRowBuilder::new(cast_options, capacity, shred), + VariantToPrimitiveArrowRowBuilder::new(cast_options, capacity, cast_mode), ), DataType::Time64(TimeUnit::Nanosecond) => Time64Nano( - VariantToPrimitiveArrowRowBuilder::new(cast_options, capacity, shred), + VariantToPrimitiveArrowRowBuilder::new(cast_options, capacity, cast_mode), ), DataType::Time64(t) => { return Err(ArrowError::InvalidArgumentError(format!( @@ -560,28 +563,28 @@ pub(crate) fn make_primitive_variant_to_arrow_row_builder<'a>( ))); } DataType::Timestamp(TimeUnit::Second, None) => TimestampSecondNtz( - VariantToTimestampNtzArrowRowBuilder::new(cast_options, capacity, shred), + VariantToTimestampNtzArrowRowBuilder::new(cast_options, capacity, cast_mode), ), DataType::Timestamp(TimeUnit::Second, tz) => TimestampSecond( - VariantToTimestampArrowRowBuilder::new(cast_options, capacity, shred, tz.clone()), + VariantToTimestampArrowRowBuilder::new(cast_options, capacity, cast_mode, tz.clone()), ), DataType::Timestamp(TimeUnit::Millisecond, None) => TimestampMilliNtz( - VariantToTimestampNtzArrowRowBuilder::new(cast_options, capacity, shred), + VariantToTimestampNtzArrowRowBuilder::new(cast_options, capacity, cast_mode), ), DataType::Timestamp(TimeUnit::Millisecond, tz) => TimestampMilli( - VariantToTimestampArrowRowBuilder::new(cast_options, capacity, shred, tz.clone()), + VariantToTimestampArrowRowBuilder::new(cast_options, capacity, cast_mode, tz.clone()), ), DataType::Timestamp(TimeUnit::Microsecond, None) => TimestampMicroNtz( - VariantToTimestampNtzArrowRowBuilder::new(cast_options, capacity, shred), + VariantToTimestampNtzArrowRowBuilder::new(cast_options, capacity, cast_mode), ), DataType::Timestamp(TimeUnit::Microsecond, tz) => TimestampMicro( - VariantToTimestampArrowRowBuilder::new(cast_options, capacity, shred, tz.clone()), + VariantToTimestampArrowRowBuilder::new(cast_options, capacity, cast_mode, tz.clone()), ), DataType::Timestamp(TimeUnit::Nanosecond, None) => TimestampNanoNtz( - VariantToTimestampNtzArrowRowBuilder::new(cast_options, capacity, shred), + VariantToTimestampNtzArrowRowBuilder::new(cast_options, capacity, cast_mode), ), DataType::Timestamp(TimeUnit::Nanosecond, tz) => TimestampNano( - VariantToTimestampArrowRowBuilder::new(cast_options, capacity, shred, tz.clone()), + VariantToTimestampArrowRowBuilder::new(cast_options, capacity, cast_mode, tz.clone()), ), DataType::Duration(_) | DataType::Interval(_) => { return Err(ArrowError::InvalidArgumentError( @@ -655,7 +658,7 @@ impl<'a> StructVariantToArrowRowBuilder<'a> { field.data_type(), cast_options, capacity, - false, + VariantCastMode::Get, )?); } Ok(Self { @@ -760,7 +763,7 @@ impl<'a> UnionVariantToArrowRowBuilder<'a> { field.data_type(), cast_options, capacity, - false, + VariantCastMode::Get, )?, len: 0, }); @@ -978,7 +981,7 @@ impl<'a> MapVariantToArrowRowBuilder<'a> { ordered: bool, cast_options: &'a CastOptions, capacity: usize, - shred: bool, + cast_mode: VariantCastMode, ) -> Result { let DataType::Struct(entry_fields) = entries_field.data_type() else { return Err(ArrowError::InvalidArgumentError(format!( @@ -999,13 +1002,13 @@ impl<'a> MapVariantToArrowRowBuilder<'a> { key_field.data_type(), cast_options, capacity, - shred, + cast_mode, )?); let value_builder = Box::new(make_typed_variant_to_arrow_row_builder( value_field.data_type(), cast_options, capacity, - shred, + cast_mode, )?); if capacity >= isize::MAX as usize { return Err(ArrowError::ComputeError( @@ -1207,12 +1210,12 @@ impl VariantPathRowBuilder<'_> { macro_rules! define_variant_to_primitive_builder { (struct $name:ident<$lifetime:lifetime $(, $generic:ident: $bound:path )?> |$array_param:ident $(, $field:ident: $field_type:ty)?| -> $builder_name:ident $(< $array_type:ty >)? { $init_expr: expr }, - |$value: ident $(, $shred: ident)?| $value_transform:expr, + |$value: ident $(, $cast_mode: ident)?| $value_transform:expr, type_name: $type_name:expr) => { pub(crate) struct $name<$lifetime $(, $generic : $bound )?> { builder: $builder_name $(<$array_type>)?, - $($shred: bool,)? + $($cast_mode: VariantCastMode,)? cast_options: &$lifetime CastOptions<$lifetime>, } @@ -1220,14 +1223,14 @@ macro_rules! define_variant_to_primitive_builder { fn new( cast_options: &$lifetime CastOptions<$lifetime>, $array_param: usize, - $($shred: bool,)? + $($cast_mode: VariantCastMode,)? // add this so that $init_expr can use it $( $field: $field_type, )? ) -> Self { Self { builder: $init_expr, cast_options, - $($shred)? + $($cast_mode)? } } @@ -1237,7 +1240,7 @@ macro_rules! define_variant_to_primitive_builder { } fn append_value(&mut self, $value: &Variant<'_, '_>) -> Result { - $(let $shred: bool = self.shred;)? + $(let $cast_mode: VariantCastMode = self.cast_mode;)? match variant_cast_with_options( $value, self.cast_options, @@ -1283,21 +1286,21 @@ define_variant_to_primitive_builder!( define_variant_to_primitive_builder!( struct VariantToBooleanArrowRowBuilder<'a> |capacity| -> BooleanBuilder { BooleanBuilder::with_capacity(capacity) }, - |value, shred| variant_to_boolean(value, shred), + |value, cast_mode| variant_to_boolean(value, cast_mode), type_name: datatypes::BooleanType::DATA_TYPE ); define_variant_to_primitive_builder!( struct VariantToPrimitiveArrowRowBuilder<'a, T:PrimitiveFromVariant> |capacity| -> PrimitiveBuilder { PrimitiveBuilder::::with_capacity(capacity) }, - |value, shred| T::from_variant(value, shred), + |value, cast_mode| T::from_variant(value, cast_mode), type_name: T::DATA_TYPE ); define_variant_to_primitive_builder!( struct VariantToTimestampNtzArrowRowBuilder<'a, T:TimestampFromVariant> |capacity| -> PrimitiveBuilder { PrimitiveBuilder::::with_capacity(capacity) }, - |value, shred| T::from_variant(value, shred), + |value, cast_mode| T::from_variant(value, cast_mode), type_name: T::DATA_TYPE ); @@ -1306,7 +1309,7 @@ define_variant_to_primitive_builder!( |capacity, tz: Option> | -> PrimitiveBuilder { PrimitiveBuilder::::with_capacity(capacity).with_timezone_opt(tz) }, - |value, shred| T::from_variant(value, shred), + |value, cast_mode| T::from_variant(value, cast_mode), type_name: T::DATA_TYPE ); @@ -1327,7 +1330,7 @@ where cast_options: &'a CastOptions<'a>, precision: u8, scale: i8, - shred: bool, + cast_mode: VariantCastMode, } impl<'a, T> VariantToDecimalArrowRowBuilder<'a, T> @@ -1340,7 +1343,7 @@ where capacity: usize, precision: u8, scale: i8, - shred: bool, + cast_mode: VariantCastMode, ) -> Result { let builder = PrimitiveBuilder::::with_capacity(capacity) .with_precision_and_scale(precision, scale)?; @@ -1349,7 +1352,7 @@ where cast_options, precision, scale, - shred, + cast_mode, }) } @@ -1359,9 +1362,13 @@ where } fn append_value(&mut self, value: &Variant<'_, '_>) -> Result { - match variant_cast_with_options(value, self.cast_options, |value| match self.shred { - true => shred_variant_to_unscaled_decimal::(value, self.precision, self.scale), - false => variant_to_unscaled_decimal::(value, self.precision, self.scale), + match variant_cast_with_options(value, self.cast_options, |value| match self.cast_mode { + VariantCastMode::Shred => { + shred_variant_to_unscaled_decimal::(value, self.precision, self.scale) + } + VariantCastMode::Get => { + variant_to_unscaled_decimal::(value, self.precision, self.scale) + } }) { Ok(Some(scaled)) => { self.builder.append_value(scaled); @@ -1502,7 +1509,7 @@ where cast_options, capacity, NullValue::ArrayElement, - shredded, + VariantCastMode::Shred, )?; ListElementBuilder::Shredded(Box::new(builder)) } else { @@ -1510,7 +1517,7 @@ where element_data_type, cast_options, capacity, - shredded, + VariantCastMode::Get, )?; ListElementBuilder::Typed(Box::new(builder)) }; @@ -1612,7 +1619,7 @@ impl<'a> VariantToFixedSizeListArrowRowBuilder<'a> { cast_options, capacity, NullValue::ArrayElement, - shredded, + VariantCastMode::Shred, )?; ListElementBuilder::Shredded(Box::new(builder)) } else { @@ -1620,7 +1627,7 @@ impl<'a> VariantToFixedSizeListArrowRowBuilder<'a> { element_data_type, cast_options, capacity, - shredded, + VariantCastMode::Get, )?; ListElementBuilder::Typed(Box::new(builder)) }; @@ -1760,6 +1767,7 @@ mod tests { use super::{ make_primitive_variant_to_arrow_row_builder, make_typed_variant_to_arrow_row_builder, }; + use crate::type_conversion::VariantCastMode; use arrow::array::{ Array, Decimal32Array, FixedSizeBinaryArray, Int32Array, ListArray, StructArray, }; @@ -1811,9 +1819,12 @@ mod tests { ]; for data_type in non_primitive_types { - let Err(err) = - make_primitive_variant_to_arrow_row_builder(&data_type, &cast_options, 1, false) - else { + let Err(err) = make_primitive_variant_to_arrow_row_builder( + &data_type, + &cast_options, + 1, + VariantCastMode::Get, + ) else { panic!("non-primitive type {data_type:?} should be rejected") }; @@ -1832,9 +1843,13 @@ mod tests { safe: false, ..Default::default() }; - let mut builder = - make_primitive_variant_to_arrow_row_builder(&DataType::Int32, &cast_options, 2, false) - .unwrap(); + let mut builder = make_primitive_variant_to_arrow_row_builder( + &DataType::Int32, + &cast_options, + 2, + VariantCastMode::Get, + ) + .unwrap(); assert!(!builder.append_value(&Variant::Null).unwrap()); assert!(builder.append_value(&Variant::Int32(42)).unwrap()); @@ -1855,7 +1870,7 @@ mod tests { &DataType::Decimal32(9, 2), &cast_options, 2, - false, + VariantCastMode::Get, ) .unwrap(); let decimal_variant: Variant<'_, '_> = VariantDecimal4::try_new(1234, 2).unwrap().into(); @@ -1879,7 +1894,7 @@ mod tests { &DataType::FixedSizeBinary(16), &cast_options, 2, - false, + VariantCastMode::Get, ) .unwrap(); let uuid = Uuid::nil(); @@ -1904,8 +1919,13 @@ mod tests { }; let list_type = DataType::List(Arc::new(Field::new("item", DataType::Int64, true))); - let mut list_builder = - make_typed_variant_to_arrow_row_builder(&list_type, &cast_options, 1, false).unwrap(); + let mut list_builder = make_typed_variant_to_arrow_row_builder( + &list_type, + &cast_options, + 1, + VariantCastMode::Get, + ) + .unwrap(); assert!(!list_builder.append_value(Variant::Null).unwrap()); let list_array = list_builder.finish().unwrap(); let list_array = list_array.as_any().downcast_ref::().unwrap(); @@ -1913,8 +1933,13 @@ mod tests { let struct_type = DataType::Struct(Fields::from(vec![Field::new("a", DataType::Int32, true)])); - let mut struct_builder = - make_typed_variant_to_arrow_row_builder(&struct_type, &cast_options, 1, false).unwrap(); + let mut struct_builder = make_typed_variant_to_arrow_row_builder( + &struct_type, + &cast_options, + 1, + VariantCastMode::Get, + ) + .unwrap(); assert!(!struct_builder.append_value(Variant::Null).unwrap()); let struct_array = struct_builder.finish().unwrap(); let struct_array = struct_array.as_any().downcast_ref::().unwrap(); diff --git a/parquet-variant-json/Cargo.toml b/parquet-variant-json/Cargo.toml index 706446da0ac9..d2983f579446 100644 --- a/parquet-variant-json/Cargo.toml +++ b/parquet-variant-json/Cargo.toml @@ -20,7 +20,6 @@ name = "parquet-variant-json" version = { workspace = true } license = { workspace = true } description = "Apache Parquet Variant to/from JSON" -homepage = { workspace = true } repository = { workspace = true } authors = { workspace = true } keywords = ["arrow", "parquet", "variant"] diff --git a/parquet-variant/Cargo.toml b/parquet-variant/Cargo.toml index df1aabd03e21..dba33c721339 100644 --- a/parquet-variant/Cargo.toml +++ b/parquet-variant/Cargo.toml @@ -20,7 +20,6 @@ name = "parquet-variant" version = { workspace = true } license = { workspace = true } description = "Apache Parquet Variant implementation in Rust" -homepage = { workspace = true } repository = { workspace = true } authors = { workspace = true } keywords = ["arrow", "parquet", "variant"] diff --git a/parquet-variant/src/decoder.rs b/parquet-variant/src/decoder.rs index b110d056ca83..a74790f7bb64 100644 --- a/parquet-variant/src/decoder.rs +++ b/parquet-variant/src/decoder.rs @@ -445,6 +445,7 @@ mod tests { ); } + #[expect(clippy::float_cmp_const)] mod float { use super::*; diff --git a/parquet-variant/src/path.rs b/parquet-variant/src/path.rs index 64aca085c534..98a6e916d859 100644 --- a/parquet-variant/src/path.rs +++ b/parquet-variant/src/path.rs @@ -322,6 +322,31 @@ mod tests { VariantPathElement::field("1"), ]); assert_eq!(path, expected); + + // an unescaped closing bracket is literal inside a quoted field + for input in [r#"["a]b"]"#, "['a]b']"] { + let path = VariantPath::try_from(input).unwrap(); + let expected = VariantPath::from_iter([VariantPathElement::field("a]b")]); + assert_eq!(path, expected); + } + + // quoted brackets do not consume the following path element + let path = VariantPath::try_from(r#"["[0]"][1]"#).unwrap(); + let expected = VariantPath::from_iter([ + VariantPathElement::field("[0]"), + VariantPathElement::index(1), + ]); + assert_eq!(path, expected); + + // escaped quotes and an unescaped bracket remain one quoted field + let path = VariantPath::try_from(r#"["\"][\""]"#).unwrap(); + let expected = VariantPath::from_iter([VariantPathElement::field(r#""][""#)]); + assert_eq!(path, expected); + + // the existing escaped spelling remains supported + let path = VariantPath::try_from(r#"["a\]b"]"#).unwrap(); + let expected = VariantPath::from_iter([VariantPathElement::field("a]b")]); + assert_eq!(path, expected); } #[test] diff --git a/parquet-variant/src/utils.rs b/parquet-variant/src/utils.rs index 782f718cfc8a..049ab82c9908 100644 --- a/parquet-variant/src/utils.rs +++ b/parquet-variant/src/utils.rs @@ -157,6 +157,7 @@ pub(crate) const fn expect_size_of(expected: usize) { /// - `\\` -> literal `\` /// - `\]` -> literal `]` /// - Any other `\x` -> literal `x` +/// - An unescaped `]` inside a quoted field is literal /// /// Outside brackets, no escaping is supported. /// @@ -169,6 +170,7 @@ pub(crate) const fn expect_size_of(expected: usize) { /// - `"foo[1].bar"` -> field `foo`, index 1, field `bar` /// - `"['a.b']"` -> field `a.b` (dot is literal inside bracket) /// - `"['a\]b']"` -> field `a]b` (escaped `]` +/// - `"['a]b']"` -> field `a]b` (`]` is literal inside quotes) /// - etc. /// /// # Errors @@ -235,6 +237,7 @@ fn parse_in_bracket(s: &str, i: usize) -> Result<(VariantPathElement<'_>, usize) let mut unescaped = String::new(); let mut chars = s[start..].char_indices(); let mut end = None; + let mut quote = None; while let Some((offset, c)) = chars.next() { match c { @@ -245,8 +248,16 @@ fn parse_in_bracket(s: &str, i: usize) -> Result<(VariantPathElement<'_>, usize) } // Trailing backslash will be handled as 'unclosed [' below } - ']' => { - // Unescaped ']' ends the bracket + '\'' | '"' if quote == Some(c) => { + quote = None; + unescaped.push(c); + } + '\'' | '"' if quote.is_none() && unescaped.is_empty() => { + quote = Some(c); + unescaped.push(c); + } + ']' if quote.is_none() => { + // An unescaped ']' outside quotes ends the bracket end = Some(start + offset); break; } diff --git a/parquet/Cargo.toml b/parquet/Cargo.toml index cbf3c74ec2af..52346e7134dd 100644 --- a/parquet/Cargo.toml +++ b/parquet/Cargo.toml @@ -20,7 +20,6 @@ name = "parquet" version = { workspace = true } license = { workspace = true } description = "Apache Parquet implementation in Rust" -homepage = { workspace = true } repository = { workspace = true } authors = { workspace = true } keywords = ["arrow", "parquet", "hadoop"] @@ -343,6 +342,11 @@ name = "arrow_reader_peak_memory" required-features = ["arrow", "test_common", "experimental"] harness = false +[[bench]] +name = "arrow_writer_peak_memory" +required-features = ["arrow"] +harness = false + [lib] bench = false diff --git a/parquet/benches/arrow_statistics.rs b/parquet/benches/arrow_statistics.rs index a8a1f2b1e5ab..af201656d77a 100644 --- a/parquet/benches/arrow_statistics.rs +++ b/parquet/benches/arrow_statistics.rs @@ -18,7 +18,7 @@ //! Benchmarks of benchmark for extracting arrow statistics from parquet use arrow::array::{ArrayRef, DictionaryArray, Float64Array, StringArray, UInt64Array}; -use arrow_array::{Int32Array, Int64Array, RecordBatch}; +use arrow_array::{Decimal128Array, Int32Array, Int64Array, RecordBatch, StringViewArray}; use arrow_schema::{ DataType::{self, *}, Field, Schema, @@ -26,7 +26,11 @@ use arrow_schema::{ use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main}; use parquet::{ arrow::arrow_reader::ArrowReaderOptions, - file::{metadata::PageIndexPolicy, properties::WriterProperties}, + file::{ + metadata::{PageIndexPolicy, page_index::PageIndexBuilder}, + page_index::index_reader::decode_column_index, + properties::WriterProperties, + }, }; use parquet::{ arrow::{ArrowWriter, arrow_reader::ArrowReaderBuilder}, @@ -253,5 +257,128 @@ fn criterion_benchmark(c: &mut Criterion) { } } -criterion_group!(benches, criterion_benchmark); +/// Makes one column with `rows` values, where every 7th value is null. +fn make_page_index_column(data_type: &DataType, rows: usize) -> ArrayRef { + let valid = |i: usize| !i.is_multiple_of(7); + match data_type { + Int64 => Arc::new(Int64Array::from_iter( + (0..rows).map(|i| valid(i).then_some(i as i64 * 3)), + )), + Utf8 => Arc::new(StringArray::from_iter( + (0..rows).map(|i| valid(i).then(|| format!("value-{i:08}"))), + )), + Utf8View => Arc::new(StringViewArray::from_iter( + (0..rows).map(|i| valid(i).then(|| format!("value-{i:08}"))), + )), + Decimal128(precision, scale) => Arc::new( + Decimal128Array::from_iter((0..rows).map(|i| valid(i).then_some(i as i128 * 1001))) + .with_precision_and_scale(*precision, *scale) + .unwrap(), + ), + _ => unimplemented!("{data_type}"), + } +} + +/// Writes a file with many small data pages and returns its bytes. +fn create_page_index_file( + data_type: &DataType, + row_groups: usize, + rows_per_group: usize, +) -> Vec { + let schema = Arc::new(Schema::new(vec![Field::new( + "col", + data_type.clone(), + true, + )])); + let props = WriterProperties::builder() + .set_max_row_group_row_count(Some(rows_per_group)) + .set_data_page_row_count_limit(10) + .set_write_batch_size(10) + .set_statistics_enabled(EnabledStatistics::Page) + .build(); + let mut buffer = Vec::new(); + let mut writer = ArrowWriter::try_new(&mut buffer, schema.clone(), Some(props)).unwrap(); + let column = make_page_index_column(data_type, row_groups * rows_per_group); + let batch = RecordBatch::try_new(schema, vec![column]).unwrap(); + // The page row limit is only checked between writes, so write in small slices + for offset in (0..batch.num_rows()).step_by(10) { + writer.write(&batch.slice(offset, 10)).unwrap(); + } + writer.close().unwrap(); + buffer +} + +/// Measures getting page statistics from the stored column index bytes by +/// building `ColumnIndexMetaData` and converting it to Arrow arrays. +fn page_index_benchmark(c: &mut Criterion) { + let row_groups = 20; + let data_types = [Int64, Utf8, Utf8View, Decimal128(20, 2)]; + // 10 rows per page, so 100 or 500 pages per row group: 2000 or 10000 pages + let rows_per_group_options = [1000, 5000]; + + for (data_type, rows_per_group) in data_types + .iter() + .flat_map(|t| rows_per_group_options.map(|rows| (t.clone(), rows))) + { + let data = bytes::Bytes::from(create_page_index_file( + &data_type, + row_groups, + rows_per_group, + )); + let options = ArrowReaderOptions::new().with_page_index_policy(PageIndexPolicy::from(true)); + let reader = ArrowReaderBuilder::try_new_with_options(data.clone(), options).unwrap(); + let metadata = reader.metadata().clone(); + let converter = + StatisticsConverter::try_new("col", reader.schema(), reader.parquet_schema()).unwrap(); + let column = converter.parquet_column_index().unwrap(); + let physical_type = reader.parquet_schema().column(column).physical_type(); + + // (number of pages, stored column index bytes) for each row group + let column_indexes: Vec<(usize, &[u8])> = metadata + .row_groups() + .iter() + .enumerate() + .map(|(rg, row_group)| { + let range = row_group.column(column).column_index_range().unwrap(); + let num_pages = metadata + .page_index_for_row_group(rg) + .num_data_pages(column) + .unwrap(); + (num_pages, &data[range.start as usize..range.end as usize]) + }) + .collect(); + let row_group_indices: Vec = (0..row_groups).collect(); + let num_columns = reader.parquet_schema().num_columns(); + + let mut group = c.benchmark_group(format!( + "Decode page index statistics for {data_type} ({} pages)", + column_indexes.iter().map(|(n, _)| n).sum::() + )); + group.bench_function("full page index", |b| { + b.iter(|| { + let mut builder = PageIndexBuilder::new(row_groups, num_columns); + for (rg, (_, bytes)) in column_indexes.iter().enumerate() { + let index = decode_column_index(bytes, physical_type).unwrap(); + builder.put_column_index(index, rg, column); + } + let page_index = builder.build(); + let _ = converter + .data_page_mins(&page_index, &row_group_indices) + .unwrap(); + let _ = converter + .data_page_maxes(&page_index, &row_group_indices) + .unwrap(); + let _ = converter + .data_page_null_counts(&page_index, &row_group_indices) + .unwrap(); + let _ = converter + .data_page_nan_counts(&page_index, &row_group_indices) + .unwrap(); + }) + }); + group.finish(); + } +} + +criterion_group!(benches, criterion_benchmark, page_index_benchmark); criterion_main!(benches); diff --git a/parquet/benches/arrow_writer_peak_memory.rs b/parquet/benches/arrow_writer_peak_memory.rs new file mode 100644 index 000000000000..fc8d7ea18b36 --- /dev/null +++ b/parquet/benches/arrow_writer_peak_memory.rs @@ -0,0 +1,1071 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Criterion benchmark measuring the *memory footprint* of `ArrowWriter` when +//! encoding an Arrow `RecordBatch` to Parquet. +//! +//! Uses a thread-local tracking allocator and two custom criterion measurements: +//! [`PeakMemory`] for the live high-water mark and [`AllocatedBytes`] for +//! cumulative bytes requested. The benchmark quantifies writer memory over and +//! above the input batch and configuration throughout the writer's lifecycle. +//! +//! How the measurement isolates *writer* memory: +//! * The `RecordBatch` is built lazily before `iter_batched`, and writer +//! properties are cloned in its **setup** closure. Both run *outside* the +//! measured window. Their live allocations therefore land in the +//! per-iteration baseline (subtracted out of `PeakMemory`) and are never +//! counted by `AllocatedBytes`. +//! * `ArrowWriter::try_new`, `writer.write(&batch)`, and `writer.close()` run +//! inside the measured window, including writer initialization, encoding, +//! finalization, and destruction. +//! * The output sink is [`std::io::Empty`], which discards serialized bytes, so +//! the figure reflects the writer's internal working set (column encoders, +//! dictionary interners, level/encoding buffers, retained page data up to the +//! row-group flush) rather than the size of the produced file. +//! * `BatchSize::PerIteration` is mandatory: it forces criterion to bracket a +//! single `routine()` call between `Measurement::start()`/`end()`, so each +//! sample re-anchors the peak and isolates one complete writer lifecycle. +//! +//! Constraints: +//! * The tracking allocator is **thread-local**, so this only measures the +//! synchronous, single-threaded `ArrowWriter`. A parallel/async writer would +//! allocate on other threads and be silently undercounted. + +use std::alloc::Layout; +use std::cell::{Cell, OnceCell}; +use std::io::Empty; +use std::sync::Arc; + +use arrow::datatypes::*; +use arrow::record_batch::RecordBatch; +use arrow::util::data_gen::create_random_batch; +use arrow_array::{ + Array, ArrayRef, BooleanArray, Decimal128Array, DictionaryArray, FixedSizeBinaryArray, + Int8Array, Int32Array, Int64Array, ListArray, RunArray, StringArray, StructArray, +}; +use arrow_buffer::{NullBuffer, OffsetBuffer}; +use criterion::measurement::{Measurement, ValueFormatter}; +use criterion::{BatchSize, BenchmarkGroup, Criterion, criterion_group, criterion_main}; +use parquet::arrow::ArrowWriter; +use parquet::errors::Result; +use parquet::file::properties::{CdcOptions, WriterProperties, WriterVersion}; + +// --------------------------------------------------------------------------- +// Thread-local tracking allocator +// --------------------------------------------------------------------------- + +thread_local! { + static LIVE_BYTES: Cell = const { Cell::new(0) }; + static PEAK_BYTES: Cell = const { Cell::new(0) }; + static ALLOCATED_BYTES: Cell = const { Cell::new(0) }; +} + +struct TrackingAllocator { + inner: std::alloc::System, +} + +#[global_allocator] +static GLOBAL: TrackingAllocator = TrackingAllocator { + inner: std::alloc::System, +}; + +fn add_live_bytes(size: usize) { + LIVE_BYTES.with(|live| { + let new = live.get().saturating_add(size); + live.set(new); + PEAK_BYTES.with(|peak| { + if new > peak.get() { + peak.set(new); + } + }); + }); +} + +fn subtract_live_bytes(size: usize) { + LIVE_BYTES.with(|live| { + live.set(live.get().saturating_sub(size)); + }); +} + +fn add_allocated_bytes(size: usize) { + ALLOCATED_BYTES.with(|allocated| { + allocated.set(allocated.get().saturating_add(size)); + }); +} + +#[expect(unsafe_code)] +unsafe impl std::alloc::GlobalAlloc for TrackingAllocator { + unsafe fn alloc(&self, layout: Layout) -> *mut u8 { + let ptr = unsafe { self.inner.alloc(layout) }; + if !ptr.is_null() { + add_live_bytes(layout.size()); + add_allocated_bytes(layout.size()); + } + ptr + } + + unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) { + subtract_live_bytes(layout.size()); + unsafe { self.inner.dealloc(ptr, layout) }; + } + + unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 { + let new_ptr = unsafe { self.inner.realloc(ptr, layout, new_size) }; + if !new_ptr.is_null() { + let old_size = layout.size(); + add_allocated_bytes(new_size); + if new_size > old_size { + add_live_bytes(new_size - old_size); + } else { + subtract_live_bytes(old_size - new_size); + } + } + new_ptr + } +} + +fn reset_peak() { + PEAK_BYTES.with(|peak| { + LIVE_BYTES.with(|live| { + peak.set(live.get()); + }); + }); +} + +fn peak_bytes() -> usize { + PEAK_BYTES.with(|peak| peak.get()) +} + +fn live_bytes() -> usize { + LIVE_BYTES.with(|live| live.get()) +} + +fn reset_allocated() { + ALLOCATED_BYTES.with(|allocated| allocated.set(0)); +} + +fn allocated_bytes() -> usize { + ALLOCATED_BYTES.with(|allocated| allocated.get()) +} + +// --------------------------------------------------------------------------- +// Criterion custom byte measurements +// --------------------------------------------------------------------------- + +struct BytesFormatter; + +const BYTE_UNITS: &[(u32, &str)] = &[ + (60, "EiB"), + (50, "PiB"), + (40, "TiB"), + (30, "GiB"), + (20, "MiB"), + (10, "KiB"), + (0, "B"), +]; + +fn bytes_per_unit(exponent: u32) -> f64 { + (1_u64 << exponent) as f64 +} + +fn scale_bytes(typical: f64, values: &mut [f64]) -> &'static str { + for &(exponent, unit) in BYTE_UNITS { + let scale = bytes_per_unit(exponent); + if typical >= scale { + for v in values.iter_mut() { + *v /= scale; + } + return unit; + } + } + unreachable!("BYTE_UNITS contains B") +} + +impl ValueFormatter for BytesFormatter { + fn scale_values(&self, typical: f64, values: &mut [f64]) -> &'static str { + scale_bytes(typical, values) + } + + fn scale_throughputs( + &self, + typical: f64, + _throughput: &criterion::Throughput, + values: &mut [f64], + ) -> &'static str { + scale_bytes(typical, values) + } + + fn scale_for_machines(&self, values: &mut [f64]) -> &'static str { + // Machine-readable: always bytes + let _ = values; + "B" + } +} + +struct PeakMemory; + +impl Measurement for PeakMemory { + type Intermediate = usize; + type Value = usize; + + fn start(&self) -> Self::Intermediate { + reset_peak(); + live_bytes() + } + + fn end(&self, baseline: Self::Intermediate) -> Self::Value { + peak_bytes().saturating_sub(baseline) + } + + fn add(&self, v1: &Self::Value, v2: &Self::Value) -> Self::Value { + *v1 + *v2 + } + + fn zero(&self) -> Self::Value { + 0 + } + + fn to_f64(&self, value: &Self::Value) -> f64 { + *value as f64 + } + + fn formatter(&self) -> &dyn ValueFormatter { + &BytesFormatter + } +} + +struct AllocatedBytes; + +impl Measurement for AllocatedBytes { + type Intermediate = (); + type Value = usize; + + fn start(&self) -> Self::Intermediate { + reset_allocated(); + } + + fn end(&self, _baseline: Self::Intermediate) -> Self::Value { + allocated_bytes() + } + + fn add(&self, v1: &Self::Value, v2: &Self::Value) -> Self::Value { + *v1 + *v2 + } + + fn zero(&self) -> Self::Value { + 0 + } + + fn to_f64(&self, value: &Self::Value) -> f64 { + *value as f64 + } + + fn formatter(&self) -> &dyn ValueFormatter { + &BytesFormatter + } +} + +// --------------------------------------------------------------------------- +// Test data generation for representative writer input shapes. +// --------------------------------------------------------------------------- + +const BATCH_SIZE: usize = 1024 * 1024; + +fn create_primitive_bench_batch( + size: usize, + null_density: f32, + true_density: f32, +) -> Result { + let fields = vec![ + Field::new("_1", DataType::Int32, true), + Field::new("_2", DataType::Int64, true), + Field::new("_3", DataType::UInt32, true), + Field::new("_4", DataType::UInt64, true), + Field::new("_5", DataType::Float32, true), + Field::new("_6", DataType::Float64, true), + Field::new("_7", DataType::Date64, true), + ]; + let schema = Schema::new(fields); + Ok(create_random_batch( + Arc::new(schema), + size, + null_density, + true_density, + )?) +} + +fn create_primitive_bench_batch_non_null( + size: usize, + null_density: f32, + true_density: f32, +) -> Result { + let fields = vec![ + Field::new("_1", DataType::Int32, false), + Field::new("_2", DataType::Int64, false), + Field::new("_3", DataType::UInt32, false), + Field::new("_4", DataType::UInt64, false), + Field::new("_5", DataType::Float32, false), + Field::new("_6", DataType::Float64, false), + Field::new("_7", DataType::Date64, false), + ]; + let schema = Schema::new(fields); + Ok(create_random_batch( + Arc::new(schema), + size, + null_density, + true_density, + )?) +} + +fn create_bool_bench_batch_non_null( + size: usize, + null_density: f32, + true_density: f32, +) -> Result { + let fields = vec![Field::new("_1", DataType::Boolean, false)]; + let schema = Schema::new(fields); + Ok(create_random_batch( + Arc::new(schema), + size, + null_density, + true_density, + )?) +} + +fn create_bool_dictionary_bench_batch(size: usize) -> Result { + let keys = Int8Array::from_iter((0..size).map(|i| (i % 10 != 9).then_some((i % 2) as i8))); + let values = BooleanArray::from(vec![false, true]); + let array: ArrayRef = Arc::new(DictionaryArray::new(keys, Arc::new(values) as ArrayRef)); + Ok(RecordBatch::try_from_iter([("col", array)])?) +} + +fn create_string_bench_batch_non_null( + size: usize, + null_density: f32, + true_density: f32, +) -> Result { + let fields = vec![ + Field::new("_1", DataType::Utf8, false), + Field::new("_2", DataType::LargeUtf8, false), + ]; + let schema = Schema::new(fields); + Ok(create_random_batch( + Arc::new(schema), + size, + null_density, + true_density, + )?) +} + +/// High-null byte-array columns (Utf8 / LargeUtf8 / Utf8View / BinaryView). +fn create_byte_array_sparse_bench_batch(size: usize, null_density: f32) -> Result { + let fields = vec![ + Field::new("_1", DataType::Utf8, true), + Field::new("_2", DataType::LargeUtf8, true), + Field::new("_3", DataType::Utf8View, true), + Field::new("_4", DataType::BinaryView, true), + ]; + let schema = Schema::new(fields); + Ok(create_random_batch( + Arc::new(schema), + size, + null_density, + 0.5, + )?) +} + +/// Many short, fixed-width strings — exercises the BYTE_ARRAY dense hot path. +fn create_short_string_bench_batch(size: usize) -> Result { + let array = Arc::new(StringArray::from_iter_values( + (0..size).map(|i| format!("{i:08}")), + )) as _; + Ok(RecordBatch::try_from_iter([("col", array)])?) +} + +/// FixedSizeBinary column → FIXED_LEN_BYTE_ARRAY (raw FLBA path). +fn create_fixed_size_binary_bench_batch( + size: usize, + byte_width: i32, + null_density: f32, +) -> Result { + let width = byte_width as usize; + let nulls = nulls_for_density(size, null_density); + let array = FixedSizeBinaryArray::try_from_sparse_iter_with_size( + (0..size).map(|i| { + nulls.as_ref().is_none_or(|n| n.is_valid(i)).then(|| { + let mut v = vec![0u8; width]; + let b = (i as u64).to_le_bytes(); + let n = width.min(8); + v[..n].copy_from_slice(&b[..n]); + v + }) + }), + byte_width, + )?; + Ok(RecordBatch::try_from_iter([( + "col", + Arc::new(array) as ArrayRef, + )])?) +} + +/// Decimal128(38, 10) column → FIXED_LEN_BYTE_ARRAY(16) (converted FLBA path). +fn create_decimal128_bench_batch(size: usize, null_density: f32) -> Result { + let values: Vec = (0..size).map(|i| i as i128).collect(); + let nulls = nulls_for_density(size, null_density); + let array = Decimal128Array::new(values.into(), nulls).with_precision_and_scale(38, 10)?; + Ok(RecordBatch::try_from_iter([( + "col", + Arc::new(array) as ArrayRef, + )])?) +} + +fn create_decimal128_dictionary_bench_batch( + size: usize, + null_density: f32, + true_density: f32, +) -> Result { + Ok(create_random_batch( + dictionary_schema(DataType::Decimal128(38, 10)), + size, + null_density, + true_density, + )?) +} + +fn create_decimal128_dictionary_bench_batch_1pct_cardinality( + size: usize, + null_density: f32, +) -> Result { + let cardinality = dictionary_cardinality_1pct(size); + let keys = dictionary_keys(size, cardinality, null_density); + let values = Decimal128Array::from_iter_values((0..cardinality).map(|i| i as i128)) + .with_precision_and_scale(38, 10)?; + let array: ArrayRef = Arc::new(DictionaryArray::::new(keys, Arc::new(values))); + Ok(RecordBatch::try_from_iter([("col", array)])?) +} + +fn create_int64_dictionary_bench_batch( + size: usize, + null_density: f32, + true_density: f32, +) -> Result { + Ok(create_random_batch( + dictionary_schema(DataType::Int64), + size, + null_density, + true_density, + )?) +} + +fn create_int64_dictionary_bench_batch_1pct_cardinality( + size: usize, + null_density: f32, +) -> Result { + let cardinality = dictionary_cardinality_1pct(size); + let schema = dictionary_schema(DataType::Int64); + let keys = dictionary_keys(size, cardinality, null_density); + let values = Int64Array::from_iter_values((0..cardinality).map(|i| i as i64)); + let array = DictionaryArray::::new(keys, Arc::new(values)); + Ok(RecordBatch::try_new( + schema, + vec![Arc::new(array) as ArrayRef], + )?) +} + +fn create_string_dictionary_bench_batch( + size: usize, + null_density: f32, + true_density: f32, +) -> Result { + let fields = vec![Field::new( + "_1", + DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)), + true, + )]; + let schema = Schema::new(fields); + Ok(create_random_batch( + Arc::new(schema), + size, + null_density, + true_density, + )?) +} + +fn create_string_dictionary_bench_batch_1pct_cardinality( + size: usize, + null_density: f32, +) -> Result { + let cardinality = dictionary_cardinality_1pct(size); + let schema = dictionary_schema(DataType::Utf8); + let keys = dictionary_keys(size, cardinality, null_density); + let values = StringArray::from_iter_values((0..cardinality).map(|i| format!("value_{i:08}"))); + let array = DictionaryArray::::new(keys, Arc::new(values)); + Ok(RecordBatch::try_new( + schema, + vec![Arc::new(array) as ArrayRef], + )?) +} + +/// A single run-end-encoded column with long runs of few distinct values — +/// REE's intended regime. `make_values` produces the per-run value array. +fn create_run_end_encoded_bench_batch( + run_ends_type: DataType, + size: usize, + run_len: usize, + make_values: impl Fn(usize) -> ArrayRef, +) -> RecordBatch { + let num_runs = size.div_ceil(run_len); + let mut acc = 0usize; + let run_ends: Vec = (0..num_runs) + .map(|_| { + acc = (acc + run_len).min(size); + acc as i64 + }) + .collect(); + let values = make_values(num_runs); + let run_array: ArrayRef = match run_ends_type { + DataType::Int32 => { + let ends = Int32Array::from_iter_values(run_ends.iter().map(|&v| v as i32)); + Arc::new(RunArray::::try_new(&ends, &values).unwrap()) + } + DataType::Int64 => { + let ends = Int64Array::from_iter_values(run_ends.iter().copied()); + Arc::new(RunArray::::try_new(&ends, &values).unwrap()) + } + other => panic!("unsupported REE run-ends type for bench: {other:?}"), + }; + let schema = Arc::new(Schema::new(vec![Field::new( + "_1", + run_array.data_type().clone(), + true, + )])); + RecordBatch::try_new(schema, vec![run_array]).unwrap() +} + +/// Struct-of-leaves run values for the non-leaf REE cases: 16 distinct +/// (int64, utf8) pairs cycled over `n` runs, one in eight struct-null. +/// Mirrors `arrow_writer.rs`. +fn make_ree_struct_values(n: usize) -> ArrayRef { + let a = Int64Array::from_iter_values((0..n).map(|i| (i % 16) as i64)); + let b = StringArray::from_iter_values((0..n).map(|i| format!("category_{:02}", i % 16))); + let fields = Fields::from(vec![ + Field::new("a", DataType::Int64, true), + Field::new("b", DataType::Utf8, true), + ]); + let nulls = NullBuffer::from_iter((0..n).map(|i| i % 8 != 7)); + Arc::new(StructArray::new( + fields, + vec![Arc::new(a), Arc::new(b)], + Some(nulls), + )) +} + +/// List run values for the non-leaf REE cases: each run value is a list of +/// 8 ints drawn from 16 distinct lists. Mirrors `arrow_writer.rs`. +fn make_ree_list_values(n: usize) -> ArrayRef { + let values = + Int32Array::from_iter_values((0..n * 8).map(|i| ((i / 8) % 16 * 8 + i % 8) as i32)); + let offsets = OffsetBuffer::from_lengths(std::iter::repeat_n(8usize, n)); + let field = Arc::new(Field::new_list_field(DataType::Int32, true)); + Arc::new(ListArray::new(field, offsets, Arc::new(values), None)) +} + +/// A batch of `size` lists of 4 elements over a run-end-encoded child with +/// 256-row runs (16 distinct values, one run in eight null). +/// Mirrors `arrow_writer.rs`. +fn create_list_of_int32_ree_batch(size: usize) -> RecordBatch { + let child_len = size * 4; + let num_runs = child_len.div_ceil(256); + let run_ends = Int32Array::from_iter_values( + (0..num_runs).map(|i| (((i + 1) * 256).min(child_len)) as i32), + ); + let values = + Int32Array::from_iter((0..num_runs).map(|i| (i % 8 != 7).then_some((i % 16) as i32))); + let ree: ArrayRef = Arc::new(RunArray::::try_new(&run_ends, &values).unwrap()); + let offsets = OffsetBuffer::from_lengths(std::iter::repeat_n(4usize, size)); + let field = Arc::new(Field::new_list_field(ree.data_type().clone(), true)); + let list: ArrayRef = Arc::new(ListArray::new(field, offsets, ree, None)); + let schema = Arc::new(Schema::new(vec![Field::new( + "_1", + list.data_type().clone(), + true, + )])); + RecordBatch::try_new(schema, vec![list]).unwrap() +} + +/// A list-of-runs column with null struct-parent rows at the beginning, +/// middle, and end, separating the visible child rows into two ranges. +fn create_fragmented_list_of_int32_ree_batch(size: usize) -> RecordBatch { + let list_batch = create_list_of_int32_ree_batch(size); + let list = list_batch.column(0).clone(); + let fields = Fields::from(vec![Field::new("items", list.data_type().clone(), true)]); + let nulls = NullBuffer::from_iter((0..size).map(|i| i != 0 && i != size / 2 && i + 1 != size)); + let structs: ArrayRef = Arc::new(StructArray::new(fields, vec![list], Some(nulls))); + let schema = Arc::new(Schema::new(vec![Field::new( + "_1", + structs.data_type().clone(), + true, + )])); + RecordBatch::try_new(schema, vec![structs]).unwrap() +} + +/// A batch of `size` outer lists of 2 elements over a run-end-encoded child +/// (128-element runs) whose run values are inner lists of 4 ints. +/// Mirrors `arrow_writer.rs`. +fn create_list_of_list_int32_ree_batch(size: usize) -> RecordBatch { + let child_len = size * 2; + let num_runs = child_len.div_ceil(128); + let run_ends = Int32Array::from_iter_values( + (0..num_runs).map(|i| (((i + 1) * 128).min(child_len)) as i32), + ); + let inner_values = + Int32Array::from_iter_values((0..num_runs * 4).map(|i| ((i / 4) % 16 * 4 + i % 4) as i32)); + let inner_offsets = OffsetBuffer::from_lengths(std::iter::repeat_n(4usize, num_runs)); + let inner_field = Arc::new(Field::new_list_field(DataType::Int32, true)); + let inner = ListArray::new(inner_field, inner_offsets, Arc::new(inner_values), None); + let ree: ArrayRef = Arc::new(RunArray::::try_new(&run_ends, &inner).unwrap()); + let offsets = OffsetBuffer::from_lengths(std::iter::repeat_n(2usize, size)); + let field = Arc::new(Field::new_list_field(ree.data_type().clone(), true)); + let list: ArrayRef = Arc::new(ListArray::new(field, offsets, ree, None)); + let schema = Arc::new(Schema::new(vec![Field::new( + "_1", + list.data_type().clone(), + true, + )])); + RecordBatch::try_new(schema, vec![list]).unwrap() +} + +/// `RunEndEncoded>>>`. Mirrors +/// `arrow_writer.rs`. +fn create_ree_list_ree_list_batch(size: usize) -> RecordBatch { + let outer_runs = size.div_ceil(256); + let inner_len = outer_runs * 2; + let inner_runs = inner_len.div_ceil(128); + + let values = Int32Array::from_iter_values( + (0..inner_runs * 4).map(|i| ((i / 4) % 16 * 4 + i % 4) as i32), + ); + let value_offsets = OffsetBuffer::from_lengths(std::iter::repeat_n(4usize, inner_runs)); + let value_field = Arc::new(Field::new_list_field(DataType::Int32, true)); + let value_lists = ListArray::new(value_field, value_offsets, Arc::new(values), None); + + let inner_ends = Int32Array::from_iter_values( + (0..inner_runs).map(|i| (((i + 1) * 128).min(inner_len)) as i32), + ); + let inner: ArrayRef = + Arc::new(RunArray::::try_new(&inner_ends, &value_lists).unwrap()); + let outer_value_offsets = OffsetBuffer::from_lengths(std::iter::repeat_n(2usize, outer_runs)); + let outer_value_field = Arc::new(Field::new_list_field(inner.data_type().clone(), true)); + let outer_values: ArrayRef = Arc::new(ListArray::new( + outer_value_field, + outer_value_offsets, + inner, + None, + )); + + let outer_ends = + Int32Array::from_iter_values((0..outer_runs).map(|i| (((i + 1) * 256).min(size)) as i32)); + let array: ArrayRef = + Arc::new(RunArray::::try_new(&outer_ends, &outer_values).unwrap()); + let schema = Arc::new(Schema::new(vec![Field::new( + "_1", + array.data_type().clone(), + true, + )])); + RecordBatch::try_new(schema, vec![array]).unwrap() +} + +fn dictionary_schema(value_type: DataType) -> Arc { + Arc::new(Schema::new(vec![Field::new( + "_1", + DataType::Dictionary(Box::new(DataType::Int32), Box::new(value_type)), + true, + )])) +} + +fn dictionary_cardinality_1pct(size: usize) -> usize { + (size / 100).max(1) +} + +fn dictionary_keys(size: usize, cardinality: usize, null_density: f32) -> Int32Array { + let keys = (0..size) + .map(|i| (i % cardinality) as i32) + .collect::>(); + Int32Array::new(keys.into(), nulls_for_density(size, null_density)) +} + +fn nulls_for_density(size: usize, null_density: f32) -> Option { + if null_density == 0. { + return None; + } + let null_threshold = (null_density.clamp(0., 1.) * 10_000.) as usize; + Some(NullBuffer::from( + (0..size) + .map(|i| (i * 9973) % 10_000 >= null_threshold) + .collect::>(), + )) +} + +// --------------------------------------------------------------------------- +// Benchmark driver +// --------------------------------------------------------------------------- + +/// Writer property variants to measure. Kept small: `default` for parquet 1.0 +/// encoders, `parquet_2` for v2 / DELTA encoder coverage, and `cdc` for the +/// content-defined-chunking path. +fn writer_props() -> Vec<(&'static str, WriterProperties)> { + vec![ + ("default", WriterProperties::builder().build()), + ( + "parquet_2", + WriterProperties::builder() + .set_writer_version(WriterVersion::PARQUET_2_0) + .build(), + ), + ( + "cdc", + WriterProperties::builder() + .set_content_defined_chunking(Some(CdcOptions::default())) + .build(), + ), + ] +} + +/// Measure one `(batch, props)` cell. Input construction and the props clone are +/// outside the measured window. Writer construction, `write`, and `close` into a +/// discarding [`Empty`] sink, including writer destruction, are measured. +fn bench_cell( + group: &mut BenchmarkGroup, + prop_name: &str, + benchmark: &BatchBenchmark, + batch: &OnceCell, + props: &WriterProperties, +) { + group.bench_function(prop_name, |b| { + // Filtered-out Criterion IDs never invoke this closure, so exact + // filters build only the selected input instead of every 1M-row batch. + let batch = batch.get_or_init(|| (benchmark.make_batch)()); + b.iter_batched( + || props.clone(), + |props| { + let mut writer = + ArrowWriter::try_new(Empty::default(), batch.schema(), Some(props)) + .expect("writer construction"); + writer.write(batch).expect("write"); + writer.close().expect("close"); + }, + BatchSize::PerIteration, + ); + }); +} + +/// Run-end-encoded column whose run VALUES are a low-cardinality +/// `Dictionary` (run length 8 -> ~size/8 runs, 16 distinct dict +/// entries). Exercises the REE-of-dictionary path with many logical rows but +/// only `size / 8` physical runs. +fn create_ree_of_dict_bench_batch(size: usize) -> RecordBatch { + let run_len = 8usize; + let num_runs = size / run_len; + let run_ends = Int32Array::from( + (1..=num_runs) + .map(|i| (i * run_len) as i32) + .collect::>(), + ); + let dict_values = StringArray::from((0..16).map(|i| format!("val_{i:04}")).collect::>()); + let keys = Int32Array::from((0..num_runs).map(|i| (i % 16) as i32).collect::>()); + let dict = DictionaryArray::::new(keys, Arc::new(dict_values)); + let ree = RunArray::::try_new(&run_ends, &dict).unwrap(); + let field = Field::new("_1", ree.data_type().clone(), true); + let schema = Arc::new(Schema::new(vec![field])); + RecordBatch::try_new(schema, vec![Arc::new(ree)]).unwrap() +} + +/// An Int32 run-end array with `Dictionary` values, +/// exercising dictionary mapping at run granularity. +fn create_ree_of_numeric_dict_bench_batch(size: usize) -> RecordBatch { + let run_len = 8usize; + let num_runs = size / run_len; + let run_ends = Int32Array::from( + (1..=num_runs) + .map(|i| (i * run_len) as i32) + .collect::>(), + ); + let dict_values = Int64Array::from((0..16).map(|i| i as i64 * 1_000_000).collect::>()); + let keys = Int32Array::from((0..num_runs).map(|i| (i % 16) as i32).collect::>()); + let dict = DictionaryArray::::new(keys, Arc::new(dict_values)); + let ree = RunArray::::try_new(&run_ends, &dict).unwrap(); + let field = Field::new("_1", ree.data_type().clone(), true); + let schema = Arc::new(Schema::new(vec![field])); + RecordBatch::try_new(schema, vec![Arc::new(ree)]).unwrap() +} + +/// An Int32 run-end array of structs with `Dictionary` and Int64 +/// fields, exercising dictionary mapping beneath repeated struct records. +fn create_ree_struct_of_dict_bench_batch(size: usize) -> RecordBatch { + let run_len = 8usize; + let num_runs = size / run_len; + let run_ends = Int32Array::from( + (1..=num_runs) + .map(|i| (i * run_len) as i32) + .collect::>(), + ); + let dict_values = StringArray::from((0..16).map(|i| format!("tag_{i:04}")).collect::>()); + let keys = Int32Array::from((0..num_runs).map(|i| (i % 16) as i32).collect::>()); + let tag = DictionaryArray::::new(keys, Arc::new(dict_values)); + let val = Int64Array::from((0..num_runs).map(|i| i as i64).collect::>()); + let fields = Fields::from(vec![ + Field::new("tag", tag.data_type().clone(), true), + Field::new("val", DataType::Int64, true), + ]); + let values = StructArray::new(fields, vec![Arc::new(tag), Arc::new(val)], None); + let ree = RunArray::::try_new(&run_ends, &values).unwrap(); + let field = Field::new("_1", ree.data_type().clone(), true); + let schema = Arc::new(Schema::new(vec![field])); + RecordBatch::try_new(schema, vec![Arc::new(ree)]).unwrap() +} + +struct BatchBenchmark { + name: &'static str, + make_batch: Box RecordBatch>, +} + +impl BatchBenchmark { + fn new(name: &'static str, make_batch: impl Fn() -> RecordBatch + 'static) -> Self { + Self { + name, + make_batch: Box::new(make_batch), + } + } +} + +macro_rules! lazy_batches { + ($(($name:expr, $batch:expr $(,)?)),* $(,)?) => { + vec![$(BatchBenchmark::new($name, || $batch)),*] + }; +} + +fn create_batches() -> Vec { + lazy_batches![ + ( + "primitive_non_null", + create_primitive_bench_batch_non_null(BATCH_SIZE, 0.0, 0.75).unwrap(), + ), + ( + "primitive_nullable", + create_primitive_bench_batch(BATCH_SIZE, 0.25, 0.75).unwrap(), + ), + ( + "primitive_sparse_99pct_null", + create_primitive_bench_batch(BATCH_SIZE, 0.99, 0.75).unwrap(), + ), + ( + "bool_non_null", + create_bool_bench_batch_non_null(BATCH_SIZE, 0.25, 0.75).unwrap(), + ), + ( + "string_non_null", + create_string_bench_batch_non_null(BATCH_SIZE, 0.0, 0.75).unwrap(), + ), + ( + "short_string_non_null", + create_short_string_bench_batch(BATCH_SIZE).unwrap(), + ), + ( + "byte_array_sparse_99pct_null", + create_byte_array_sparse_bench_batch(BATCH_SIZE, 0.99).unwrap(), + ), + ( + "fixed_size_binary_non_null", + create_fixed_size_binary_bench_batch(BATCH_SIZE, 16, 0.0).unwrap(), + ), + ( + "fixed_size_binary_sparse", + create_fixed_size_binary_bench_batch(BATCH_SIZE, 16, 0.25).unwrap(), + ), + ( + "decimal128_non_null", + create_decimal128_bench_batch(BATCH_SIZE, 0.0).unwrap(), + ), + ( + "decimal128_sparse", + create_decimal128_bench_batch(BATCH_SIZE, 0.25).unwrap(), + ), + ( + "int64_dictionary", + create_int64_dictionary_bench_batch(BATCH_SIZE, 0.25, 0.75).unwrap(), + ), + ( + "int64_dictionary_1pct_cardinality", + create_int64_dictionary_bench_batch_1pct_cardinality(BATCH_SIZE, 0.25).unwrap(), + ), + ( + "string_dictionary", + create_string_dictionary_bench_batch(BATCH_SIZE, 0.25, 0.75).unwrap(), + ), + ( + "string_dictionary_1pct_cardinality", + create_string_dictionary_bench_batch_1pct_cardinality(BATCH_SIZE, 0.25).unwrap(), + ), + ( + "decimal128_dictionary", + create_decimal128_dictionary_bench_batch(BATCH_SIZE, 0.25, 0.75).unwrap(), + ), + ( + "decimal128_dictionary_1pct_cardinality", + create_decimal128_dictionary_bench_batch_1pct_cardinality(BATCH_SIZE, 0.25).unwrap(), + ), + ( + "bool_dictionary", + create_bool_dictionary_bench_batch(BATCH_SIZE).unwrap(), + ), + ( + "string_ree_low_cardinality", + create_run_end_encoded_bench_batch(DataType::Int32, BATCH_SIZE, 256, |n| { + Arc::new(StringArray::from_iter_values( + (0..n).map(|i| format!("category_{:02}", i % 16)), + )) + }), + ), + ( + "int64_ree_low_cardinality", + create_run_end_encoded_bench_batch(DataType::Int32, BATCH_SIZE, 256, |n| { + Arc::new(Int64Array::from_iter_values( + (0..n).map(|i| (i % 16) as i64), + )) + }), + ), + ( + "fixed_size_binary_ree_low_cardinality", + create_run_end_encoded_bench_batch(DataType::Int32, BATCH_SIZE, 256, |n| { + Arc::new( + FixedSizeBinaryArray::try_from_iter((0..n).map(|i| { + let mut b = [0u8; 16]; + b[0] = (i % 16) as u8; + b + })) + .unwrap(), + ) + }), + ), + ( + "string_ree_of_dict", + create_ree_of_dict_bench_batch(BATCH_SIZE), + ), + ( + "numeric_ree_of_dict", + create_ree_of_numeric_dict_bench_batch(BATCH_SIZE), + ), + ( + "ree_struct_of_dict", + create_ree_struct_of_dict_bench_batch(BATCH_SIZE), + ), + ( + "int32_ree_99pct_null", + create_run_end_encoded_bench_batch(DataType::Int32, BATCH_SIZE, 256, |n| { + Arc::new(Int32Array::from_iter( + (0..n).map(|i| (i % 100 == 0).then_some((i % 16) as i32)), + )) + }), + ), + ( + "string_ree_99pct_null", + create_run_end_encoded_bench_batch(DataType::Int32, BATCH_SIZE, 256, |n| { + Arc::new(StringArray::from_iter((0..n).map(|i| { + (i % 100 == 0).then(|| format!("category_{:02}", i % 16)) + }))) + }), + ), + ( + "struct_ree", + create_run_end_encoded_bench_batch(DataType::Int32, BATCH_SIZE, 256, |n| { + make_ree_struct_values(n) + }), + ), + ( + "list_int32_ree", + create_run_end_encoded_bench_batch(DataType::Int32, BATCH_SIZE, 256, |n| { + make_ree_list_values(n) + }), + ), + ( + "list_of_int32_ree", + create_list_of_int32_ree_batch(BATCH_SIZE), + ), + ( + "struct_of_list_of_int32_ree_fragmented", + create_fragmented_list_of_int32_ree_batch(BATCH_SIZE), + ), + ( + "list_of_list_int32_ree", + create_list_of_list_int32_ree_batch(BATCH_SIZE), + ), + ( + "ree_list_ree_list", + create_ree_list_ree_list_batch(BATCH_SIZE), + ), + ] +} + +fn add_benches(c: &mut Criterion, measurement_name: &str) { + let batches = create_batches(); + let props = writer_props(); + + for benchmark in &batches { + let mut group = c.benchmark_group(format!( + "arrow_writer_{measurement_name}/{}", + benchmark.name + )); + let batch = OnceCell::new(); + for (prop_name, prop) in &props { + bench_cell(&mut group, prop_name, benchmark, &batch, prop); + } + group.finish(); + } +} + +fn add_peak_memory_benches(c: &mut Criterion) { + add_benches(c, "peak_memory"); +} + +fn add_allocated_bytes_benches(c: &mut Criterion) { + add_benches(c, "allocated_bytes"); +} + +fn peak_memory_criterion() -> Criterion { + Criterion::default() + .with_measurement(PeakMemory) + .sample_size(10) +} + +fn allocated_bytes_criterion() -> Criterion { + Criterion::default() + .with_measurement(AllocatedBytes) + .sample_size(10) +} + +criterion_group! { + name = peak_memory; + config = peak_memory_criterion(); + targets = add_peak_memory_benches +} + +criterion_group! { + name = cumulative_allocated_bytes; + config = allocated_bytes_criterion(); + targets = add_allocated_bytes_benches +} +criterion_main!(peak_memory, cumulative_allocated_bytes); diff --git a/parquet/benches/parquet_round_trip.rs b/parquet/benches/parquet_round_trip.rs index 825b32ea9586..bf82a8a952c1 100644 --- a/parquet/benches/parquet_round_trip.rs +++ b/parquet/benches/parquet_round_trip.rs @@ -332,8 +332,10 @@ fn read_write(c: &mut Criterion, spec: ParquetFileSpec, msg: &str) { file_from_spec(spec, &mut buffer); c.bench_function(&format!("write {msg}"), |b| { - buffer.clear(); - b.iter(|| file_from_spec(spec, &mut buffer)) + b.iter(|| { + buffer.clear(); + file_from_spec(spec, &mut buffer) + }) }); let file_bytes = Bytes::from(buffer); diff --git a/parquet/src/arrow/array_reader/byte_array.rs b/parquet/src/arrow/array_reader/byte_array.rs index ea9296c7cd6e..1641fd021977 100644 --- a/parquet/src/arrow/array_reader/byte_array.rs +++ b/parquet/src/arrow/array_reader/byte_array.rs @@ -414,12 +414,12 @@ impl ByteArrayDecoderPlain { self.offset = end_offset; read += 1; } - self.max_remaining_values -= to_read; + self.max_remaining_values -= read; if self.validate_utf8 { output.check_valid_utf8(initial_values_length)?; } - Ok(to_read) + Ok(read) } pub fn skip(&mut self, to_skip: usize) -> Result { @@ -753,4 +753,18 @@ mod tests { assert_eq!(decoder.skip_values(1024).unwrap(), 0); } } + + #[test] + fn test_plain_decoder_reports_values_actually_read() { + // The page claims to contain two values, but its buffer contains only + // one complete PLAIN-encoded BYTE_ARRAY value. + let buffer = Bytes::from_static(&[3, 0, 0, 0, b'f', b'o', b'o']); + let mut decoder = ByteArrayDecoderPlain::new(buffer, 2, Some(2), false); + let mut output = OffsetBuffer::::with_capacity(2); + + assert_eq!(decoder.read(&mut output, 2).unwrap(), 1); + assert_eq!(output.values.as_slice(), b"foo"); + assert_eq!(output.offsets.as_slice(), &[0, 3]); + assert_eq!(decoder.max_remaining_values, 1); + } } diff --git a/parquet/src/arrow/array_reader/cached_array_reader.rs b/parquet/src/arrow/array_reader/cached_array_reader.rs index 012fa2548170..980d67f31bf6 100644 --- a/parquet/src/arrow/array_reader/cached_array_reader.rs +++ b/parquet/src/arrow/array_reader/cached_array_reader.rs @@ -87,6 +87,9 @@ pub struct CachedArrayReader { local_cache: HashMap, /// Statistics to report on the Cache behavior metrics: ArrowReaderMetrics, + /// Exclusive upper bound of the batch ids already removed from the shared + /// cache by [`Self::cleanup_consumed_batches`]. + cleaned_up_to: usize, } impl CachedArrayReader { @@ -111,6 +114,7 @@ impl CachedArrayReader { role, local_cache: HashMap::new(), metrics, + cleaned_up_to: 0, } } @@ -168,22 +172,28 @@ impl CachedArrayReader { /// Remove batches from cache that have been completely consumed /// This is only called for Consumer role readers - fn cleanup_consumed_batches(&self) { + fn cleanup_consumed_batches(&mut self) { let current_batch_id = self.get_batch_id_from_position(self.outer_position); + // `outer_position` only moves forward, so the watermark can never run + // ahead of the batch the reader is currently on. + debug_assert!(current_batch_id.val >= self.cleaned_up_to); // Remove batches that are at least one batch behind the current position // This ensures we don't remove batches that might still be needed for the current batch // We can safely remove batch_id if current_batch_id > batch_id + 1 - if current_batch_id.val > 1 { - let mut cache = self.shared_cache.write().unwrap(); - for batch_id_to_remove in 0..(current_batch_id.val - 1) { - cache.remove( - self.column_idx, - BatchID { - val: batch_id_to_remove, - }, - ); - } + if current_batch_id.val <= 1 { + return; + } + let end = current_batch_id.val - 1; + // Everything below `cleaned_up_to` was removed by an earlier call. + if end <= self.cleaned_up_to { + return; + } + let start = self.cleaned_up_to; + self.cleaned_up_to = end; + let mut cache = self.shared_cache.write().unwrap(); + for val in start..end { + cache.remove(self.column_idx, BatchID { val }); } } } @@ -631,6 +641,65 @@ mod tests { assert!(cache.read().unwrap().get(0, BatchID { val: 2 }).is_some()); } + #[test] + fn test_consumer_cleanup_after_skip() { + // A consumer that skips several batches (e.g. rows filtered out by a + // predicate) must still remove the skipped batches from the shared + // cache, even though it never fetched them itself: one cleanup call + // covers the whole skipped range. + let metrics = ArrowReaderMetrics::disabled(); + let cache = Arc::new(RwLock::new(RowGroupCache::new(3, usize::MAX))); // Batch size 3 + + // Producer populates batches 0..=3 (12 values). + let producer_values: Vec = (1..=12).collect(); + let mut producer = CachedArrayReader::new( + Box::new(MockArrayReader::new(producer_values.clone())), + cache.clone(), + 0, + CacheRole::Producer, + metrics.clone(), + ); + for _ in 0..4 { + producer.read_records(3).unwrap(); + producer.consume_batch().unwrap(); + } + for batch in 0..4 { + assert!( + cache + .read() + .unwrap() + .get(0, BatchID { val: batch }) + .is_some() + ); + } + + // Consumer skips batches 0..=2 outright and reads batch 3. + let mut consumer = CachedArrayReader::new( + Box::new(MockArrayReader::new(producer_values)), + cache.clone(), + 0, + CacheRole::Consumer, + metrics, + ); + assert_eq!(consumer.skip_records(9).unwrap(), 9); + assert_eq!(consumer.read_records(3).unwrap(), 3); + let array = consumer.consume_batch().unwrap(); + assert_eq!(array.len(), 3); + + // current_batch_id = 12 / 3 = 4, so cleanup covers batches 0..3 in a + // single call, including the ones the consumer never fetched. + for batch in 0..3 { + assert!( + cache + .read() + .unwrap() + .get(0, BatchID { val: batch }) + .is_none() + ); + } + assert!(cache.read().unwrap().get(0, BatchID { val: 3 }).is_some()); + } + #[test] fn test_producer_keeps_batches() { let metrics = ArrowReaderMetrics::disabled(); diff --git a/parquet/src/arrow/arrow_reader/filter.rs b/parquet/src/arrow/arrow_reader/filter.rs index 3fd5e1d650be..75c2604a53a6 100644 --- a/parquet/src/arrow/arrow_reader/filter.rs +++ b/parquet/src/arrow/arrow_reader/filter.rs @@ -16,8 +16,12 @@ // under the License. use crate::arrow::ProjectionMask; -use arrow_array::{BooleanArray, RecordBatch}; +use crate::arrow::arrow_reader::{RowSelection, RowSelectionPolicy}; +use crate::schema::types::SchemaDescriptor; +use arrow_array::{Array, BooleanArray, RecordBatch}; +use arrow_buffer::BooleanBuffer; use arrow_schema::ArrowError; +use arrow_select::filter::{SlicesIterator, filter_record_batch, prep_null_mask_filter}; use std::fmt::{Debug, Formatter}; /// A predicate operating on [`RecordBatch`] @@ -153,7 +157,10 @@ where /// This design has a couple of implications: /// /// * [`RowFilter`] can be used to skip entire pages, and thus IO, in addition to CPU decode overheads -/// * Columns may be decoded multiple times if they appear in multiple [`ProjectionMask`] +/// * Columns may be decoded multiple times if they appear in multiple [`ProjectionMask`]. +/// Consecutive predicates whose projection is the same single top-level, non-repeated +/// column are evaluated together on one decoded stream, so ordering such predicates +/// next to each other avoids decoding that column again /// * IO will be deferred until needed by a [`ProjectionMask`] /// /// As such there is a trade-off between a single large predicate, or multiple predicates, @@ -198,4 +205,477 @@ impl RowFilter { pub fn into_predicates(self) -> Vec> { self.predicates } + + /// Fuse consecutive predicates on the same single top-level, non-repeated leaf. + /// This avoids repeated decoding or predicate-cache replay of that column. + pub(crate) fn fuse_same_projection(self, parquet_schema: &SchemaDescriptor) -> Self { + let mut predicates: Vec> = + Vec::with_capacity(self.predicates.len()); + let mut group: Vec> = Vec::new(); + let mut flush_group = |group: &mut Vec>| { + if group.len() > 1 && can_fuse_projection(group[0].projection(), parquet_schema) { + let group = std::mem::take(group); + predicates.push(Box::new(FusedPredicate::new(group))); + } else { + predicates.append(group); + } + }; + + for predicate in self.predicates { + if group + .last() + .is_some_and(|last| last.projection() != predicate.projection()) + { + flush_group(&mut group); + } + group.push(predicate); + } + flush_group(&mut group); + + Self { predicates } + } +} + +/// Restrict fusion to one top-level, non-repeated leaf to limit compaction costs. +fn can_fuse_projection(projection: &ProjectionMask, parquet_schema: &SchemaDescriptor) -> bool { + let mut leaf_indices = + (0..parquet_schema.num_columns()).filter(|idx| projection.leaf_included(*idx)); + let Some(leaf_idx) = leaf_indices.next() else { + return false; + }; + if leaf_indices.next().is_some() { + return false; + } + + let column = parquet_schema.column(leaf_idx); + column.path().parts().len() == 1 && column.max_rep_level() == 0 +} + +/// Evaluate same-projection predicates in order on one decoded batch. +/// Later predicates see only surviving rows, sliced when contiguous or compacted +/// otherwise. +struct FusedPredicate { + /// At least two predicates, all with the same projection. + predicates: Vec>, + /// Chooses how intermediate selections are composed. The output is always + /// a [`BooleanArray`], so this is independent of the reader's + /// [`RowSelectionPolicy`], which governs decoding. + composition_policy: RowSelectionPolicy, +} + +impl FusedPredicate { + /// Create a fused predicate from at least two predicates that share one + /// projection. + fn new(predicates: Vec>) -> Self { + debug_assert!(predicates.len() > 1); + debug_assert!( + predicates + .windows(2) + .all(|pair| pair[0].projection() == pair[1].projection()) + ); + Self { + predicates, + composition_policy: RowSelectionPolicy::default(), + } + } +} + +impl Debug for FusedPredicate { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + write!( + f, + "FusedPredicate {{ {} predicates }}", + self.predicates.len() + ) + } +} + +impl ArrowPredicate for FusedPredicate { + fn projection(&self) -> &ProjectionMask { + self.predicates[0].projection() + } + + fn evaluate(&mut self, batch: RecordBatch) -> Result { + let num_rows = batch.num_rows(); + // Positions in the original batch; None until a predicate rejects rows. + let mut selection: Option = None; + let mut filtered_batch = batch; + let last_predicate_idx = self.predicates.len() - 1; + + for (idx, predicate) in self.predicates.iter_mut().enumerate() { + let filter = evaluate_predicate(predicate.as_mut(), filtered_batch.clone())?; + // No mapping is needed if all preceding predicates accepted every row. + if idx == last_predicate_idx && selection.is_none() { + return Ok(filter); + } + let true_count = filter.true_count(); + if true_count == 0 { + return Ok(BooleanArray::new(BooleanBuffer::new_unset(num_rows), None)); + } + if true_count == filter.len() { + continue; + } + + let predicate_selection = RowSelection::from_boolean_buffer(filter.values().clone()); + selection = Some(match selection.take() { + // Only the accumulated selection drives the composition + // algorithm, so adapt it once, right before it is used. + Some(prev) => self + .composition_policy + .apply(prev) + .and_then(&predicate_selection), + None => predicate_selection, + }); + if idx != last_predicate_idx { + filtered_batch = narrow_batch(&filtered_batch, &filter, true_count)?; + } + } + + let mask = match selection { + Some(selection) => selection.into_boolean_buffer(), + None => BooleanBuffer::new_set(num_rows), + }; + debug_assert_eq!(mask.len(), num_rows); + Ok(BooleanArray::new(mask, None)) + } +} + +/// Validate the predicate result length and treat nulls as false. +fn evaluate_predicate( + predicate: &mut dyn ArrowPredicate, + batch: RecordBatch, +) -> Result { + let input_rows = batch.num_rows(); + let filter = predicate.evaluate(batch)?; + if filter.len() != input_rows { + return Err(ArrowError::InvalidArgumentError(format!( + "ArrowPredicate predicate returned {} rows, expected {input_rows}", + filter.len() + ))); + } + Ok(match filter.null_count() { + 0 => filter, + _ => prep_null_mask_filter(&filter), + }) +} + +/// Restrict `batch` to the rows `filter` accepts, slicing zero-copy when they +/// form one contiguous range and compacting otherwise. +/// +/// `filter` must be null-free and accept some but not all rows. +fn narrow_batch( + batch: &RecordBatch, + filter: &BooleanArray, + true_count: usize, +) -> Result { + let mut slices = SlicesIterator::new(filter); + let (start, end) = slices + .next() + .expect("a partially selected filter has a true slice"); + if end - start == true_count { + Ok(batch.slice(start, end - start)) + } else { + filter_record_batch(batch, filter) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::arrow::array_reader::ArrayReader; + use crate::arrow::array_reader::StructArrayReader; + use crate::arrow::array_reader::test_util::make_int32_page_reader; + use crate::arrow::arrow_reader::ReadPlanBuilder; + use crate::schema::parser::parse_message_type; + use arrow_array::Int32Array; + use arrow_array::cast::AsArray; + use arrow_array::types::Int32Type; + use arrow_schema::{DataType, Field, Fields, Schema}; + use std::sync::Arc; + + fn test_schema() -> SchemaDescriptor { + let schema = parse_message_type( + "message schema { + REQUIRED INT32 a; + REQUIRED INT32 b; + REQUIRED INT32 c; + REQUIRED GROUP nested { REQUIRED INT32 d; } + REPEATED INT32 e; + }", + ) + .unwrap(); + SchemaDescriptor::new(Arc::new(schema)) + } + + #[test] + fn can_fuse_projection_requires_one_top_level_leaf() { + let schema = test_schema(); + + let one_leaf = ProjectionMask::leaves(&schema, [1]); + let two_leaves = ProjectionMask::leaves(&schema, [0, 2]); + let nested_leaf = ProjectionMask::leaves(&schema, [3]); + let repeated_leaf = ProjectionMask::leaves(&schema, [4]); + assert!(can_fuse_projection(&one_leaf, &schema)); + assert!(!can_fuse_projection(&two_leaves, &schema)); + assert!(!can_fuse_projection(&nested_leaf, &schema)); + assert!(!can_fuse_projection(&repeated_leaf, &schema)); + assert!(!can_fuse_projection( + &ProjectionMask::none(schema.num_columns()), + &schema + )); + } + + fn accept_all(projection: ProjectionMask) -> Box { + Box::new(ArrowPredicateFn::new(projection, |batch| { + Ok(BooleanArray::from(vec![true; batch.num_rows()])) + })) + } + + #[test] + fn fuse_same_projection_groups_consecutive_eligible_predicates() { + let schema = test_schema(); + let a = ProjectionMask::leaves(&schema, [0]); + let b = ProjectionMask::leaves(&schema, [1]); + let ac = ProjectionMask::leaves(&schema, [0, 2]); + let nested = ProjectionMask::leaves(&schema, [3]); + + let filter = RowFilter::new(vec![ + // fused + accept_all(a.clone()), + accept_all(a.clone()), + // single predicate: kept as-is + accept_all(b.clone()), + // two leaves: not eligible for fusion + accept_all(ac.clone()), + accept_all(ac.clone()), + // nested leaf: not eligible for fusion + accept_all(nested.clone()), + accept_all(nested.clone()), + // fused + accept_all(a.clone()), + accept_all(a.clone()), + accept_all(a.clone()), + ]) + .fuse_same_projection(&schema); + + let projections: Vec<_> = filter + .predicates() + .iter() + .map(|predicate| predicate.projection().clone()) + .collect(); + assert_eq!( + projections, + vec![a.clone(), b, ac.clone(), ac, nested.clone(), nested, a] + ); + assert_eq!(format!("{filter:?}"), "RowFilter { 7 predicates: }"); + } + + #[test] + fn fuse_same_projection_keeps_single_predicates_and_empty_filters() { + let schema = test_schema(); + let a = ProjectionMask::leaves(&schema, [0]); + let b = ProjectionMask::leaves(&schema, [1]); + + let filter = + RowFilter::new(vec![accept_all(a), accept_all(b)]).fuse_same_projection(&schema); + assert_eq!(filter.predicates().len(), 2); + + let filter = RowFilter::new(vec![]).fuse_same_projection(&schema); + assert!(filter.predicates().is_empty()); + } + + fn int_batch(values: impl IntoIterator) -> RecordBatch { + let schema = Arc::new(Schema::new(vec![Field::new("v", DataType::Int32, false)])); + let values = Int32Array::from_iter_values(values); + RecordBatch::try_new(schema, vec![Arc::new(values)]).unwrap() + } + + /// A predicate over the `v` column that asserts `precondition` holds for + /// every row it sees before returning `f(v)`. + fn int_predicate( + precondition: impl Fn(i32) -> bool + Send + 'static, + f: impl Fn(i32) -> Option + Send + 'static, + ) -> Box { + Box::new(ArrowPredicateFn::new(ProjectionMask::all(), move |batch| { + let values = batch.column(0).as_primitive::(); + assert!( + values.values().iter().all(|v| precondition(*v)), + "predicate saw rows rejected by a predecessor" + ); + Ok(values.values().iter().map(|v| f(*v)).collect()) + })) + } + + /// Fragmented survivors: roughly every other row. + fn even() -> Box { + int_predicate(|_| true, |v| Some(v % 2 == 0)) + } + + /// Clustered survivors: one contiguous range. + fn in_range() -> Box { + int_predicate(|_| true, |v| Some((10..40).contains(&v))) + } + + /// Applies after `even`: mixes nulls (rejected) into the result. + fn divisible_by_three_after_even() -> Box { + int_predicate( + |v| v % 2 == 0, + |v| if v % 7 == 0 { None } else { Some(v % 3 == 0) }, + ) + } + + fn all_true() -> Box { + int_predicate(|_| true, |_| Some(true)) + } + + fn all_false() -> Box { + int_predicate(|_| true, |_| Some(false)) + } + + fn never_called() -> Box { + Box::new(ArrowPredicateFn::new(ProjectionMask::all(), |_| { + panic!("predicate evaluated after every row was rejected") + })) + } + + /// Evaluate `predicates` one after another, each on the rows kept by its + /// predecessors, the way `RowFilter` applies separate predicates. + fn sequential(batch: &RecordBatch, predicates: &mut [Box]) -> Vec { + let mut kept = vec![true; batch.num_rows()]; + for predicate in predicates { + // Like `RowFilter`, stop once every row has been rejected. + if !kept.contains(&true) { + break; + } + let survivors = filter_record_batch(batch, &BooleanArray::from(kept.clone())).unwrap(); + let filter = predicate.evaluate(survivors).unwrap(); + for (filter_idx, keep) in kept.iter_mut().filter(|keep| **keep).enumerate() { + *keep = filter.is_valid(filter_idx) && filter.value(filter_idx); + } + } + kept + } + + const POLICIES: [RowSelectionPolicy; 4] = [ + RowSelectionPolicy::Selectors, + RowSelectionPolicy::Mask, + RowSelectionPolicy::Auto { threshold: 32 }, + RowSelectionPolicy::Auto { threshold: 2 }, + ]; + + type PredicateChain = fn() -> Vec>; + + #[test] + fn fused_predicate_matches_sequential_evaluation() { + let cases: Vec = vec![ + || vec![even(), divisible_by_three_after_even()], + || vec![even(), in_range(), divisible_by_three_after_even()], + || vec![in_range(), even()], + || vec![even(), all_true(), divisible_by_three_after_even()], + || vec![all_true(), even()], + || vec![all_true(), all_true()], + || vec![even(), all_false(), never_called()], + || vec![all_false(), never_called()], + ]; + + for num_rows in [0, 1, 7, 97, 200] { + let batch = int_batch(0..num_rows); + for make_predicates in &cases { + let expected = sequential(&batch, &mut make_predicates()); + for policy in POLICIES { + let mut fused = FusedPredicate { + predicates: make_predicates(), + composition_policy: policy, + }; + let actual = fused.evaluate(batch.clone()).unwrap(); + assert_eq!(actual.null_count(), 0); + assert_eq!( + actual.values().iter().collect::>(), + expected, + "{num_rows} rows, {policy:?}" + ); + } + } + } + } + + #[test] + fn fused_predicate_rejects_wrong_result_length() { + let short = Box::new(ArrowPredicateFn::new(ProjectionMask::all(), |_| { + Ok(BooleanArray::from(vec![true])) + })); + let mut fused = FusedPredicate::new(vec![all_true(), short]); + let err = fused.evaluate(int_batch(0..4)).unwrap_err(); + assert!( + err.to_string() + .contains("ArrowPredicate predicate returned 1 rows, expected 4"), + "{err}" + ); + } + + #[test] + fn narrow_batch_slices_contiguous_survivors() { + let batch = int_batch(0..6); + let filter = BooleanArray::from(vec![false, false, true, true, true, false]); + let narrowed = narrow_batch(&batch, &filter, 3).unwrap(); + assert_eq!( + narrowed.column(0).as_primitive::().values(), + &[2, 3, 4] + ); + // Slicing shares the input buffer instead of copying it. + assert_eq!(narrowed.column(0).to_data().buffers()[0].as_ptr(), unsafe { + batch.column(0).to_data().buffers()[0].as_ptr().add(2 * 4) + }); + + let filter = BooleanArray::from(vec![false, true, false, true, false, false]); + let narrowed = narrow_batch(&batch, &filter, 2).unwrap(); + assert_eq!( + narrowed.column(0).as_primitive::().values(), + &[1, 3] + ); + } + + /// Fusing predicates into one `ReadPlanBuilder::with_predicate` call + /// produces the same selection as applying them one at a time. + #[test] + fn fused_predicate_matches_predicate_major_read_plan() { + let data: Vec = (0..97).collect(); + let make_reader = || { + let levels = vec![0; data.len()]; + let leaf = make_int32_page_reader(&data, &levels, &levels, 0, 0, None); + let struct_type = + DataType::Struct(Fields::from(vec![Field::new("c0", DataType::Int32, false)])); + Box::new(StructArrayReader::new( + struct_type, + vec![leaf], + 0, + 0, + false, + None, + )) as Box + }; + let make_predicates = || vec![even(), divisible_by_three_after_even()]; + + let prior = RowSelection::from_filters(&[BooleanArray::from( + (0..data.len()).map(|idx| idx % 5 != 0).collect::>(), + )]); + for initial in [None, Some(prior)] { + let mut sequential = ReadPlanBuilder::new(7).with_selection(initial.clone()); + for predicate in &mut make_predicates() { + sequential = sequential + .with_predicate(make_reader(), predicate.as_mut()) + .unwrap(); + } + + for policy in POLICIES { + let mut fused = FusedPredicate::new(make_predicates()); + let plan = ReadPlanBuilder::new(7) + .with_selection(initial.clone()) + .with_row_selection_policy(policy) + .with_predicate(make_reader(), &mut fused) + .unwrap(); + assert_eq!(plan.selection(), sequential.selection(), "{policy:?}"); + } + } + } } diff --git a/parquet/src/arrow/arrow_reader/read_plan.rs b/parquet/src/arrow/arrow_reader/read_plan.rs index 5ebfbc489054..6daabd7e184f 100644 --- a/parquet/src/arrow/arrow_reader/read_plan.rs +++ b/parquet/src/arrow/arrow_reader/read_plan.rs @@ -20,7 +20,7 @@ use crate::arrow::array_reader::ArrayReader; use crate::arrow::arrow_reader::selection::{ - LoadedRowRanges, RowSelectionInner, RowSelectionPolicy, RowSelectionStrategy, + LoadedRowRanges, RowSelectionPolicy, RowSelectionStrategy, }; use crate::arrow::arrow_reader::{ ArrowPredicate, ParquetRecordBatchReader, RowSelection, RowSelectionCursor, RowSelector, @@ -157,17 +157,7 @@ impl ReadPlanBuilder { /// /// Guarantees to return either `Selectors` or `Mask`, never `Auto`. pub(crate) fn resolve_selection_strategy(&self) -> RowSelectionStrategy { - match self.row_selection_policy { - RowSelectionPolicy::Selectors => RowSelectionStrategy::Selectors, - RowSelectionPolicy::Mask => RowSelectionStrategy::Mask, - RowSelectionPolicy::Auto { threshold, .. } => { - let Some(selection) = self.selection.as_ref() else { - return RowSelectionStrategy::Selectors; - }; - - selection.auto_selection_strategy(threshold) - } - } + self.row_selection_policy.resolve(self.selection.as_ref()) } /// Evaluates an [`ArrowPredicate`], updating this plan's `selection` @@ -306,7 +296,7 @@ impl ReadPlanBuilder { } = self; let row_selection_cursor = selection - .map(|s| build_cursor(s.trim(), selection_strategy, loaded_row_ranges)) + .map(|s| selection_strategy.build_cursor(s.trim(), loaded_row_ranges)) .unwrap_or_else(RowSelectionCursor::new_all); ReadPlan { @@ -316,28 +306,6 @@ impl ReadPlanBuilder { } } -/// Lower a [`RowSelection`] to the cursor form requested by the resolved strategy. -fn build_cursor( - selection: RowSelection, - strategy: RowSelectionStrategy, - loaded_row_ranges: Option>, -) -> RowSelectionCursor { - match (strategy, selection.into_inner()) { - (RowSelectionStrategy::Mask, RowSelectionInner::Mask(mask)) => { - RowSelectionCursor::new_mask_from_buffer((*mask).into_mask(), loaded_row_ranges) - } - (RowSelectionStrategy::Mask, RowSelectionInner::Selectors(selectors)) => { - RowSelectionCursor::new_mask_from_selectors(selectors, loaded_row_ranges) - } - (RowSelectionStrategy::Selectors, RowSelectionInner::Selectors(selectors)) => { - RowSelectionCursor::new_selectors(selectors) - } - (RowSelectionStrategy::Selectors, RowSelectionInner::Mask(mask)) => { - RowSelectionCursor::new_selectors((*mask).into_selectors()) - } - } -} - /// Builder for [`ReadPlan`] that applies a limit and offset to the read plan /// /// See [`ReadPlanBuilder::limited`] to create this builder. diff --git a/parquet/src/arrow/arrow_reader/selection/algebra.rs b/parquet/src/arrow/arrow_reader/selection/algebra.rs index e0fda42b821e..0f9b445fa260 100644 --- a/parquet/src/arrow/arrow_reader/selection/algebra.rs +++ b/parquet/src/arrow/arrow_reader/selection/algebra.rs @@ -27,6 +27,16 @@ use arrow_buffer::{BooleanBuffer, BooleanBufferBuilder, MutableBuffer, bit_util} use std::cmp::Ordering; use std::iter::Peekable; +// Use word-at-a-time expansion for masks with at least 8192 rows, roughly +// 75% outer selectivity and 5% inner selectivity. Smaller or sparser inputs +// use set indices to avoid scanning every output word. +// In `outer.and_then(inner)`, outer selects from the original rows, and inner +// selects from those selected rows: inner.len() == outer.count_set_bits(). +// The selectivity constants are divisors: at most 1/4 dropped, at least 1/20 kept. +const AND_THEN_DENSE_MASK_MIN_LEN: usize = 8192; +const AND_THEN_DENSE_MASK_MAX_DROPPED_FRACTION: usize = 4; +const AND_THEN_DENSE_MASK_MIN_INNER_FRACTION: usize = 20; + /// Applies `second` to the rows selected by `first`, both selector-backed. pub(super) fn and_then_row_selections( first: &[RowSelector], @@ -405,6 +415,26 @@ fn and_then_masks(mask: &BooleanBuffer, other: &BooleanBuffer) -> BooleanBuffer return mask.clone(); } + if should_use_dense_mask(mask.len(), selected_count, other_true_count) { + and_then_dense_masks(mask, other) + } else { + and_then_sparse_masks(mask, other) + } +} + +/// Checks the length and selectivity thresholds for dense expansion. +#[inline] +fn should_use_dense_mask(mask_len: usize, selected_count: usize, other_true_count: usize) -> bool { + mask_len >= AND_THEN_DENSE_MASK_MIN_LEN + && mask_len - selected_count <= mask_len.div_ceil(AND_THEN_DENSE_MASK_MAX_DROPPED_FRACTION) + && other_true_count >= selected_count / AND_THEN_DENSE_MASK_MIN_INNER_FRACTION +} + +/// Maps each set bit in `other` to the corresponding set bit in `mask`. +/// Requires `other.len() == mask.count_set_bits()`. +fn and_then_sparse_masks(mask: &BooleanBuffer, other: &BooleanBuffer) -> BooleanBuffer { + debug_assert_eq!(other.len(), mask.count_set_bits()); + let mut builder = BooleanBufferBuilder::new(mask.len()); let mut outer_set_indices = mask.set_indices(); let mut next_selected_ordinal = 0usize; @@ -430,11 +460,51 @@ fn and_then_masks(mask: &BooleanBuffer, other: &BooleanBuffer) -> BooleanBuffer builder.finish() } +/// Scatters the next `mask_word.count_ones()` bits of `other` into each mask word. +/// Requires `other.len() == mask.count_set_bits()`. +#[inline(never)] +fn and_then_dense_masks(mask: &BooleanBuffer, other: &BooleanBuffer) -> BooleanBuffer { + debug_assert_eq!(other.len(), mask.count_set_bits()); + + let mut other_chunks = other.bit_chunks().iter_padded(); + let mut other_remaining = other.len(); + let mut pending = 0_u128; + let mut pending_len = 0; + let mut output = MutableBuffer::with_capacity(mask.len().div_ceil(8)); + + for mask_word in mask + .bit_chunks() + .iter_padded() + .take(mask.len().div_ceil(64)) + { + let selected = mask_word.count_ones() as usize; + while pending_len < selected { + let chunk = other_chunks + .next() + .expect("validated other length matches selected row count"); + let chunk_len = other_remaining.min(64); + pending |= (chunk as u128) << pending_len; + pending_len += chunk_len; + other_remaining -= chunk_len; + } + + let values = pending as u64; + output.extend_from_slice(&bit_util::expand(values, mask_word).to_le_bytes()); + pending >>= selected; + pending_len -= selected; + } + + debug_assert_eq!(other_remaining, 0); + debug_assert_eq!(pending_len, 0); + output.truncate(mask.len().div_ceil(8)); + BooleanBuffer::new(output.into(), 0, mask.len()) +} + #[cfg(test)] mod tests { use super::*; use arrow_array::BooleanArray; - use rand::{RngExt, rng}; + use rand::{RngExt, SeedableRng, rng, rngs::StdRng}; #[test] fn test_and() { @@ -728,6 +798,108 @@ mod tests { ); } + struct AndThenTest { + outer: BooleanBuffer, + inner: BooleanBuffer, + dense_mask_expected: bool, + } + + impl AndThenTest { + fn run(self) { + let Self { + outer, + inner, + dense_mask_expected, + } = self; + assert_eq!(inner.len(), outer.count_set_bits()); + assert_eq!( + should_use_dense_mask(outer.len(), inner.len(), inner.count_set_bits()), + dense_mask_expected + ); + + let mut inner_idx = 0; + let expected = BooleanBuffer::from_iter((0..outer.len()).map(|i| { + if !outer.value(i) { + return false; + } + let value = inner.value(inner_idx); + inner_idx += 1; + value + })); + + if dense_mask_expected { + assert_eq!(and_then_dense_masks(&outer, &inner), expected); + } + + let outer = RowSelection::from_boolean_buffer(outer); + let inner = RowSelection::from_boolean_buffer(inner); + let actual = outer.and_then(&inner); + assert_eq!(actual.as_mask().unwrap(), &expected); + } + } + + #[test] + fn test_dense_mask_and_then_mask_with_offsets() { + for len in [8192, 8193, 10_000] { + for outer_stride in [100, 20, 4] { + let outer_offset = 3; + let outer = BooleanBuffer::from_iter( + (0..len + outer_offset).map(|i| (i + 1) % outer_stride != 0), + ) + .slice(outer_offset, len); + + let inner_offset = 5; + let inner_len = outer.count_set_bits(); + let inner = BooleanBuffer::from_iter( + (0..inner_len + inner_offset).map(|i| (i + 2) % 97 != 0), + ) + .slice(inner_offset, inner_len); + + AndThenTest { + outer, + inner, + dense_mask_expected: true, + } + .run(); + } + } + } + + #[test] + fn test_dense_mask_and_then_mask_randomized_word_boundaries() { + let mut rng = StdRng::seed_from_u64(0x67e4_a91b_239d_c805); + + for (len, outer_offset, inner_offset) in [(8192, 0, 0), (8192, 3, 5), (8193, 7, 1)] { + let outer = + BooleanBuffer::from_iter((0..len + outer_offset).map(|_| rng.random_bool(0.8))) + .slice(outer_offset, len); + let inner_len = outer.count_set_bits(); + let inner = BooleanBuffer::from_iter( + (0..inner_len + inner_offset).map(|_| rng.random_bool(0.5)), + ) + .slice(inner_offset, inner_len); + + AndThenTest { + outer, + inner, + dense_mask_expected: true, + } + .run(); + } + } + + #[test] + fn test_dense_mask_and_then_policy() { + let len = 8192; + let selected = len * 3 / 4; + let min_inner = selected / AND_THEN_DENSE_MASK_MIN_INNER_FRACTION; + + assert!(!should_use_dense_mask(len - 1, selected, min_inner)); + assert!(!should_use_dense_mask(len, selected - 1, min_inner)); + assert!(!should_use_dense_mask(len, selected, min_inner - 1)); + assert!(should_use_dense_mask(len, selected, min_inner)); + } + #[test] fn test_selector_and_then_mask() { let outer = @@ -990,4 +1162,42 @@ mod tests { ); } } + + #[test] + fn dense_composition_with_zero_words_and_offsets() { + for mask_len in [8192, 8193, 65537] { + for (mask_offset, inner_offset) in [(0, 0), (3, 5), (63, 65), (65, 129)] { + for stride in [4, 100] { + let mask = + BooleanBuffer::from_iter((0..mask_offset + mask_len + 73).map(|i| { + !(mask_offset..mask_offset + mask_len).contains(&i) + || (i - mask_offset) % stride != 0 + })) + .slice(mask_offset, mask_len); + let len = mask.count_set_bits(); + for start in [0, 63, len / 2] { + for select_last in [false, true] { + let count = len / 10; + let inner = + BooleanBuffer::from_iter((0..inner_offset + len + 73).map(|i| { + if !(inner_offset..inner_offset + len).contains(&i) { + return true; + } + let j = i - inner_offset; + (start..start + count).contains(&j) + || select_last && j == len - 1 + })) + .slice(inner_offset, len); + AndThenTest { + outer: mask.clone(), + inner, + dense_mask_expected: true, + } + .run(); + } + } + } + } + } + } } diff --git a/parquet/src/arrow/arrow_reader/selection/cursor.rs b/parquet/src/arrow/arrow_reader/selection/cursor.rs index 4d08197718eb..2be8b0d48fc2 100644 --- a/parquet/src/arrow/arrow_reader/selection/cursor.rs +++ b/parquet/src/arrow/arrow_reader/selection/cursor.rs @@ -22,7 +22,6 @@ //! matching [`RowSelectionCursor`], which keeps the per-reader position while //! the selection itself stays immutable. -use super::boolean::boolean_mask_from_selectors; use super::{RowSelection, RowSelector}; use crate::errors::ParquetError; use arrow_array::BooleanArray; @@ -51,6 +50,38 @@ impl Default for RowSelectionPolicy { } } +impl RowSelectionPolicy { + /// Resolve this policy for a selection without changing its backing. + pub(crate) fn resolve(self, selection: Option<&RowSelection>) -> RowSelectionStrategy { + match self { + Self::Selectors => RowSelectionStrategy::Selectors, + Self::Mask => RowSelectionStrategy::Mask, + Self::Auto { threshold } => selection + .map(|selection| selection.auto_selection_strategy(threshold)) + .unwrap_or(RowSelectionStrategy::Selectors), + } + } + + /// Materialize a selection according to this policy, preserving its full + /// length, including trailing skips needed for selection composition. + pub(crate) fn apply(self, selection: RowSelection) -> RowSelection { + match ( + self.resolve(Some(&selection)), + selection.as_mask().is_some(), + ) { + (RowSelectionStrategy::Mask, true) | (RowSelectionStrategy::Selectors, false) => { + selection + } + (RowSelectionStrategy::Mask, false) => { + RowSelection::from_boolean_buffer(selection.into_boolean_buffer()) + } + (RowSelectionStrategy::Selectors, true) => { + RowSelection::from(Vec::::from(selection)) + } + } + } +} + /// Fully resolved strategy for materializing [`RowSelection`] during execution. /// /// This is determined by [`RowSelectionPolicy`], including selector density for @@ -63,6 +94,24 @@ pub(crate) enum RowSelectionStrategy { Mask, } +impl RowSelectionStrategy { + /// Build a cursor using this strategy. The selection must be trimmed of + /// trailing skips before constructing a mask cursor. + pub(crate) fn build_cursor( + self, + selection: RowSelection, + loaded_row_ranges: Option>, + ) -> RowSelectionCursor { + match self { + Self::Mask => RowSelectionCursor::new_mask_from_buffer( + selection.into_boolean_buffer(), + loaded_row_ranges, + ), + Self::Selectors => RowSelectionCursor::new_selectors(selection.into()), + } + } +} + /// Cursor for iterating a [`RowSelection`] during execution within a /// [`ReadPlan`](crate::arrow::arrow_reader::ReadPlan). /// @@ -79,25 +128,6 @@ pub enum RowSelectionCursor { } impl RowSelectionCursor { - /// Create a [`MaskCursor`] cursor backed by a bitmask, from an existing set of selectors - pub(crate) fn new_mask_from_selectors( - selectors: Vec, - loaded_row_ranges: Option>, - ) -> Self { - debug_assert!( - selectors - .last() - .map(|selector| !selector.skip) - .unwrap_or(true), - "Mask selectors must not end with a skip" - ); - Self::Mask(MaskCursor { - mask: boolean_mask_from_selectors(&selectors), - position: 0, - loaded_row_ranges, - }) - } - /// Create a [`MaskCursor`] cursor backed by an existing bitmask. pub(crate) fn new_mask_from_buffer( mask: BooleanBuffer, @@ -377,14 +407,55 @@ impl LoadedRowRanges { mod tests { use super::*; + #[test] + fn selection_respects_policy() { + let mask = BooleanBuffer::from(vec![true, true, false, false, true, true, false, false]); + for (policy, expect_mask) in [ + (RowSelectionPolicy::Mask, true), + (RowSelectionPolicy::Selectors, false), + (RowSelectionPolicy::Auto { threshold: 2 }, false), + (RowSelectionPolicy::Auto { threshold: 3 }, true), + ] { + for selection in [ + RowSelection::from_boolean_buffer(mask.clone()), + RowSelection::from(vec![ + RowSelector::select(2), + RowSelector::skip(2), + RowSelector::select(2), + RowSelector::skip(2), + ]), + ] { + let selection = policy.apply(selection); + assert_eq!(selection.as_mask().is_some(), expect_mask, "{policy:?}"); + assert_eq!(selection.into_boolean_buffer(), mask); + } + } + } + + #[test] + fn adaptive_selection_uses_selectors_for_long_runs() { + let mask = BooleanBuffer::from_iter((0..1_024).map(|idx| (256..768).contains(&idx))); + let selection = + RowSelectionPolicy::default().apply(RowSelection::from_boolean_buffer(mask)); + assert!(selection.as_mask().is_none()); + } + + #[test] + fn adaptive_selection_keeps_fragmented_masks() { + let mask = BooleanBuffer::from_iter((0..1_024).map(|idx| idx % 2 == 0)); + let selection = + RowSelectionPolicy::default().apply(RowSelection::from_boolean_buffer(mask)); + assert!(selection.as_mask().is_some()); + } + #[test] fn test_loaded_mask_chunk_stops_at_trimmed_mask_end() { let loaded = LoadedRowRanges::from_selection(RowSelection::from_consecutive_ranges( std::iter::once(0..5), 10, )); - let RowSelectionCursor::Mask(mut cursor) = RowSelectionCursor::new_mask_from_selectors( - vec![RowSelector::select(1)], + let RowSelectionCursor::Mask(mut cursor) = RowSelectionStrategy::Mask.build_cursor( + RowSelection::from(vec![RowSelector::select(1)]), Some(loaded.into()), ) else { unreachable!() @@ -397,13 +468,13 @@ mod tests { #[test] fn test_next_mask_chunk_until_cursor_is_empty() { - let RowSelectionCursor::Mask(mut cursor) = RowSelectionCursor::new_mask_from_selectors( - vec![ + let RowSelectionCursor::Mask(mut cursor) = RowSelectionStrategy::Mask.build_cursor( + RowSelection::from(vec![ RowSelector::skip(2), RowSelector::select(2), RowSelector::skip(1), RowSelector::select(1), - ], + ]), None, ) else { unreachable!() diff --git a/parquet/src/arrow/arrow_reader/selection/mod.rs b/parquet/src/arrow/arrow_reader/selection/mod.rs index 74e6727255a7..df9c4cde08f8 100644 --- a/parquet/src/arrow/arrow_reader/selection/mod.rs +++ b/parquet/src/arrow/arrow_reader/selection/mod.rs @@ -49,7 +49,8 @@ use algebra::{ }; pub use boolean::MaskRunIter; use boolean::{ - MaskSelection, limit_mask, mask_has_at_least_runs, offset_mask, split_off_mask, trim_mask, + MaskSelection, boolean_mask_from_selectors, limit_mask, mask_has_at_least_runs, offset_mask, + split_off_mask, trim_mask, }; pub(crate) use cursor::{LoadedRowRanges, MaskCursor, RowSelectionStrategy}; pub use cursor::{RowSelectionCursor, RowSelectionPolicy}; @@ -238,6 +239,14 @@ impl RowSelection { self.inner } + /// Consume this selection and return its bitmap representation. + pub(crate) fn into_boolean_buffer(self) -> BooleanBuffer { + match self.inner { + RowSelectionInner::Mask(mask) => mask.into_mask(), + RowSelectionInner::Selectors(selectors) => boolean_mask_from_selectors(&selectors), + } + } + /// Choose the automatic materialisation strategy without converting between /// selector and mask backing. #[inline] diff --git a/parquet/src/arrow/arrow_reader/tests/row_selection.rs b/parquet/src/arrow/arrow_reader/tests/row_selection.rs index d9286f84dbfc..55310d0a1cbc 100644 --- a/parquet/src/arrow/arrow_reader/tests/row_selection.rs +++ b/parquet/src/arrow/arrow_reader/tests/row_selection.rs @@ -53,13 +53,12 @@ fn get_expected_batches( if let Some(last_start) = last_start.take() { expected_batches.push(column.slice(last_start, row_offset - last_start)) } - row_offset += to_read } false => { last_start.get_or_insert(row_offset); - row_offset += to_read } } + row_offset += to_read; } } diff --git a/parquet/src/arrow/arrow_writer/mod.rs b/parquet/src/arrow/arrow_writer/mod.rs index 6b5f29d0f113..a72584bbd807 100644 --- a/parquet/src/arrow/arrow_writer/mod.rs +++ b/parquet/src/arrow/arrow_writer/mod.rs @@ -401,14 +401,10 @@ impl ArrowWriter { continue; } - if let Some(avg_row_bytes) = current_bytes - .checked_div(buffered_rows) - .filter(|avg_row_bytes| *avg_row_bytes > 0) + let avg_row_bytes = current_bytes / buffered_rows; + if let Some(rows_that_fit) = (max_bytes - current_bytes).checked_div(avg_row_bytes) { // At this point, `current_bytes < max_bytes` (checked above) - let remaining_bytes = max_bytes - current_bytes; - let rows_that_fit = remaining_bytes.checked_div(avg_row_bytes).unwrap_or(0); - if candidate_rows > rows_that_fit { if rows_that_fit > 0 { split_at = Some(rows_that_fit); diff --git a/parquet/src/arrow/push_decoder/mod.rs b/parquet/src/arrow/push_decoder/mod.rs index 24561c33b4a4..57cb2b92128b 100644 --- a/parquet/src/arrow/push_decoder/mod.rs +++ b/parquet/src/arrow/push_decoder/mod.rs @@ -320,6 +320,11 @@ impl ParquetPushDecoderBuilder { max_predicate_cache_size, } = self; + // Evaluate eligible same-projection predicates as one predicate. + let filter = filter.map(|filter| { + filter.fuse_same_projection(parquet_metadata.file_metadata().schema_descr()) + }); + let has_predicates = filter .as_ref() .is_some_and(|filter| !filter.predicates.is_empty()); @@ -972,7 +977,7 @@ mod test { }; use crate::arrow::{ArrowWriter, ProjectionMask}; use crate::errors::ParquetError; - use crate::file::metadata::ParquetMetaDataPushDecoder; + use crate::file::metadata::{PageIndexPolicy, ParquetMetaDataPushDecoder}; use crate::file::properties::WriterProperties; use arrow::compute::kernels::cmp::{gt, lt}; use arrow_array::cast::AsArray; @@ -1144,6 +1149,74 @@ mod test { expect_finished(decoder.try_decode()); } + /// Push only the pages needed for each batch of a row group rather than + /// the entire row group, and check whether the decoder can produce that + /// batch before the rest of the row group's pages have been pushed. + #[test] + fn test_decoder_first_pages_only() { + let metadata = test_file_parquet_metadata_with_offset_index(); + let mut decoder = ParquetPushDecoderBuilder::try_new_decoder(Arc::clone(&metadata)) + .unwrap() + // Each data page has 100 rows, so a batch of 100 rows needs only + // the first data page of each column + .with_batch_size(100) + .build() + .unwrap(); + + // Row group 0: the decoder asks for the entire column chunk of each + // of the three columns "a", "b", and "c" + let ranges = expect_needs_data(decoder.try_decode()); + assert_eq!(ranges, vec![4..1860, 1860..3716, 3716..11062]); + + // Compute the ranges that cover only the first data page of each + // column (and the dictionary page which precedes it) + let page_index = metadata.page_index_for_row_group(0); + let row_group = metadata.row_group(0); + let mut first_page_ranges = vec![]; + let mut second_page_ranges = vec![]; + for (idx, column) in row_group.columns().iter().enumerate() { + let (start, len) = column.byte_range(); + let locations = page_index.page_locations(idx).unwrap(); + assert_eq!(locations.len(), 2, "expected 2 data pages per column chunk"); + let second_page_start = locations[1].offset as u64; + first_page_ranges.push(start..second_page_start); + second_page_ranges.push(second_page_start..start + len); + } + // Note the first range for each column includes the dictionary page + assert_eq!(first_page_ranges, vec![4..1734, 1860..3590, 3716..10936]); + assert_eq!( + second_page_ranges, + vec![1734..1860, 3590..3716, 10936..11062] + ); + + // Push only the first page of each column. This is all the data + // needed to decode the first batch of 100 rows. + push_ranges_to_decoder(&mut decoder, first_page_ranges); + + // Note will likely change as part of + // + + // decoder still reports it needs the (complete) ranges + // it originally asked for, and does not produce a batch. + let ranges = expect_needs_data(decoder.try_decode()); + assert_eq!(ranges, vec![4..1860, 1860..3716, 3716..11062]); + + // Pushing the second pages as separate ranges does not help either: + // the decoder does not coalesce adjacent pushed ranges, so a requested + // range is only satisfied by a single pushed buffer that covers it. + push_ranges_to_decoder(&mut decoder, second_page_ranges); + let ranges = expect_needs_data(decoder.try_decode()); + assert_eq!(ranges, vec![4..1860, 1860..3716, 3716..11062]); + + // Only once the exact ranges originally requested are pushed does the + // decoder produce batches. + push_ranges_to_decoder(&mut decoder, ranges); + let batch = expect_data(decoder.try_decode()); + assert_eq!(batch, TEST_BATCH.slice(0, 100)); + let batch = expect_data(decoder.try_decode()); + assert_eq!(batch, TEST_BATCH.slice(100, 100)); + } + /// Decode multiple columns "a" and "b", expect that the decoder requests /// only a single request per row group #[test] @@ -1370,6 +1443,105 @@ mod test { expect_finished(decoder.try_decode()); } + /// Consecutive filters on the same projection share one decode stream. + #[test] + fn test_decoder_same_projection_filters() { + let builder = + ParquetPushDecoderBuilder::try_new_decoder(test_file_parquet_metadata()).unwrap(); + let schema_descr = builder.metadata().file_metadata().schema_descr_ptr(); + let projection_a = ProjectionMask::columns(&schema_descr, ["a"]); + + let row_filter_gt = ArrowPredicateFn::new(projection_a.clone(), |batch: RecordBatch| { + let column = batch.column(0).as_primitive::(); + gt(column, &Int64Array::new_scalar(175)) + }); + let row_filter_lt = ArrowPredicateFn::new(projection_a, |batch: RecordBatch| { + let column = batch.column(0).as_primitive::(); + assert!(column.iter().flatten().all(|value| value > 175)); + lt(column, &Int64Array::new_scalar(190)) + }); + + let mut decoder = builder + .with_projection(ProjectionMask::columns(&schema_descr, ["c"])) + .with_row_filter(RowFilter::new(vec![ + Box::new(row_filter_gt), + Box::new(row_filter_lt), + ])) + .with_batch_size(50) + .build() + .unwrap(); + + // One filter-column request, followed by the selected output page. + let ranges = expect_needs_data(decoder.try_decode()); + push_ranges_to_decoder(&mut decoder, ranges); + let ranges = expect_needs_data(decoder.try_decode()); + push_ranges_to_decoder(&mut decoder, ranges); + + let batch = expect_data(decoder.try_decode()); + let expected = TEST_BATCH.slice(176, 14).project(&[2]).unwrap(); + assert_eq!(batch, expected); + + // Row group 1 is rejected by the fused predicates, so no output-column + // request is made. + let ranges = expect_needs_data(decoder.try_decode()); + push_ranges_to_decoder(&mut decoder, ranges); + expect_finished(decoder.try_decode()); + } + + /// Fused same-projection filters honour the output limit: evaluation + /// stops once enough rows survive the whole group. + #[test] + fn test_decoder_same_projection_filters_with_limit() { + use std::sync::atomic::{AtomicUsize, Ordering}; + + let builder = + ParquetPushDecoderBuilder::try_new_decoder(test_file_parquet_metadata()).unwrap(); + let schema_descr = builder.metadata().file_metadata().schema_descr_ptr(); + let projection_a = ProjectionMask::columns(&schema_descr, ["a"]); + + let rows_evaluated = Arc::new(AtomicUsize::new(0)); + let rows_evaluated_for_predicate = Arc::clone(&rows_evaluated); + let row_filter_gt = + ArrowPredicateFn::new(projection_a.clone(), move |batch: RecordBatch| { + rows_evaluated_for_predicate.fetch_add(batch.num_rows(), Ordering::Relaxed); + let column = batch.column(0).as_primitive::(); + gt(column, &Int64Array::new_scalar(175)) + }); + let row_filter_lt = ArrowPredicateFn::new(projection_a, |batch: RecordBatch| { + let column = batch.column(0).as_primitive::(); + lt(column, &Int64Array::new_scalar(190)) + }); + + let mut decoder = builder + .with_projection(ProjectionMask::columns(&schema_descr, ["c"])) + .with_row_filter(RowFilter::new(vec![ + Box::new(row_filter_gt), + Box::new(row_filter_lt), + ])) + .with_batch_size(10) + .with_limit(5) + .build() + .unwrap(); + + // One filter-column request for the fused group, then the output page. + let ranges = expect_needs_data(decoder.try_decode()); + push_ranges_to_decoder(&mut decoder, ranges); + let ranges = expect_needs_data(decoder.try_decode()); + push_ranges_to_decoder(&mut decoder, ranges); + + let batch = expect_data(decoder.try_decode()); + let expected = TEST_BATCH.slice(176, 5).project(&[2]).unwrap(); + assert_eq!(batch, expected); + + // The limit was satisfied by row group 0. + expect_finished(decoder.try_decode()); + + // Row 181 is the 6th match, so at most the batch holding it (rows + // 180..189) is evaluated. + let evaluated = rows_evaluated.load(Ordering::Relaxed); + assert!(evaluated <= 190, "evaluated {evaluated} rows"); + } + /// Decode with a filter that uses a column that is also projected, and expect /// that the filter pages are reused (don't refetch them) #[test] @@ -2706,7 +2878,7 @@ mod test { } /// return the metadata for the test file - pub fn test_file_parquet_metadata() -> Arc { + pub fn test_file_parquet_metadata() -> Arc { let mut metadata_decoder = ParquetMetaDataPushDecoder::try_new(test_file_len()).unwrap(); push_ranges_to_metadata_decoder(&mut metadata_decoder, vec![test_file_range()]); let metadata = metadata_decoder.try_decode().unwrap(); @@ -2716,6 +2888,19 @@ mod test { Arc::new(metadata) } + /// return the metadata for the test file, including the offset index + fn test_file_parquet_metadata_with_offset_index() -> Arc { + let mut metadata_decoder = ParquetMetaDataPushDecoder::try_new(test_file_len()) + .unwrap() + .with_offset_index_policy(PageIndexPolicy::Required); + push_ranges_to_metadata_decoder(&mut metadata_decoder, vec![test_file_range()]); + let metadata = metadata_decoder.try_decode().unwrap(); + let DecodeResult::Data(metadata) = metadata else { + panic!("Expected metadata to be decoded successfully"); + }; + Arc::new(metadata) + } + /// Push the given ranges to the metadata decoder, simulating reading from a file fn push_ranges_to_metadata_decoder( metadata_decoder: &mut ParquetMetaDataPushDecoder, diff --git a/parquet/src/column/chunker/cdc.rs b/parquet/src/column/chunker/cdc.rs index f524250b1b03..420bdf4f47da 100644 --- a/parquet/src/column/chunker/cdc.rs +++ b/parquet/src/column/chunker/cdc.rs @@ -1041,9 +1041,9 @@ mod arrow_tests { max_chunk_size: usize, expect_dictionary_page: bool, ) { - // Boolean and FixedSizeBinary never produce dictionary pages (matching C++) + // Boolean columns never produce dictionary pages. let expect_dict = match array.data_type() { - DataType::Boolean | DataType::FixedSizeBinary(_) => false, + DataType::Boolean => false, _ => expect_dictionary_page, }; assert_eq!( diff --git a/parquet/src/column/writer/encoder.rs b/parquet/src/column/writer/encoder.rs index 0eadecb8b038..cd0d2a267c70 100644 --- a/parquet/src/column/writer/encoder.rs +++ b/parquet/src/column/writer/encoder.rs @@ -276,8 +276,8 @@ impl ColumnValueEncoder for ColumnValueEncoderImpl { props: &WriterProperties, column_props: &ResolvedColumnProperties, ) -> Result { - let dict_supported = column_props.dictionary_enabled - && has_dictionary_support(T::get_physical_type(), props); + let dict_supported = + column_props.dictionary_enabled && has_dictionary_support(T::get_physical_type()); let dict_encoder = dict_supported.then(|| DictEncoder::new(descr.clone())); // Set either main encoder or fallback encoder. diff --git a/parquet/src/column/writer/mod.rs b/parquet/src/column/writer/mod.rs index 18def152f8e3..d7268422c494 100644 --- a/parquet/src/column/writer/mod.rs +++ b/parquet/src/column/writer/mod.rs @@ -617,12 +617,11 @@ impl<'a, E: ColumnValueEncoder> GenericColumnWriter<'a, E> { update_max(&self.descr, max, &mut self.column_metrics.max_column_value); } - // We can only set the distinct count if there are no other writes - if self.encoder.num_values() == 0 { - self.column_metrics.column_distinct_count = distinct_count; - } else { - self.column_metrics.column_distinct_count = None; - } + // Encoder counts reset per page; row metrics retain column-wide history. + let has_prior_data = self.column_metrics.total_rows_written != 0 + || self.page_metrics.num_buffered_values != 0; + self.column_metrics.column_distinct_count = + if has_prior_data { None } else { distinct_count }; let mut values_offset = 0; let mut levels_offset = 0; @@ -1978,13 +1977,10 @@ fn fallback_encoding(kind: Type, props: &WriterProperties) -> Encoding { } /// 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()) { +fn has_dictionary_support(kind: Type) -> bool { + match kind { // 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, + Type::BOOLEAN => false, _ => true, } } @@ -2571,14 +2567,38 @@ mod tests { #[test] fn test_column_writer_default_encoding_support_fixed_len_byte_array() { + for version in [WriterVersion::PARQUET_1_0, WriterVersion::PARQUET_2_0] { + let default_props = WriterProperties::builder() + .set_writer_version(version) + .build(); + let meta = column_write_and_get_metadata::( + default_props, + &[ByteArray::from(vec![1u8]).into()], + ); + assert_eq!(meta.dictionary_page_offset(), Some(0)); + } + check_encoding_write_support::( WriterVersion::PARQUET_1_0, true, &[ByteArray::from(vec![1u8]).into()], - None, - &[Encoding::PLAIN, Encoding::RLE], - &[encoding_stats(PageType::DATA_PAGE, Encoding::PLAIN, 1)], + Some(0), + &[Encoding::PLAIN, Encoding::RLE, Encoding::RLE_DICTIONARY], + &[ + encoding_stats(PageType::DICTIONARY_PAGE, Encoding::PLAIN, 1), + encoding_stats(PageType::DATA_PAGE, Encoding::RLE_DICTIONARY, 1), + ], + ); + let column_props = WriterProperties::builder() + .set_writer_version(WriterVersion::PARQUET_1_0) + .set_dictionary_enabled(false) + .set_column_dictionary_enabled(ColumnPath::from("col"), true) + .build(); + let meta = column_write_and_get_metadata::( + column_props, + &[ByteArray::from(vec![1u8]).into()], ); + assert_eq!(meta.dictionary_page_offset(), Some(0)); check_encoding_write_support::( WriterVersion::PARQUET_1_0, false, @@ -2817,6 +2837,7 @@ mod tests { let props = Arc::new( WriterProperties::builder() .set_write_page_header_statistics(true) + .set_data_page_row_count_limit(4) .build(), ); let mut writer = get_test_column_writer::(page_writer, 0, 0, props); @@ -2850,22 +2871,18 @@ mod tests { .unwrap(); let pages = reader.collect::>>().unwrap(); - assert_eq!(pages.len(), 2); + assert_eq!(pages.len(), 3); assert_eq!(pages[0].page_type(), PageType::DICTIONARY_PAGE); assert_eq!(pages[1].page_type(), PageType::DATA_PAGE); - - let page_statistics = pages[1].statistics().unwrap(); - assert_eq!( - page_statistics.min_bytes_opt().unwrap(), - 1_i32.to_le_bytes() - ); - assert_eq!( - page_statistics.max_bytes_opt().unwrap(), - 7_i32.to_le_bytes() - ); - assert_eq!(page_statistics.null_count_opt(), Some(0)); - assert!(page_statistics.distinct_count_opt().is_none()); + assert_eq!(pages[2].page_type(), PageType::DATA_PAGE); + for (page, min, max) in [(&pages[1], 1_i32, 4_i32), (&pages[2], 5_i32, 7_i32)] { + let stats = page.statistics().unwrap(); + assert_eq!(stats.min_bytes_opt().unwrap(), min.to_le_bytes()); + assert_eq!(stats.max_bytes_opt().unwrap(), max.to_le_bytes()); + assert_eq!(stats.null_count_opt(), Some(0)); + assert!(stats.distinct_count_opt().is_none()); + } } #[test] diff --git a/parquet/src/encodings/decoding/alp_decoder.rs b/parquet/src/encodings/decoding/alp_decoder.rs index 3847a57542b5..47f5edb4d37d 100644 --- a/parquet/src/encodings/decoding/alp_decoder.rs +++ b/parquet/src/encodings/decoding/alp_decoder.rs @@ -338,8 +338,6 @@ fn decode_range( cur.exception_positions.len() / std::mem::size_of::(), cur.exception_values.len() / Value::Exact::WIDTH, ); - #[expect(clippy::chunks_exact_to_as_chunks)] - // Requires using generic parameters in const operations: generic_const_exprs let values_chunks = cur.exception_values.chunks_exact(Value::Exact::WIDTH); for (pos_chunk, value_chunk) in cur .exception_positions diff --git a/parquet/src/encodings/encoding/mod.rs b/parquet/src/encodings/encoding/mod.rs index 05117a625a38..8b38d4940f38 100644 --- a/parquet/src/encodings/encoding/mod.rs +++ b/parquet/src/encodings/encoding/mod.rs @@ -120,15 +120,32 @@ pub(crate) mod private { "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()), + Encoding::RLE => match T::get_physical_type() { + Type::BOOLEAN => Box::new(RleValueEncoder::new()), + physical_type => return Err(unsupported_column_encoding(encoding, physical_type)), + }, + Encoding::DELTA_BINARY_PACKED => match T::get_physical_type() { + Type::INT32 | Type::INT64 => Box::new(DeltaBitPackEncoder::new()), + physical_type => return Err(unsupported_column_encoding(encoding, physical_type)), + }, + Encoding::DELTA_LENGTH_BYTE_ARRAY => match T::get_physical_type() { + Type::BYTE_ARRAY => Box::new(DeltaLengthByteArrayEncoder::new()), + physical_type => return Err(unsupported_column_encoding(encoding, physical_type)), + }, + Encoding::DELTA_BYTE_ARRAY => match T::get_physical_type() { + Type::BYTE_ARRAY | Type::FIXED_LEN_BYTE_ARRAY => { + Box::new(DeltaByteArrayEncoder::new()) + } + physical_type => return Err(unsupported_column_encoding(encoding, physical_type)), + }, Encoding::BYTE_STREAM_SPLIT => match T::get_physical_type() { Type::FIXED_LEN_BYTE_ARRAY => Box::new(VariableWidthByteStreamSplitEncoder::new( descr.type_length(), )), - _ => Box::new(ByteStreamSplitEncoder::new()), + Type::INT32 | Type::INT64 | Type::FLOAT | Type::DOUBLE => { + Box::new(ByteStreamSplitEncoder::new()) + } + physical_type => return Err(unsupported_column_encoding(encoding, physical_type)), }, Encoding::ALP => { return Err(general_err!( @@ -175,6 +192,14 @@ pub(crate) mod private { } } +fn unsupported_column_encoding(encoding: Encoding, physical_type: Type) -> ParquetError { + nyi_err!( + "Encoding {} is not supported for physical type {:?}", + encoding, + physical_type + ) +} + // ---------------------------------------------------------------------- // Plain encoding @@ -345,7 +370,6 @@ impl Encoder for RleValueEncoder { // DELTA_BINARY_PACKED encoding const MAX_PAGE_HEADER_WRITER_SIZE: usize = 32; -const DEFAULT_BIT_WRITER_SIZE: usize = 1024 * 1024; const DEFAULT_NUM_MINI_BLOCKS: usize = 4; /// Delta bit packed encoder. @@ -409,7 +433,10 @@ impl DeltaBitPackEncoder { DeltaBitPackEncoder { page_header_writer: BitWriter::new(MAX_PAGE_HEADER_WRITER_SIZE), - bit_writer: BitWriter::new(DEFAULT_BIT_WRITER_SIZE), + // Don't pre-allocate: encoders are often created eagerly (e.g. as a + // byte array fallback encoder) and may never be used. The buffer + // grows on demand and retains its capacity across pages. + bit_writer: BitWriter::new(0), total_values: 0, first_value: 0, current_value: 0, // current value to keep adding deltas @@ -849,8 +876,6 @@ mod tests { // supported encodings create_and_check_encoder::(0, Encoding::PLAIN, None); create_and_check_encoder::(0, Encoding::DELTA_BINARY_PACKED, None); - create_and_check_encoder::(0, Encoding::DELTA_LENGTH_BYTE_ARRAY, None); - create_and_check_encoder::(0, Encoding::DELTA_BYTE_ARRAY, None); create_and_check_encoder::(0, Encoding::RLE, None); // error when initializing @@ -868,6 +893,22 @@ mod tests { "Cannot initialize this encoding through this function" )), ); + create_and_check_encoder::( + 0, + Encoding::DELTA_LENGTH_BYTE_ARRAY, + Some(unsupported_column_encoding( + Encoding::DELTA_LENGTH_BYTE_ARRAY, + Type::INT32, + )), + ); + create_and_check_encoder::( + 0, + Encoding::DELTA_BYTE_ARRAY, + Some(unsupported_column_encoding( + Encoding::DELTA_BYTE_ARRAY, + Type::INT32, + )), + ); create_and_check_encoder::( 0, Encoding::ALP, @@ -885,6 +926,127 @@ mod tests { ); } + #[test] + fn column_construction_validates_encoding_fallback_matrix() { + use crate::file::properties::{WriterProperties, WriterVersion}; + use crate::file::writer::SerializedFileWriter; + use std::panic::{AssertUnwindSafe, catch_unwind}; + use std::sync::Arc; + + fn check_encoder(encoding: Encoding, supported: bool) { + let error = if supported { + None + } else if encoding == Encoding::ALP { + Some(general_err!( + "Encoding ALP only supports FLOAT and DOUBLE, got {}", + T::get_physical_type() + )) + } else { + Some(unsupported_column_encoding( + encoding, + T::get_physical_type(), + )) + }; + create_and_check_encoder::(4, encoding, error); + } + + let encodings = [ + Encoding::PLAIN, + Encoding::RLE, + Encoding::DELTA_BINARY_PACKED, + Encoding::DELTA_LENGTH_BYTE_ARRAY, + Encoding::DELTA_BYTE_ARRAY, + Encoding::BYTE_STREAM_SPLIT, + Encoding::ALP, + ]; + for physical in [ + Type::BOOLEAN, + Type::INT32, + Type::INT64, + Type::INT96, + Type::FLOAT, + Type::DOUBLE, + Type::BYTE_ARRAY, + Type::FIXED_LEN_BYTE_ARRAY, + ] { + let valid: &[Encoding] = match physical { + Type::BOOLEAN => &[Encoding::PLAIN, Encoding::RLE], + Type::INT32 | Type::INT64 => &[ + Encoding::PLAIN, + Encoding::DELTA_BINARY_PACKED, + Encoding::BYTE_STREAM_SPLIT, + ], + Type::INT96 => &[Encoding::PLAIN], + Type::FLOAT | Type::DOUBLE => { + &[Encoding::PLAIN, Encoding::BYTE_STREAM_SPLIT, Encoding::ALP] + } + Type::BYTE_ARRAY => &[ + Encoding::PLAIN, + Encoding::DELTA_LENGTH_BYTE_ARRAY, + Encoding::DELTA_BYTE_ARRAY, + ], + Type::FIXED_LEN_BYTE_ARRAY => &[ + Encoding::PLAIN, + Encoding::DELTA_BYTE_ARRAY, + Encoding::BYTE_STREAM_SPLIT, + ], + }; + // Check the factory's returned errors directly, independently of the + // infallible column writer constructor that unwraps those errors. + for encoding in encodings { + let supported = valid.contains(&encoding); + match physical { + Type::BOOLEAN => check_encoder::(encoding, supported), + Type::INT32 => check_encoder::(encoding, supported), + Type::INT64 => check_encoder::(encoding, supported), + Type::INT96 => check_encoder::(encoding, supported), + Type::FLOAT => check_encoder::(encoding, supported), + Type::DOUBLE => check_encoder::(encoding, supported), + Type::BYTE_ARRAY => check_encoder::(encoding, supported), + Type::FIXED_LEN_BYTE_ARRAY => { + check_encoder::(encoding, supported) + } + } + } + for version in [WriterVersion::PARQUET_1_0, WriterVersion::PARQUET_2_0] { + for dictionary in [false, true] { + for encoding in encodings { + let field = Arc::new( + SchemaType::primitive_type_builder("a", physical) + .with_length(4) + .build() + .unwrap(), + ); + let schema = Arc::new( + SchemaType::group_type_builder("schema") + .with_fields(vec![field]) + .build() + .unwrap(), + ); + let props = Arc::new( + WriterProperties::builder() + .set_writer_version(version) + .set_dictionary_enabled(dictionary) + .set_encoding(encoding) + .build(), + ); + let created = catch_unwind(AssertUnwindSafe(|| { + let mut writer = + SerializedFileWriter::new(Vec::new(), schema, props).unwrap(); + let mut group = writer.next_row_group().unwrap(); + group.next_column().unwrap().unwrap(); + })); + assert_eq!( + created.is_ok(), + valid.contains(&encoding), + "{physical:?} {version:?} dictionary={dictionary} {encoding:?}" + ); + } + } + } + } + } + #[test] fn test_bool() { BoolType::test(Encoding::PLAIN, TEST_SET_SIZE, -1); diff --git a/parquet/src/file/metadata/writer.rs b/parquet/src/file/metadata/writer.rs index 3a00f855bcfd..fcbc63a9c509 100644 --- a/parquet/src/file/metadata/writer.rs +++ b/parquet/src/file/metadata/writer.rs @@ -83,7 +83,7 @@ impl<'a, W: Write> ThriftMetadataWriter<'a, W> { column_metadata, row_group_idx, column_idx, - &mut self.buf, + &mut *self.buf, )?; let end_offset = self.buf.bytes_written(); // set offset and index for offset index @@ -117,7 +117,7 @@ impl<'a, W: Write> ThriftMetadataWriter<'a, W> { column_metadata, row_group_idx, column_idx, - &mut self.buf, + &mut *self.buf, )? { let end_offset = self.buf.bytes_written(); // set offset and index for offset index @@ -216,7 +216,7 @@ impl<'a, W: Write> ThriftMetadataWriter<'a, W> { // Write file metadata let start_pos = self.buf.bytes_written(); self.object_writer - .write_file_metadata(&file_meta, &mut self.buf)?; + .write_file_metadata(&file_meta, &mut *self.buf)?; let end_pos = self.buf.bytes_written(); // Write footer diff --git a/parquet/src/file/properties.rs b/parquet/src/file/properties.rs index 34dc983eacc1..e658c7edd0ed 100644 --- a/parquet/src/file/properties.rs +++ b/parquet/src/file/properties.rs @@ -2596,6 +2596,7 @@ mod tests { } #[test] + #[expect(clippy::float_cmp_const)] fn test_bloom_filter_builder_default() { let props = BloomFilterProperties::builder().build(); assert_eq!(props.fpp, DEFAULT_BLOOM_FILTER_FPP); @@ -2615,6 +2616,7 @@ mod tests { } #[test] + #[expect(clippy::float_cmp_const)] fn test_bloom_filter_builder_explicit_ndv() { let props = BloomFilterProperties::builder().with_max_ndv(1000).build(); assert_eq!(props.fpp, DEFAULT_BLOOM_FILTER_FPP); @@ -2656,6 +2658,7 @@ mod tests { } #[test] + #[expect(clippy::float_cmp_const)] fn test_column_specific_implicit_ndv_uses_row_group_size() { let custom_row_group_size: usize = 7777; let col = ColumnPath::from("col"); diff --git a/parquet/src/file/serialized_reader.rs b/parquet/src/file/serialized_reader.rs index 7345fb96457b..aab7bc8ab304 100644 --- a/parquet/src/file/serialized_reader.rs +++ b/parquet/src/file/serialized_reader.rs @@ -1154,7 +1154,6 @@ impl PageReader for SerializedPageReader { } else { *page_index += 1; } - Ok(()) } SerializedPageReaderState::Pages { page_locations, @@ -1171,10 +1170,9 @@ impl PageReader for SerializedPageReader { *page_index += 1; } } - - Ok(()) } } + Ok(()) } fn at_record_boundary(&mut self) -> Result { diff --git a/parquet/src/file/statistics.rs b/parquet/src/file/statistics.rs index fec89a8c2997..b7d2206da5a5 100644 --- a/parquet/src/file/statistics.rs +++ b/parquet/src/file/statistics.rs @@ -753,8 +753,6 @@ impl fmt::Debug for ValueStatistics { #[cfg(test)] mod tests { - use core::f32; - use super::*; #[test] diff --git a/parquet/src/file/writer.rs b/parquet/src/file/writer.rs index af60326ca387..7030cb28c671 100644 --- a/parquet/src/file/writer.rs +++ b/parquet/src/file/writer.rs @@ -745,7 +745,7 @@ impl<'a, W: Write + Send> SerializedRowGroupWriter<'a, W> { let (src_offset, src_length, write_offset) = self.begin_appended_column(&close)?; let mut read = read.take(src_length as _); - let write_length = std::io::copy(&mut read, &mut self.buf)?; + let write_length = std::io::copy(&mut read, &mut *self.buf)?; if src_length as u64 != write_length { return Err(general_err!( @@ -1064,7 +1064,7 @@ impl<'a, W: Write> SerializedPageWriter<'a, W> { page_encryptor.encrypt_page_header(&header, sink)?; } None => { - let mut protocol = ThriftCompactOutputProtocol::new(&mut self.sink); + let mut protocol = ThriftCompactOutputProtocol::new(&mut *self.sink); header.write_thrift(&mut protocol)?; } } diff --git a/parquet/src/record/reader.rs b/parquet/src/record/reader.rs index 2e2ddd4db5e0..f5886e05d1ce 100644 --- a/parquet/src/record/reader.rs +++ b/parquet/src/record/reader.rs @@ -230,10 +230,6 @@ impl TreeBuilder { 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(), path, @@ -1890,6 +1886,115 @@ mod tests { } } + #[test] + #[cfg(feature = "arrow")] + fn test_compound_map_key() { + use crate::arrow::ArrowWriter; + use arrow_array::builder::{Int32Builder, ListBuilder}; + use arrow_array::{ + Array, ArrayRef, Int32Array, MapArray, RecordBatch, StringArray, StructArray, + }; + use arrow_buffer::OffsetBuffer; + use arrow_schema::{DataType, Field as ArrowField, Fields}; + + let map_array = |keys: ArrayRef, offsets: Vec, values: Vec| { + let entries = StructArray::from(vec![ + ( + Arc::new(ArrowField::new("key", keys.data_type().clone(), false)), + keys, + ), + ( + Arc::new(ArrowField::new("value", DataType::Int32, true)), + Arc::new(Int32Array::from(values)) as ArrayRef, + ), + ]); + let entries_field = ArrowField::new("key_value", entries.data_type().clone(), false); + let map = MapArray::try_new( + Arc::new(entries_field), + OffsetBuffer::new(offsets.into()), + entries, + None, + false, + ) + .unwrap(); + Arc::new(map) as ArrayRef + }; + + // struct keys: [{(1, "a") -> 10, (2, "b") -> 20}, {}, {(3, "c") -> 30}] + let struct_keys = StructArray::new( + Fields::from(vec![ + ArrowField::new("id", DataType::Int32, false), + ArrowField::new("name", DataType::Utf8, false), + ]), + vec![ + Arc::new(Int32Array::from(vec![1, 2, 3])), + Arc::new(StringArray::from(vec!["a", "b", "c"])), + ], + None, + ); + let struct_map = map_array(Arc::new(struct_keys), vec![0, 2, 2, 3], vec![10, 20, 30]); + + // list keys: [{[1, 2] -> 10, [3] -> 20}, {}, {[] -> 30}] + let mut list_keys = ListBuilder::new(Int32Builder::new()); + list_keys.append_value([Some(1), Some(2)]); + list_keys.append_value([Some(3)]); + list_keys.append_value([]); + let list_map = map_array( + Arc::new(list_keys.finish()), + vec![0, 2, 2, 3], + vec![10, 20, 30], + ); + + let batch = + RecordBatch::try_from_iter([("struct_map", struct_map), ("list_map", list_map)]) + .unwrap(); + let mut buffer = Vec::new(); + let mut writer = ArrowWriter::try_new(&mut buffer, batch.schema(), None).unwrap(); + writer.write(&batch).unwrap(); + writer.close().unwrap(); + + let reader = SerializedFileReader::new(Bytes::from(buffer)).unwrap(); + let rows: Vec = reader + .get_row_iter(None) + .unwrap() + .map(|row| row.unwrap()) + .collect(); + + let key = |id: i32, name: &str| { + group![ + ("id".to_string(), Field::Int(id)), + ("name".to_string(), Field::Str(name.to_string())) + ] + }; + let expected_rows = vec![ + row![ + ( + "struct_map".to_string(), + map![(key(1, "a"), Field::Int(10)), (key(2, "b"), Field::Int(20))] + ), + ( + "list_map".to_string(), + map![ + (list![Field::Int(1), Field::Int(2)], Field::Int(10)), + (list![Field::Int(3)], Field::Int(20)) + ] + ), + ], + row![ + ("struct_map".to_string(), map![]), + ("list_map".to_string(), map![]), + ], + row![ + ( + "struct_map".to_string(), + map![(key(3, "c"), Field::Int(30))] + ), + ("list_map".to_string(), map![(list![], Field::Int(30))]), + ], + ]; + assert_eq!(rows, expected_rows); + } + fn test_file_reader_rows(file_name: &str, schema: Option) -> Result> { let file = get_test_file(file_name); let file_reader: Box = Box::new(SerializedFileReader::new(file)?); diff --git a/parquet/tests/arrow_writer/mod.rs b/parquet/tests/arrow_writer/mod.rs index 195e855c4086..c2803c5a1331 100644 --- a/parquet/tests/arrow_writer/mod.rs +++ b/parquet/tests/arrow_writer/mod.rs @@ -31,7 +31,7 @@ use std::fs::File; use std::io::{Read as _, Seek, SeekFrom, Write as _}; use std::sync::Arc; -use arrow::array::{ArrayRef, BinaryArray, Float64Array, Int32Array, RecordBatch}; +use arrow::array::{ArrayRef, BinaryArray, Float64Array, Int32Array, RecordBatch, StringArray}; use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; use bytes::Bytes; use parquet::arrow::arrow_writer::{ @@ -41,13 +41,11 @@ use parquet::arrow::arrow_writer::{ use parquet::arrow::{ArrowSchemaConverter, ArrowWriter}; use parquet::basic::Encoding; use parquet::errors::Result; -use parquet::file::properties::WriterProperties; +use parquet::file::properties::{WriterProperties, WriterVersion}; use parquet::file::writer::SerializedFileWriter; #[test] -#[should_panic( - expected = "DeltaBitPackDecoder only supports Int32Type, UInt32Type, Int64Type, and UInt64Type" -)] +#[should_panic(expected = "Encoding DELTA_BINARY_PACKED is not supported for physical type DOUBLE")] fn test_delta_bit_pack_type() { let props = WriterProperties::builder() .set_column_encoding("col".into(), Encoding::DELTA_BINARY_PACKED) @@ -572,3 +570,28 @@ fn page_store_spills_dictionary_pages() { (a few dictionary pages), not the ~K × dict_page of the in-memory baseline" ); } + +/// Regression test for +/// +/// Every byte array column eagerly builds its fallback encoder. With +/// `DELTA_BYTE_ARRAY` that holds two `DeltaBitPackEncoder`s, which must not +/// pre-allocate large buffers: here the columns dictionary encode, so the +/// fallback is never used at all. +#[test] +fn unused_delta_fallback_does_not_preallocate() { + const COLUMNS: usize = 100; + let column: ArrayRef = Arc::new(StringArray::from(vec!["a", "b", "a"])); + let batch = RecordBatch::try_from_iter((0..COLUMNS).map(|i| (format!("c{i}"), column.clone()))) + .unwrap(); + let props = WriterProperties::builder() + .set_writer_version(WriterVersion::PARQUET_2_0) + .set_encoding(Encoding::DELTA_BYTE_ARRAY) + .build(); + + let peak = peak_heap_bytes(|| { + let mut writer = ArrowWriter::try_new(Vec::new(), batch.schema(), Some(props)).unwrap(); + writer.write(&batch).unwrap(); + }); + + assert!(peak < COLUMNS * 64 * 1024, "peak: {peak}"); +} diff --git a/parquet_derive/Cargo.toml b/parquet_derive/Cargo.toml index 0855720628d4..0cc9d50477c9 100644 --- a/parquet_derive/Cargo.toml +++ b/parquet_derive/Cargo.toml @@ -20,7 +20,6 @@ name = "parquet_derive" version = { workspace = true } license = { workspace = true } description = "Derive macros for the Rust implementation of Apache Parquet" -homepage = { workspace = true } repository = { workspace = true } authors = { workspace = true } keywords = ["parquet"] diff --git a/parquet_derive_test/Cargo.toml b/parquet_derive_test/Cargo.toml index b7d10ab398c1..118e75cf0e68 100644 --- a/parquet_derive_test/Cargo.toml +++ b/parquet_derive_test/Cargo.toml @@ -20,7 +20,6 @@ name = "parquet_derive_test" version = { workspace = true } license = { workspace = true } description = "Integration test package for parquet-derive" -homepage = { workspace = true } repository = { workspace = true } authors = { workspace = true } keywords = [ "parquet" ] diff --git a/rust-toolchain.toml b/rust-toolchain.toml index daca2b13f389..e75546c87156 100644 --- a/rust-toolchain.toml +++ b/rust-toolchain.toml @@ -16,5 +16,5 @@ # under the License. [toolchain] -channel = "1.98.1" +channel = "1.99.0" components = ["rustfmt", "clippy"]